Spark按分组筛选指定条件行（取最短alias）的高效方法

当然有更高效的实现方式！针对你这个按id和name分组、筛选alias最短行的需求，我们可以根据你使用的是PySpark还是Pandas来选择更简洁高效的方法，下面分情况介绍：

先还原你的示例数据 #

首先我们先构建测试用的DataFrame：

# PySpark版本
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()

data = [
    (1, "abc", "short"),
    (1, "abc", "ailas-long-1"),
    (1, "abc", "another-long-alias"),
    (2, "xyz", "short_alias"),
    (2, "xyz", "same_length"),
    (3, "def", "alias_1")
]
df = spark.createDataFrame(data, schema=["id", "name", "alias"])

# Pandas版本
import pandas as pd
df = pd.DataFrame({
    "id": [1,1,1,2,2,3],
    "name": ["abc","abc","abc","xyz","xyz","def"],
    "alias": ["short","ailas-long-1","another-long-alias","short_alias","same_length","alias_1"]
})

一、PySpark 高效实现 #

1. 最优方案：groupBy + min_by（Spark 3.0+ 支持） #

Spark 3.0引入的min_by/max_by函数专门解决这种「分组后根据某个指标取对应行」的场景，直接一步到位，性能远优于Window+row_number的方式——因为它不需要生成额外的行号列，也不需要后续过滤，直接完成分组聚合：

import pyspark.sql.functions as F

result = df.groupBy("id", "name").agg(
    F.min_by("alias", F.length("alias")).alias("alias")
)

result.show()

输出完全符合你的预期：

+---+----+------------+
| id|name|       alias|
+---+----+------------+
|  1| abc|       short|
|  2| xyz|short_alias|
|  3| def|     alias_1|
+---+----+------------+

2. 旧版Spark兼容方案：sortWithinPartitions + groupBy + first #

如果你的Spark版本低于3.0，可以用这种方式，它比Window方法更高效——因为分区内排序后直接取第一行，避免了Window操作中为每一行生成row_number的额外开销：

result = df.orderBy("id", "name", F.length("alias")) \
           .groupBy("id", "name") \
           .agg(F.first("alias").alias("alias"))

result.show()

自定义筛选条件扩展 #

不管是min_by还是排序取first，都可以轻松扩展到自定义筛选条件：

• 比如要筛选alias字典序最小的行：F.min_by("alias", "alias")
• 比如要根据自定义函数计算的指标筛选：F.min_by("alias", your_custom_udf("alias"))

二、Pandas 高效实现 #

1. 最优方案：groupBy + idxmin #

利用idxmin获取每个分组中满足条件的行索引，再通过索引提取行，这是Pandas中性能最好的方式之一，因为idxmin是向量化操作：

# 先计算每个分组中alias长度最小的行的索引
min_len_idx = df.groupby(["id", "name"])["alias"].apply(lambda x: x.str.len().idxmin())
# 根据索引提取结果
result = df.loc[min_len_idx].reset_index(drop=True)

2. 简洁方案：groupBy + nsmallest #

如果追求代码简洁，可以用nsmallest直接取每个分组的前1行：

result = df.groupby(["id", "name"]).apply(
    lambda x: x.nsmallest(1, key=lambda row: len(row["alias"]))
).reset_index(drop=True)

为什么这些方法比Window+row_number更高效？ #

Window+row_number需要为每一行生成一个行号，然后再过滤行号=1的行，中间多了两次数据遍历（生成行号、过滤）；而上面的聚合/排序取首行的方法，都是直接针对分组做一次处理，减少了中间步骤和数据处理量，在大数据量场景下性能差异会更明显。

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