本文介绍在 Polars DataFrame 中一次性提取每个分组（group）的首尾两行的多种方法，重点推荐基于 gather() 的高性能方案，并对比不同实现的性能与可读性。
本文介绍在 polars dataframe 中一次性提取每个分组（group）的首尾两行的多种方法，重点推荐基于 `gather()` 的高性能方案，并对比不同实现的性能与可读性。
在 Polars 中，对分组数据分别取首行（.first()）和末行（.last()）是常见需求，但若需合并为单次操作（即每组返回两条记录：第一条 + 最后一条），直接链式调用 .first() 和 .last() 会触发两次独立聚合，效率低且结果需手动拼接。幸运的是，Polars 提供了更优雅、更高效的原生方案——核心在于利用 pl.Expr.gather() 表达式结合窗口或聚合上下文。
✅ 推荐方案一：select + gather（最高性能）
适用于希望保持原始行序、无需显式 group_by 的场景，代码简洁且速度最快（基准测试中仅需 172 μs）：
import polars as pl
df = pl.DataFrame({
"a": [1, 2, 2, 3, 4, 5],
"b": [0.5, 0.5, 4, 10, 14, 13],
"c": [True, True, True, False, False, True],
"d": ["Apple", "Apple", "Apple", "Banana", "Banana", "Banana"],
})
result = df.select(
pl.all().gather([0, -1]).over("d", mapping_strategy="explode")
)
print(result)
输出：
shape: (4, 4) ┌────────┬─────┬──────┬───────┐ │ d ┆ a ┆ b ┆ c │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ f64 ┆ bool │ ╞════════╪═════╪══════╪═══════╡ │ Apple ┆ 1 ┆ 0.5 ┆ true │ │ Apple ┆ 2 ┆ 4.0 ┆ true │ │ Banana ┆ 3 ┆ 10.0 ┆ false │ │ Banana ┆ 5 ┆ 13.0 ┆ true │ └────────┴─────┴──────┴───────┘
? 原理说明：
pl.all().gather([0, -1])对每列在"d"分组内分别取索引0（首行）和-1（末行）；over("d")触发窗口计算；mapping_strategy="explode"将每个分组的两个值展开为两行，自动对齐所有列。
✅ 推荐方案二：group_by + agg + explode（语义最清晰）
当需要明确分组逻辑或后续需与其他聚合操作组合时，此方式更具可读性与扩展性（性能次优，255 μs）：
result = (
df
.group_by("d", maintain_order=True)
.agg(pl.all().gather([0, -1]))
.explode(pl.exclude("d"))
)
注意：explode() 默认展开所有列表列，此处用 pl.exclude("d") 确保分组键 "d" 不被展开（因其本身非列表类型，但显式排除更安全）。
⚠️ 注意事项与最佳实践
-
maintain_order=True是必需的：确保group_by保留原始顺序，否则gather([0, -1])可能取到错误的“首/尾”； -
避免
.concat([first(), last()])：虽可行，但需两次分组计算，性能差且易出错（如分组数不一致）； -
慎用
filter()方案：如(df.with_row_index().filter(pl.col('index').is_in([0, -1]).over('d')))——row_index在分组内不连续，-1含义模糊，不可靠； -
大数据集请优先测试：
select + gather在多数场景下最快，但若涉及复杂表达式或自定义逻辑，group_by + agg更易维护。
综上，对于“每组取首尾行”这一高频任务，首选 select(...).gather([0, -1]).over(...) 模式：它兼具高性能、低内存开销与声明式表达力，充分体现了 Polars 向量化计算的设计优势。
