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

Question

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

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

下面是当前实现的 MWE，在 pyspark 数据帧转换为 pandas 之后。

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                                   |
+-----------------------------------------------+--------------------------------------------+