Spark 3.1实现DataFrame与权重字典的点积计算需求

针对Spark 3.1无法使用spark.pandas的场景，你可以通过原生PySpark的列表达式实现与pandasdot等价的点积计算，核心思路是将权重字典的键值对转换成「列×权重」的表达式，再对这些表达式求和。

完整实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 初始化SparkSession
spark = SparkSession.builder.appName("DotProductCalculation").getOrCreate()

# 定义权重字典
dic = {'SG_actions': 1, 'SO_actions': 2, 'GS_actions': 3}

# 创建测试DataFrame
data = [
    (1, 2, 0, 0),
    (1, 0, 1, 1),
    (2, 1, 1, 1)
]
schema = ["Ag", "SG_actions", "SO_actions", "GS_actions"]
df = spark.createDataFrame(data, schema=schema)

# 构建点积计算表达式：遍历字典，生成列×权重的表达式，再求和
dot_product_expr = sum(col(col_name) * weight for col_name, weight in dic.items())

# 添加新列并展示结果
result_df = df.withColumn("New col", dot_product_expr)
result_df.show()

代码解释 #

1. 表达式构建：通过生成器遍历权重字典的每个键值对，对匹配的列执行列值×权重的计算，最后用sum()将所有乘积结果累加，完全等价于pandas中的dot点积操作。
2. 兼容性：全程使用PySpark 3.1原生API，无需依赖spark.pandas或更高版本特性。
3. 结果验证：运行代码后会输出与你期望一致的DataFrame，包含新增的点积计算列。

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