PySpark DataFrame中按类别实现向量元素级求和与归一化

原始DataFrame #

data = [{"category": 'A', "bigram": 'delicious spaghetti', "vector": [0.01, -0.02, 0.03], 'all_vector' : 2},
        {"category": 'A', "bigram": 'delicious dinner', "vector": [0.04, 0.05, 0.06], 'all_vector' : 2},
        {"category": 'B', "bigram": 'new blog', "vector": [-0.14, -0.15, -0.16], 'all_vector' : 2},
        {"category": 'B', "bigram": 'bright sun', "vector": [0.071, -0.09, 0.063], 'all_vector' : 2}
        ]

sdf = spark.createDataFrame(data)

输出展示：

+----------+-------------------+--------+---------------------+
|all_vector|bigram             |category|vector               |
+----------+-------------------+--------+---------------------+
|2         |delicious spaghetti|A       |[0.01, -0.02, 0.03]  |
|2         |delicious dinner   |A       |[0.04, 0.05, 0.06]   |
|2         |new blog           |B       |[-0.14, -0.15, -0.16]|
|2         |bright sun         |B       |[0.071, -0.09, 0.063]|
+----------+-------------------+--------+---------------------+

需求 #

• 按category列分组
• 对每组的vector列执行元素级求和（即每组内所有向量的对应位置元素相加）
• 将求和结果除以all_vector完成归一化，生成norm_vector列
• 同一分组的所有行均使用该组的归一化向量

问题代码（无法运行） #

@udf_annotator(returnType=ArrayType(FloatType()))
def result_vector(vector, all_vector):  
    lst = [sum(x) for x in zip(*vector)] / all_vector
    return lst

sdf_new  = sdf\
    .withColumn('norm_vector', result_vector(F.col('vector'), F.col('all_vector')))\
    .withColumn('rank', F.row_number().over(Window.partitionBy('category')))\
    .where(F.col('rank') == 1)

代码问题分析 #

1. UDF参数错误：UDF接收的是单条记录的vector（单个数组），而非分组后的所有向量集合，zip(*vector)会因输入类型不匹配报错
2. 列表除法不支持：Python列表无法直接与数字做除法，需逐个元素运算
3. 结果过滤错误：添加row_number()并过滤rank==1会丢失分组内其他行，不符合需求

正确解决方案 #

方案一：固定长度向量（Spark内置函数） #

适用于已知向量长度的场景，直接按索引求和：

from pyspark.sql import functions as F

# 分组计算向量元素和与归一化向量
grouped_df = sdf.groupBy("category") \
    .agg(
        F.first("all_vector").alias("all_vector"),
        F.array(
            F.sum(F.col("vector")[0]),
            F.sum(F.col("vector")[1]),
            F.sum(F.col("vector")[2])
        ).alias("sum_vector")
    ) \
    .withColumn("norm_vector", F.transform("sum_vector", lambda x: x / F.col("all_vector")))

# 将归一化向量关联回原DataFrame
result_df = sdf.join(grouped_df.select("category", "norm_vector"), on="category", how="left")

# 展示结果
result_df.show(truncate=False)

方案二：任意长度向量（通用版） #

适用于向量长度不固定的场景，使用aggregate和transform实现通用元素级聚合：

from pyspark.sql import functions as F

# 分组计算通用长度向量的元素和与归一化向量
grouped_df = sdf.groupBy("category") \
    .agg(
        F.first("all_vector").alias("all_vector"),
        # 聚合所有向量：元素级求和
        F.aggregate(
            F.collect_list("vector"),
            # 初始化全0向量，长度与分组内第一个向量一致
            F.array_repeat(F.lit(0.0), F.size(F.first("vector"))),
            # 累加每个向量的对应元素
            lambda acc, x: F.transform(F.arrays_zip(acc, x), lambda pair: pair["0"] + pair["1"])
        ).alias("sum_vector")
    ) \
    .withColumn("norm_vector", F.transform("sum_vector", lambda x: x / F.col("all_vector")))

# 关联回原DataFrame
result_df = sdf.join(grouped_df.select("category", "norm_vector"), on="category", how="left")

result_df.show(truncate=False)

输出结果 #

+----------+-------------------+--------+---------------------+-----------------------+
|all_vector|bigram             |category|vector               |norm_vector            |
+----------+-------------------+--------+---------------------+-----------------------+
|2         |delicious spaghetti|A       |[0.01, -0.02, 0.03]  |[0.025, 0.015, 0.045]  |
|2         |delicious dinner   |A       |[0.04, 0.05, 0.06]   |[0.025, 0.015, 0.045]  |
|2         |new blog           |B       |[-0.14, -0.15, -0.16]|[-0.0345, -0.12, -0.0485]|
|2         |bright sun         |B       |[0.071, -0.09, 0.063]|[-0.0345, -0.12, -0.0485]|
+----------+-------------------+--------+---------------------+-----------------------+

注：手动计算验证：

• A组求和：[0.01+0.04, -0.02+0.05, 0.03+0.06] = [0.05, 0.03, 0.09]，除以2得[0.025, 0.015, 0.045]
• B组求和：[-0.14+0.071, -0.15+(-0.09), -0.16+0.063] = [-0.069, -0.24, -0.097]，除以2得[-0.0345, -0.12, -0.0485]

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