如何用Polars将DataFrame所有组统一为指定大小并填充null

针对Polars 1.1.0版本，这里提供两种实现思路，都能完成将分组后各组统一填充到指定大小（缺失值用null补充）的需求：

方法一：分组后直接扩展（简洁直观） #

通过group_by结合apply，对每个分组单独扩展至目标长度，不足部分用None填充：

import polars as pl

def make_groups_uniform(df: pl.DataFrame, group_by: str, group_size: int) -> pl.DataFrame:
    return df.group_by(group_by, maintain_order=True).apply(
        lambda group: pl.DataFrame({
            group_by: [group[group_by][0]] * group_size,
            "value": group["value"].to_list() + [None] * (group_size - len(group))
        })
    )

# 测试示例
dfa = pl.DataFrame(data={'group': ['a', 'a', 'a', 'b', 'b', 'c'],
                         'value': ['a1', 'a2', 'a3', 'b1', 'b2', 'c1']})
print(make_groups_uniform(dfa, group_by='group', group_size=4))

运行后会输出你期望的结果，maintain_order=True保证分组顺序和原数据一致。

方法二：基于行号连接（高效适用于大数据集） #

如果处理的数据集较大，apply的性能可能不足，可通过生成分组的完整行号范围，再与原数据左连接实现：

import polars as pl

def make_groups_uniform(df: pl.DataFrame, group_by: str, group_size: int) -> pl.DataFrame:
    # 生成每个分组的所有目标行号（0到group_size-1）
    group_full_rows = df.select(group_by).unique(maintain_order=True).with_columns(
        pl.int_range(0, group_size).alias("row_num").explode()
    )
    # 给原数据添加组内行号
    df_with_row = df.with_row_index().group_by(group_by).agg(
        pl.col("value").sort_by("index").alias("value"),
        pl.int_range(0, pl.count()).alias("row_num")
    ).explode(["value", "row_num"])
    # 左连接补全缺失值
    return group_full_rows.join(df_with_row, on=[group_by, "row_num"], how="left").drop("row_num")

# 测试示例
dfa = pl.DataFrame(data={'group': ['a', 'a', 'a', 'b', 'b', 'c'],
                         'value': ['a1', 'a2', 'a3', 'b1', 'b2', 'c1']})
print(make_groups_uniform(dfa, group_by='group', group_size=4))

这种方法利用Polars的向量化操作，避免了Python层面的循环，性能更优。

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