本文介绍如何在 Polars 中对分组数据（如按 cik 分组）内部，针对不同 form 类型（如 '10-Q' 和 '10-K'）分别保留指定数量的最新记录（如 1 个最新 10-Q、最多 10 个最新 10-K），避免低效的 map_groups，推荐使用 over() + int_range() 的向量化方案。
本文介绍如何在 polars 中对分组数据（如按 `cik` 分组）内部，针对不同 `form` 类型（如 `'10-q'` 和 `'10-k'`）分别保留指定数量的最新记录（如 1 个最新 `10-q`、最多 10 个最新 `10-k`），避免低效的 `map_groups`，推荐使用 `over()` + `int_range()` 的向量化方案。
在 Polars 中实现「分组内按条件筛选」是常见但易踩坑的操作。原始代码使用 group_by(...).map_groups(...) 配合 sort + filter + head/slice，虽功能正确，但存在明显缺陷：map_groups 禁用查询优化器、强制物化每组数据，性能差且难以维护。
更优雅、高性能的解法是利用 Polars 的窗口函数能力——通过 pl.Expr.over() 在逻辑分组上计算行序号，并结合布尔掩码一次性完成过滤。
✅ 推荐方案：over() + int_range() 向量化过滤
核心思路分三步：
- 全局预排序：先按 cik（升序）、form（升序）、period（降序）排序，确保同 cik 同 form 的记录中，最新日期排在最前；
- 组内编号：使用 pl.int_range(pl.len()).over("cik", "form") 为每个 (cik, form) 组内的行分配从 0 开始的连续索引；
-
条件过滤：基于 form 值和组内索引构建布尔表达式，例如：
- 10-Q：仅保留索引为 0 的行（即最新一条）；
- 10-K：保留索引
import polars as pl
import datetime
# 构造示例数据（2 个 cik，各含 3 条 10-Q 和 2 条 10-K）
df = pl.DataFrame({
"cik": [0] * 5 + [1] * 5,
"form": (["10-Q"] * 3 + ["10-K"] * 2) * 2,
"period": [datetime.date(2021, 1, 1 + day) for day in range(10)],
})
# ✅ 高效实现：每 cik 组内保留 1 条最新 10-Q + 最多 10 条最新 10-K
result = (
df
.sort(["cik", "form", "period"], descending=[False, False, True])
.filter(
(
(pl.col("form") == "10-Q") & (pl.int_range(pl.len()).over("cik", "form") == 0)
) | (
(pl.col("form") == "10-K") & (pl.int_range(pl.len()).over("cik", "form") <blockquote><p>? 提示：pl.int_range(pl.len()).over("cik", "form") 是关键——它不依赖物理分组，而是基于当前排序结果，在每个 (cik, form) 逻辑窗口内生成 0, 1, 2, ... 序号，完全向量化、零 Python 回调。</p></blockquote><h3>⚠️ 注意事项与最佳实践</h3>
- 务必先排序再 over：over 本身不保证顺序，必须显式 sort 确保时间字段降序排列，否则“最新”将失效；
- 避免 map_groups 除非必要：map_groups 会绕过 Polars 查询优化器，性能通常比等效的 over/window 表达式低 3–10 倍；
- over 分组键需完整覆盖逻辑维度：此处用 ("cik", "form") 而非仅 "cik"，确保 10-Q 和 10-K 的序号各自独立重置；
- 空组安全：若某 cik 下无 10-Q 或 10-K，该条件分支自动不匹配，不影响其他记录；
- 扩展性强：如需支持动态数量（如按 cik 规模调整保留条数），可结合 pl.col("cik").n_unique().over("cik") 或外部映射字典。
✅ 总结
| 方法 | 性能 | 可读性 | 可维护性 | 是否推荐 |
|---|---|---|---|---|
| group_by(...).map_groups(...) | ❌ 低（强制物化） | ⚠️ 中（Python 逻辑嵌套） | ❌ 差（难调试/优化） | 否 |
| sort + over("cik","form") + filter | ✅ 高（全向量化） | ✅ 高（声明式逻辑） | ✅ 优（易修改/复用） | 强烈推荐 |
该模式适用于所有「分组内 Top-N」场景（如每用户最近 3 笔订单、每品类销量前 5 商品），掌握 over 与 int_range 的组合，是写出高效 Polars 代码的关键一步。
