PySpark无shuffle统计数组列预设词表元素计数方法

需求说明 #

已知固定完整词表，对DataFrame中数组类型列的每行数据，统计词表中每个元素的出现次数，生成和词表长度一致的计数向量。要求全程不触发shuffle，避免使用groupBy、全局聚合类操作，跳过低效的中间映射表构建步骤。

基础配置与样例数据 #

预设词表 #

vocab = ['A', 'B', 'C', 'D', 'E']

样例数据构造 #

import pandas as pd
from pyspark.sql import SparkSession
import pyspark.sql.functions as F

data = {'date': ['2014-01-01', '2014-01-02', '2014-01-03'],
     'flat': ['A;A;B', 'D;B;E;B;B', 'B;A']}

data = pd.DataFrame(data)
data['date'] = pd.to_datetime(data['date'])

spark = SparkSession.builder \
    .master('local[*]') \
    .config("spark.driver.memory", "500g") \
    .appName('pandasToSparkDF-app') \
    .getOrCreate()
spark.conf.set("spark.sql.execution.arrow.enabled", "true")
spark.sparkContext.setLogLevel("OFF")

df = spark.createDataFrame(data)
new_frame = df.withColumn("list", F.split("flat", "\;"))

期望输出 #

+-------------------+-----------+---------------------+
|               date| flat      | counts              |
+-------------------+-----------+---------------------+
|2014-01-01 00:00:00|A;A;B      |[2, 1, 0, 0, 0]      |
|2014-01-02 00:00:00|D;B;E;B;B  |[0, 3, 0, 1, 1]      |
|2014-01-03 00:00:00|B;A        |[1, 1, 0, 0, 0]      |
+-------------------+-----------+---------------------+

原有实现的问题 #

原有参考实现逻辑为：对每行数组先去重得到distinct_items，统计每个去重元素的出现次数，再将去重元素和次数组装为count_map映射，最后遍历词表从map中取值生成计数数组。该方案中间步骤多，额外生成多列临时数据，大规模数据集、多列计算场景下性能损耗明显。
原有实现参考代码如下：

from pyspark.sql import functions as F
df=spark.createDataFrame(data)
df.withColumn("list", F.split("flat","\;"))\
  .withColumn("distinct_items", F.array_distinct("list")) \
  .withColumn("occurrences", F.expr("""array_sort(transform(distinct_items, x-> aggregate(list, 0,(acc,t)->acc+IF(t=x,1,0))))"""))\
  .withColumn("count_map", F.map_from_arrays("distinct_items", "occurrences"))\
  .withColumn(
      "counts",
      F.array(
          [
              F.when(
                  F.col("count_map")
                      .getItem(v)
                      .isNull(),
                      0,
                  )
                  .otherwise(
                      F.col("count_map").getItem(v)
                  )
              for v in vocab
           ]
      )).drop("occurrences", "distinct_items").show()

高效无shuffle实现方案 #

直接基于Spark SQL内置的高阶函数，一步生成计数数组，完全跳过中间去重、map构建步骤，所有计算均为行级操作，不会触发shuffle。

实现代码 #

from pyspark.sql import functions as F

# 构造词表常量数组的SQL表达式
vocab_sql = ",".join([f"'{v}'" for v in vocab])

result = df.withColumn("list", F.split("flat", "\;")) \
    .withColumn(
        "counts",
        F.expr(f"""
            transform(
                array({vocab_sql}),
                target -> aggregate(
                    list, 
                    0, 
                    (acc, elem) -> acc + IF(elem = target, 1, 0)
                )
            )
        """)
    )

result.show(truncate=False)

逻辑说明 #

• 首先将字符串格式的flat字段按分号拆分为数组列list
• 用transform函数遍历由固定词表生成的常量数组，逐个处理词表中的元素target
• 嵌套aggregate函数遍历当前行的list数组，累加和target相等的元素个数，最终直接得到和词表顺序一致、长度一致的计数数组
• 全程无任何shuffle操作，无临时中间列生成，性能相比原有实现提升明显，适配大规模数据场景。

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