基于重置累计和分组：Polars中按文件大小阈值分组文件

需求：现有包含file_size字段的Polars DataFrame，需按照指定阈值（示例为3）对文件分组，每组的累计文件大小不得超过该阈值。分组规则为：依次累加文件大小，若加入当前文件后累计值超过阈值，则将该文件归入新组；否则归入当前组。

输入DataFrame示例 #

import polars as pl

data = {
    "bucket": ["bucket1"] * 8,
    "full_path": ["file1.txt","file2.txt","file3.txt","file4.txt","file5.txt","file6.txt","file7.txt","file8.txt"],
    "ETag": ["11c9d17dc657860c447e655fa79172f","21c9d17dc657860c447e655fa79172g","31c9d17dc657860c447e655fa79172f","41c9d17dc657860c447e655fa79172f","51c9d17dc657860c447e655fa79172f","61c9d17dc657860c447e655fa79172f","71c9d17dc657860c447e655fa79172f","81c9d17dc657860c447e655fa79172f"],
    "file_size": [1, 2, 2,2, 1, 3, 1, 2]
}
df = pl.DataFrame(data)
print(df)

输出结构：

shape: (8, 4)
┌─────────┬───────────┬─────────────────────────────────┬───────────┐
│ bucket  ┆ full_path ┆ ETag                            ┆ file_size │
│ ---     ┆ ---       ┆ ---                             ┆ ---       │
│ str     ┆ str       ┆ str                             ┆ i64       │
╞═════════╪═══════════╪═════════════════════════════════╪═══════════╡
│ bucket1 ┆ file1.txt ┆ 11c9d17dc657860c447e655fa79172f ┆ 1         │
│ bucket1 ┆ file2.txt ┆ 21c9d17dc657860c447e655fa79172g ┆ 2         │
│ bucket1 ┆ file3.txt ┆ 31c9d17dc657860c447e655fa79172f ┆ 2         │
│ bucket1 ┆ file4.txt ┆ 41c9d17dc657860c447e655fa79172f ┆ 2         │
│ bucket1 ┆ file5.txt ┆ 51c9d17dc657860c447e655fa79172f ┆ 1         │
│ bucket1 ┆ file6.txt ┆ 61c9d17dc657860c447e655fa79172f ┆ 3         │
│ bucket1 ┆ file7.txt ┆ 71c9d17dc657860c447e655fa79172f ┆ 1         │
│ bucket1 ┆ file8.txt ┆ 81c9d17dc657860c447e655fa79172f ┆ 2         │
└─────────┴───────────┴─────────────────────────────────┴───────────┘

预期输出（阈值为3时） #

shape: (8, 5)
┌─────────┬───────────┬─────────────────────────────────┬───────────┬──────────────┐
│ bucket  ┆ full_path ┆ ETag                            ┆ file_size ┆ group_number │
│ ---     ┆ ---       ┆ ---                             ┆ ---       ┆ ---          │
│ str     ┆ str       ┆ str                             ┆ i64       ┆ i64          │
╞═════════╪═══════════╪═════════════════════════════════╪═══════════╪══════════════╡
│ bucket1 ┆ file1.txt ┆ 11c9d17dc657860c447e655fa79172f ┆ 1         ┆ 1            │
│ bucket1 ┆ file2.txt ┆ 21c9d17dc657860c447e655fa79172g ┆ 2         ┆ 1            │
│ bucket1 ┆ file3.txt ┆ 31c9d17dc657860c447e655fa79172f ┆ 2         ┆ 2            │
│ bucket1 ┆ file4.txt ┆ 41c9d17dc657860c447e655fa79172f ┆ 2         ┆ 3            │
│ bucket1 ┆ file5.txt ┆ 51c9d17dc657860c447e655fa79172f ┆ 1         ┆ 3            │
│ bucket1 ┆ file6.txt ┆ 61c9d17dc657860c447e655fa79172f ┆ 3         ┆ 4            │
│ bucket1 ┆ file7.txt ┆ 71c9d17dc657860c447e655fa79172f ┆ 1         ┆ 5            │
│ bucket1 ┆ file8.txt ┆ 81c9d17dc657860c447e655fa79172f ┆ 2         ┆ 5            │
└─────────┴───────────┴─────────────────────────────────┴───────────┴──────────────┘

实现方案 #

方法1：循环遍历（适合小数据量） #

通过遍历每个文件大小，维护当前组的累计大小和组编号，逻辑直观易懂：

threshold = 3

# 添加行索引保证遍历顺序（若DataFrame已按需求排序可省略）
df = df.with_row_index()

current_sum = 0
current_group = 1
group_numbers = []

for size in df['file_size']:
    if current_sum + size > threshold:
        current_group += 1
        current_sum = size
    else:
        current_sum += size
    group_numbers.append(current_group)

# 添加分组编号并移除索引列
result_df = df.with_columns(pl.Series(name='group_number', values=group_numbers)).drop('index')
print(result_df)

方法2：Polars向量化操作（适合大数据量） #

利用Polars的fold函数实现纯向量化计算，避免Python循环，性能更优：

threshold = 3

result_df = df.with_columns(
    pl.fold(
        # 初始化累加器：当前组累计大小、当前组编号
        acc=pl.struct(current_sum=pl.lit(0), group=pl.lit(1)),
        # 定义累加逻辑
        function=lambda acc, x: pl.when(acc["current_sum"] + x > threshold)
        .then(pl.struct(current_sum=x, group=acc["group"] + 1))
        .otherwise(pl.struct(current_sum=acc["current_sum"] + x, group=acc["group"])),
        # 要处理的列
        exprs=[pl.col("file_size")]
    ).over(pl.int_range(0, pl.count()))  # 按行顺序执行累加
    .struct.rename_fields(["current_sum", "group_number"])  # 重命名字段
).drop("current_sum")  # 移除临时累计大小列

print(result_df)

两种方法均能得到符合要求的分组结果，方法2更适合处理大规模数据集，充分利用Polars的向量化性能优势。

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