PySpark中如何统计连续双词对数量？求内置实现方案

嗨，很高兴能帮你解决这个问题！你提到Scala里可以用sliding（应该是你说的slicing）来实现连续双词对的统计，其实PySpark不管是用RDD还是DataFrame API，都有对应的内置方案，完全不需要自定义函数~

先纠正你代码里的小问题 #

你推测的代码里用了splicing，但PySpark里并没有这个方法，对应的其实是RDD的sliding方法，或者DataFrame API里的数组操作函数。下面我分两种场景给你具体的实现：

场景1：用RDD实现（类似Scala的RDD sliding） #

如果你的数据是以RDD形式存在的，分两种情况：

情况A：统计每个句子内部的连续双词对 #

假设你的RDD每个元素是一个分词后的词数组（比如["I", "love", "PySpark"]），可以这样做：

# 对每个词数组应用滑动窗口大小为2，然后展开所有双词对
bigrams_rdd = rdd.flatMap(lambda words: words.sliding(2))
# 转换为键值对并统计数量
bigram_counts = bigrams_rdd.map(lambda bg: (tuple(bg), 1)).reduceByKey(lambda x, y: x + y)

情况B：统计全局连续的双词对（跨句子） #

如果你的RDD每个元素是单个词（比如从文本行flatMap分词后得到的），可以直接用RDD的sliding方法：

# 生成全局滑动窗口大小为2的双词对
bigrams_rdd = rdd.sliding(2)
# 统计计数
bigram_counts = bigrams_rdd.map(lambda bg: (tuple(bg), 1)).reduceByKey(lambda x, y: x + y)

场景2：用DataFrame API实现（更推荐，优化更好） #

如果你的数据是DataFrame，且有一列（比如叫words）是分词后的数组类型，用DataFrame API会更高效，因为PySpark会自动优化执行计划：

from pyspark.sql import functions as F

# 第一步：生成每个句子的连续双词对数组
df_with_bigrams = df.withColumn("bigrams", F.expr("""
    transform(
        range(0, size(words) - 1),
        i -> struct(words[i] as first_word, words[i+1] as second_word)
    )
"""))

# 第二步：展开双词对数组，然后分组计数
bigram_counts_df = df_with_bigrams.select(F.explode("bigrams").alias("bigram")) \
    .groupBy("bigram.first_word", "bigram.second_word") \
    .count() \
    .withColumnRenamed("count", "frequency")

或者更简洁的写法，用arrays_zip结合slice：

# 生成两个数组：去掉最后一个词的数组，和去掉第一个词的数组
df_sliced = df.withColumn("prev_words", F.slice(F.col("words"), 1, F.size(F.col("words")) - 1)) \
    .withColumn("next_words", F.slice(F.col("words"), 2, F.size(F.col("words")) - 1))

# 把两个数组拉链成双词对数组，再展开统计
bigram_counts_df = df_sliced.select(F.explode(F.arrays_zip("prev_words", "next_words")).alias("bigram")) \
    .groupBy("bigram.prev_words", "bigram.next_words") \
    .count()

这两种方法都是PySpark的内置实现，不需要自己写自定义函数，而且性能比自定义函数要好很多，尤其是DataFrame API的方式，能充分利用PySpark的优化器。

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