PySpark按ID分组并自定义计算A/B值的实现方法

针对你的需求，核心思路是先把每个ID下A、B类型的数据转成宽表（同一行展示A和B的数值），再做除法计算，最后整理成目标格式。以下是具体实现步骤和代码：

步骤1：将长表转宽表 #

通过groupBy按ID分组，pivot将type列的A、B值转为列名，再用agg提取对应value_one和value_two的数值（因为每个ID每个type仅一行，用first即可）：

from pyspark.sql import functions as F

# 假设原始DataFrame名为df
wide_df = df.groupBy("id") \
            .pivot("type") \
            .agg(F.first("value_one").alias("value_one"),
                 F.first("value_two").alias("value_two"))

# 转宽表后结构示例：
# | id | A_value_one | A_value_two | B_value_one | B_value_two |

步骤2：计算C类型的数值 #

基于宽表计算A/B的结果，同时指定type为'C'，用round函数控制小数位数（比如保留两位）：

result_df = wide_df.select(
    "id",
    (F.col("A_value_one") / F.col("B_value_one")).alias("value_one"),
    F.lit("C").alias("type"),
    (F.col("A_value_two") / F.col("B_value_two")).alias("value_two")
).withColumn("value_one", F.round(F.col("value_one"), 2)) \
 .withColumn("value_two", F.round(F.col("value_two"), 2))

步骤3：查看结果 #

执行result_df.show()就能得到目标数据表：

+---+---------+----+---------+
| id|value_one|type|value_two|
+---+---------+----+---------+
|  1|      0.5|   C|      1.0|
|  2|      3.0|   C|     0.67|
+---+---------+----+---------+

补充说明 #

• 若担心出现除以零的情况，可添加判断逻辑：F.when(F.col("B_value_one") != 0, F.col("A_value_one")/F.col("B_value_one")).otherwise(0)
• 若原始数据存在空值，可在agg时用F.coalesce处理，避免计算出错

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