Polars分组内过滤优化咨询：按规则保留指定表单数据

你的需求可以通过Polars的窗口函数实现，相比原代码的两次分组+map_groups，窗口函数是向量化操作，对大型DataFrame的性能提升明显，同时代码更简洁：

优化后代码（推荐版本） #

def filter_sub_for_11_rows_per_cik(df_):
    return (
        df_
        # 按cik+form分组，组内按period降序生成行号（最新的记录行号为1）
        .with_row_number(
            name='row_num',
            over=['cik', 'form'],
            order_by=pl.col('period').desc()
        )
        # 筛选条件：10-Q留最新1条，10-K留最新最多10条
        .filter(
            (pl.col('form') == '10-Q') & (pl.col('row_num') == 1)
            | (pl.col('form') == '10-K') & (pl.col('row_num') <= 10)
        )
        # 删掉临时行号列
        .drop('row_num')
    )

代码说明 #

1. 窗口行号生成：with_row_number配合over=['cik', 'form']，在每个cik和form的分组内，按period降序生成从1开始的行号，最新的记录会被标记为行号1。
2. 精准筛选：直接通过条件过滤出需要的行，无需拆分DataFrame再合并，逻辑连贯。
3. 性能优势：避免了原代码中map_groups的逐组lambda执行（这对大型数据集性能影响很大），全程是Polars的向量化操作，内存占用更低、运行速度更快。

另一种排序后筛选的版本 #

如果你偏好先全局排序的方式，也可以用下面的写法，逻辑更直观：

def filter_sub_for_11_rows_per_cik(df_):
    return (
        df_
        # 按cik升序、form升序、period降序全局排序
        .sort(['cik', 'form', 'period'], descending=[False, False, True])
        # 给每个cik+form组生成行号
        .with_row_index(name='row_num', offset=1, over=['cik', 'form'])
        .filter(
            (pl.col('form') == '10-Q') & (pl.col('row_num') == 1)
            | (pl.col('form') == '10-K') & (pl.col('row_num') <= 10)
        )
        .drop('row_num')
    )

这个版本先全局排序，再给每个分组生成行号，效果和上面的版本一致，但对于超大型DataFrame，窗口内排序的版本性能会更优。

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