PySpark按分组基于另一列新增列：求分组最大价格问题

问题描述 #

现有如下结构的PySpark DataFrame：

id item price
1  x    2.0
1  y
2  z    1.0

需要按id分组，为每一行添加对应分组的price列最大值，得到结果：

id item price max_price
1  x    2.0   2.0
1  y          2.0
2  z    1.0   1.0

尝试使用以下代码时，因price列为字符串类型无法转换而报错：

df = df.withColumn('price', max('price').over(Window.partitionBy('id')))

解决方案 #

核心思路 #

先将字符串类型的price转换为数值类型（自动保留空值），再通过窗口函数计算分组最大值。

实现代码 #

方法1：分步处理（更直观）

from pyspark.sql.types import DoubleType
from pyspark.sql.window import Window
import pyspark.sql.functions as F

# 1. 将price列转为数值类型，空值保持为null
df = df.withColumn("price_numeric", F.col("price").cast(DoubleType()))

# 2. 按id分组计算最大值，新增max_price列
window = Window.partitionBy("id")
df = df.withColumn("max_price", F.max("price_numeric").over(window))\
       .drop("price_numeric")  # 删除临时列

方法2：一步完成（简化写法）

无需创建临时列，直接在窗口函数中完成类型转换：

from pyspark.sql.window import Window
import pyspark.sql.functions as F

window = Window.partitionBy("id")
df = df.withColumn(
    "max_price",
    F.max(F.col("price").cast("double")).over(window)
)

说明 #

• cast("double")或cast(DoubleType())都能将字符串类型的数值转为双精度浮点数，空值会被保留为null，max函数会自动忽略空值，不影响最大值计算。
• 最终结果保留原始price列的内容，新增的max_price列对应分组内非空price的最大值。

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