如何在PySpark中使用bigram/n-gram算法获取字符级n元组列表？

• Spark MLlib的NGram组件默认是词级粒度，输入列要求为「分词后单个词汇组成的数组」，只有数组长度≥n时才会生成对应n-gram。你将完整字符串zhang1997作为数组唯一元素传入，数组长度为1小于n=2，自然输出空数组。

全程使用Spark内置算子实现，性能远高于自定义UDF，无需跨Python进程序列化开销：

1. 先把输入字符串拆分为单个字符组成的数组，过滤拆分产生的首尾空字符
2. 调用NGram组件生成字符级2-gram
3. 对生成的2-gram去重后统计长度，空文本直接返回0

可运行代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import split, expr, length, when
from pyspark.ml.feature import NGram

# 初始化Spark会话
spark = SparkSession.builder.appName("char_2gram_calc").getOrCreate()

# 构造测试数据，包含正常输入和空输入
test_df = spark.createDataFrame([(1, "zhang1997"), (2, "")], ["id", "text"])

# 步骤1：拆分字符串为单个字符的数组
char_split_df = test_df.withColumn(
    "char_list",
    # 拆分后过滤空值，避免异常
    expr("filter(split(text, ''), x -> trim(x) != '')")
)

# 步骤2：生成字符级2-gram
ngram_transformer = NGram(n=2, inputCol="char_list", outputCol="char_2grams")
ngram_result_df = ngram_transformer.transform(char_split_df)

# 步骤3：计算最终结果
final_df = ngram_result_df.withColumn(
    "result",
    when(length("text") == 0, 0).otherwise(expr("size(array_distinct(char_2grams))"))
)

# 输出验证结果
final_df.select("id", "text", "char_2grams", "result").show(truncate=False)

输出结果说明 #

测试输入zhang1997对应的char_2grams列会输出['zh','ha','an','ng','g1','19','99','97']，result列值为8；空输入对应的result列值为0，完全匹配需求。

内容的提问来源于stack exchange，提问作者Mario