如何将大型Polars DataFrame按列值拆分为指定数量的子DataFrame

问题背景 #

我有一个大型Polars DataFrame，希望根据数据规模将其拆分为n个自定义数量（例如2个、3个或5个）的子DataFrame。每个Identifier对应多条观测记录，之前的做法是按特定Identifier逐个拆分，但实际场景中Identifier数量庞大，无法手动指定，因此需要实现类似逻辑，将DataFrame拆分为近似均等的指定数量子DataFrame，同时尽量保证同一Identifier的记录不会被拆分到不同子集中。

方法1：按行索引直接拆分（快速高效，不保证Identifier完整性） #

如果不需要确保同一Identifier的所有记录在同一个子DataFrame中，这种方法最简单直接，直接按总行数均分：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    'Identifier': [1234,1234, 2345,2345, 3456,3456, 4567,4567],
    'DateColumn': ['2022-02-13','2022-02-14', '2022-02-13','2022-02-14',
                   '2022-02-13','2022-02-14', '2022-02-13','2022-02-14']
})

def split_df_by_rows(df, n_splits):
    total_rows = len(df)
    split_size = total_rows // n_splits
    splits = []
    start = 0
    # 处理行数无法整除的情况，前几个子集多分配一行
    for i in range(n_splits):
        end = start + split_size + (1 if i < total_rows % n_splits else 0)
        splits.append(df[start:end])
        start = end
    return splits

# 拆分为2个子DataFrame
split_dfs = split_df_by_rows(df, 2)
for i, sub_df in enumerate(split_dfs):
    print(f"子DataFrame {i+1}:")
    print(sub_df)

输出结果：

子DataFrame 1:
┌────────────┬────────────┐
│ Identifier ┆ DateColumn │
│ ---        ┆ ---        │
│ i64        ┆ str        │
╞════════════╪════════════╡
│ 1234       ┆ 2022-02-13 │
│ 1234       ┆ 2022-02-14 │
│ 2345       ┆ 2022-02-13 │
│ 2345       ┆ 2022-02-14 │
└────────────┴────────────┘
子DataFrame 2:
┌────────────┬────────────┐
│ Identifier ┆ DateColumn │
│ ---        ┆ ---        │
│ i64        ┆ str        │
╞════════════╪════════════╡
│ 3456       ┆ 2022-02-13 │
│ 3456       ┆ 2022-02-14 │
│ 4567       ┆ 2022-02-13 │
│ 4567       ┆ 2022-02-14 │
└────────────┴────────────┘

方法2：按Identifier分组后分配（保证同一Identifier在同一子集） #

如果需要确保同一个Identifier的所有记录都在同一个子DataFrame中，先统计每个Identifier的行数，再将Identifier分配到不同子集，确保每个子集的总数据量近似均等：

import polars as pl

def split_df_by_identifier(df, n_splits, id_col='Identifier'):
    # 统计每个Identifier的记录数
    id_counts = df.group_by(id_col).agg(pl.count().alias('count'))
    # 计算累计行数，用于分配分组
    id_counts = id_counts.with_columns(pl.cumsum('count').alias('cum_count'))
    
    total_rows = df.height
    # 计算每个子集的行数阈值
    split_thresholds = [total_rows * (i+1)/n_splits for i in range(n_splits)]
    
    # 给每个Identifier分配对应的组号
    group_assignments = []
    current_group = 0
    for _, row in id_counts.iter_rows(named=True):
        while current_group < n_splits and row['cum_count'] > split_thresholds[current_group]:
            current_group += 1
        group_assignments.append((row[id_col], current_group))
    
    # 创建组号映射字典
    group_map = dict(group_assignments)
    # 给原DataFrame添加组号列
    df_with_group = df.with_columns(pl.col(id_col).map_dict(group_map).alias('split_group'))
    # 按组拆分并移除组号列
    split_dfs = [df_with_group.filter(pl.col('split_group') == i).drop('split_group') for i in range(n_splits)]
    return split_dfs

# 拆分为2个子DataFrame
split_dfs = split_df_by_identifier(df, 2)
for i, sub_df in enumerate(split_dfs):
    print(f"子DataFrame {i+1}:")
    print(sub_df)

当Identifier行数不均等时，比如修改示例DataFrame：

df = pl.DataFrame({
    'Identifier': [1234,1234,1234, 2345,2345, 3456],
    'DateColumn': ['2022-02-13','2022-02-14','2022-02-15', 
                   '2022-02-13','2022-02-14', '2022-02-13']
})

拆分后会自动调整分组，保证同一Identifier不被拆分：

子DataFrame 1:
┌────────────┬────────────┐
│ Identifier ┆ DateColumn │
│ ---        ┆ ---        │
│ i64        ┆ str        │
╞════════════╪════════════╡
│ 1234       ┆ 2022-02-13 │
│ 1234       ┆ 2022-02-14 │
│ 1234       ┆ 2022-02-15 │
└────────────┴────────────┘
子DataFrame 2:
┌────────────┬────────────┐
│ Identifier ┆ DateColumn │
│ ---        ┆ ---        │
│ i64        ┆ str        │
╞════════════╪════════════╡
│ 2345       ┆ 2022-02-13 │
│ 2345       ┆ 2022-02-14 │
│ 3456       ┆ 2022-02-13 │
└────────────┴────────────┘

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