PySpark多行语句中如何引用变量作为分母计算占比

假设你已经完成分组聚合，得到包含major和n_students列的DataFrame，要添加prop列（专业学生占总人数的比例），只需借助PySpark的lit()函数将本地变量current_students转为Spark可识别的常量，再进行列运算即可：

步骤1：导入必要函数 #

from pyspark.sql.functions import col, lit, round

步骤2：定义总学生数变量 #

current_students = 2055

步骤3：添加比例列 #

假设聚合后的DataFrame名为agg_current（如果你的聚合结果存在原current中，直接替换为current即可）：

# 基础版：计算比例
agg_current = agg_current.withColumn("prop", col("n_students") / lit(current_students))

# 进阶版：保留4位小数（可选）
agg_current = agg_current.withColumn("prop", round(col("n_students") / lit(current_students), 4))

示例输出 #

执行agg_current.show()会得到类似如下结果：

+-------+----------+------+
|  major|n_students|  prop|
+-------+----------+------+
|Physics|       320|0.1557|
|Math   |       450|0.2190|
|CS     |       580|0.2822|
+-------+----------+------+

关键说明 #

• lit(current_students)的作用是将Python本地变量转换为PySpark的常量列，让分布式的DataFrame能引用这个数值计算，避免硬编码2055带来的维护问题。
• col("n_students")用于引用DataFrame中的现有列，确保运算在Spark分布式环境中执行。

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