多分类列拆分并按日期聚合统计的高效实现方案咨询

问题背景 #

需将多列分类数据按date字段聚合，把每个分类列的不同取值拆分为单独的统计列，统计每个分类的出现次数。原方案用Spark SQL的case when硬编码实现，当分类列数量达10个左右时，代码冗余且执行效率低下。

输入示例 #

输出示例 #

方案1：Spark DataFrame 动态透视（Pivot）实现 #

Spark的pivot操作针对行列转换做了专门优化，比手动编写大量case when更高效，尤其适合多分类列场景。步骤如下：

1. 将多列分类数据转为长格式（Unpivot），把每个分类列转为键值对形式；
2. 按date和分类键分组计数；
3. 透视转回宽格式，得到目标结构。

代码实现 #

from pyspark.sql import functions as F

# 假设输入数据已加载为DataFrame df
# df = spark.read.table("prep_data")

# 1. Unpivot：将多列转为长表结构
unpivot_df = df.select(
    F.col("date"),
    F.explode(
        F.array(
            *[F.struct(F.lit(col).alias("col_name"), F.col(col).alias("category")) 
              for col in ["Value1", "Value2", "Value3"]]
        )
    ).alias("data")
).select(
    "date",
    "data.col_name",
    "data.category"
)

# 2. 分组计数 + 3. Pivot转回宽表
result_df = unpivot_df.groupBy("date")\
    .pivot("col_name")\
    .pivot("category")\
    .count()\
    .fillna(0)  # 未出现的分类填充为0

# 重命名列以匹配输出格式
result_df = result_df.select(
    "date",
    *[F.col(c).alias(f"{c.split('_')[0]}_{c.split('_')[1]}") 
      for c in result_df.columns if c != "date"]
)

result_df.show()

优势 #

• 无需硬编码每个分类值，自动识别所有分类；
• Spark对pivot有执行优化，减少重复计算逻辑，多列场景下性能提升明显。

方案2：动态生成Spark SQL语句 #

如果偏好SQL写法，可通过Python动态生成透视所需的SQL代码，避免手动编写冗余的case when，同时保证执行效率。

代码实现 #

# 获取所有分类列及对应的唯一值
category_cols = ["Value1", "Value2", "Value3"]
category_values = {}
for col in category_cols:
    distinct_vals = [row[0] for row in df.select(col).distinct().collect()]
    category_values[col] = distinct_vals

# 动态生成count(case when ...)语句
case_when_clauses = []
for col in category_cols:
    for val in category_values[col]:
        clause = f"COUNT(CASE WHEN {col} = '{val}' THEN 1 END) AS {col}_{val}"
        case_when_clauses.append(clause)

# 拼接完整SQL
sql_query = f"""
SELECT 
    date,
    {', '.join(case_when_clauses)}
FROM prep_data
GROUP BY date
"""

# 执行SQL
result_df = spark.sql(sql_query)
result_df.show()

优势 #

• 保留SQL写法习惯，同时避免手动编写冗余代码；
• 分类列或值变化时无需修改代码，易维护。

性能对比说明 #

• 原case when方案：每个分类列的每个值都要单独判断，多列时会产生大量重复过滤逻辑，Spark执行计划存在冗余分支；
• 动态Pivot/动态SQL方案：Pivot通过统一的分组聚合逻辑处理所有分类，动态SQL生成最优结构，减少重复计算，在10+分类列场景下执行效率显著提升。

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