PySpark：按Type和Date分组获取Number最大值对应行

问题描述 #

现有如下PySpark DataFrame：

Type | Number |    Date    | Value |
------------------------------------
   A |      1 | 2022-10-01 |     5 |
   A |      2 | 2022-10-01 |     8 |
   A |      3 | 2022-11-23 |     4 |
   B |      1 | 2022-02-02 |     1 |
   B |      2 | 2022-02-04 |     9 |
   B |      3 | 2022-02-04 |     3 |
   B |      4 | 2022-02-04 |     1 |

需要按Type和Date分组，获取每组内Number最大值对应的完整行，期望结果：

Type | Number |    Date    | Value |
------------------------------------
   A |      2 | 2022-10-01 |     8 |
   A |      3 | 2022-11-23 |     4 |
   B |      1 | 2022-02-02 |     1 |
   B |      4 | 2022-02-04 |     1 |

尝试了以下代码但未成功：

from pyspark.sql.functions import *
from pyspark.sql import Window

w = Window.partitionBy("Type", "Date")
df_result = (
    df.withColumn("MaxNumber", max("Number").over(w))
    .where(col("Number") == col("MaxNumber"))
    .drop("MaxNumber")
)
df_result.display()

解决方案 #

你的代码逻辑本身是正确的，未成功通常是数据类型或格式问题，以下是调整方案：

1. 修正数据类型（最常见原因） #

如果Number列是字符串类型，max()会按字符串规则排序（比如"10"会小于"2"），导致结果不符合预期。先将Number转换为数值类型：

from pyspark.sql.functions import *
from pyspark.sql import Window

# 转换Number为整数类型
df = df.withColumn("Number", col("Number").cast("int"))

w = Window.partitionBy("Type", "Date")
df_result = (
    df.withColumn("MaxNumber", max("Number").over(w))
    .where(col("Number") == col("MaxNumber"))
    .drop("MaxNumber")
)
df_result.display()

2. 可选：用row_number()确保单条结果 #

如果需要每组只返回一行（当存在多个相同最大值时），可以用row_number()按Number降序排序后取第一行：

from pyspark.sql.functions import *
from pyspark.sql import Window

w = Window.partitionBy("Type", "Date").orderBy(col("Number").desc())
df_result = (
    df.withColumn("row_rank", row_number().over(w))
    .where(col("row_rank") == 1)
    .drop("row_rank")
)
df_result.display()

这种方式会在多条最大值记录中只保留排序后的第一条，而原代码会保留所有最大值对应的行。

3. 验证分组列格式 #

确保Date列是date类型而非字符串，避免因格式不一致导致分组错误：

# 转换Date列为标准日期类型
df = df.withColumn("Date", col("Date").cast("date"))

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