pyspark后缀替换避免某些单词而不映射到pandas或rdd

Question

我继承了一个程序，它修改pyspark数据帧中的一些字符串。其中一个步骤涉及从字符串中的某些单词中删除后缀，并附加一个异常列表，即使这些异常有后缀，也会被单独保留。目前，这是通过使用udf将数据帧转换为pandas来完成的，然后在读回pyspark之前，对生成的pandas数据帧中的字符串应用自定义函数。不幸的是，对需求的更改意味着代码在任何情况下都不能使用udf或映射到rdd。我需要直接在pyspark中执行相同的功能

后缀移除函数逐字读取字符串，检查单词是否在例外列表中，如果不在例外列表中，检查接受的后缀（这是一个严格的列表，我不能只使用现有的词干分析器），如果有，检查词干单词是否超过4个字符，如果是，则执行替换

以下是pyspark数据帧转换为pandas后当前实现的MWE

import pandas as pd


exception_list = ['WOODLAND', 'FISHING', 'LAUGHING']
suffix_list = ['ING', 'AND']

cols = ['input']
data = [
    ["CAT DOG FROG WOODLAND FARMLAND LAUGHING UNICORN"],
    ["BOG FISHING CARTING MISSING AND SOGGY"],
    ["SEARCHING"],
    ["FINDING"],
    ["SING SINGING"]
]

df = pd.DataFrame(data, columns=cols)
df.head()


def strip_sufx_word(word, suffix, exception, min_stem_length=4):
    for sufx in suffix:
        if word[-len(sufx):] == sufx:
            if len(word[:-len(sufx)])>=min_stem_length:
                if word not in exception:
                    word = word[:-len(sufx)]
    return word


def strip_sufx_string(phrase, suffix, exception):
    new_phrase = [strip_sufx_word(word, suffix, exception)
                  for word in phrase.split()]
    return ' '.join(new_phrase)


df['output'] = df['input'].apply(strip_sufx_string,
                                 suffix=suffix_list,
                                 exception=exception_list)

df.head()

有没有办法在pyspark中做到这一点？我愿意使用RegexTokenizer之类的东西，稍后再加入它，并创建附加的真理列，这些列在之后会被删除。它只需要在数据帧没有离开pyspark或映射到其他任何东西的情况下完成

Answer 1

高阶函数在这里很有用：

import pyspark.sql.functions as F

exception_list = ['WOODLAND', 'FISHING', 'LAUGHING']
suffix_list = ['ING', 'AND']
min_stem_length = 4

result = sdf.withColumn(
    'exception_list', 
    F.array(*[F.lit(w) for w in exception_list])
).withColumn(
    'suffix_list', 
    F.array(*[F.lit(w) for w in suffix_list])
).withColumn(
    'output', 
    F.expr(f"""
        concat_ws(' ', 
            transform(
                split(input, ' '), 
                word -> 
                    aggregate(
                        suffix_list, 
                        word, 
                        (acc, s) -> 
                            case when substring(acc, -length(s)) = s 
                                 and length(substring(acc, 1, length(acc)-length(s))) >= {min_stem_length} 
                                 and not array_contains(exception_list, acc) 
                                 then substring(acc, 1, length(acc)-length(s)) 
                                 else acc 
                            end
                     )
            )
        )
    """
    )
).drop('exception_list', 'suffix_list')

result.show(truncate=False)
+                       -+                      +
|input                                          |output                                      |
+                       -+                      +
|CAT DOG FROG WOODLAND FARMLAND LAUGHING UNICORN|CAT DOG FROG WOODLAND FARML LAUGHING UNICORN|
|BOG FISHING CARTING MISSING AND SOGGY          |BOG FISHING CART MISS AND SOGGY             |
|SEARCHING                                      |SEARCH                                      |
|FINDING                                        |FIND                                        |
|SING SINGING                                   |SING SING                                   |
+                       -+                      +