如何按symbol分组并按指定起始索引拆分Polars DataFrame

问题背景 #

现有如下Polars DataFrame：

import datetime
import polars as pl

df_orig = pl.DataFrame(
    {
        "symbol": [*["A"] * 10, *["B"] * 8],
        "date": [
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 10),
                eager=True,
            ),
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 8),
                eager=True,
            ),
        ],
        "data": [*range(10), *range(8)],
    }
)

df_helper = pl.DataFrame({"symbol": ["A", "B"], "start_idx": [[0, 5], [0, 3]]})

chunk_size = 5

with pl.Config(tbl_rows=30):
    print(df_orig)
    print(df_helper)

原始数据输出 #

df_orig内容：

shape: (18, 3)
┌────────┬─────────────────────┬──────┐
│ symbol ┆ date                ┆ data │
│ ---    ┆ ---                 ┆ ---  │
│ str    ┆ datetime[μs]        ┆ i64  │
╞════════╪═════════════════════╪══════╡
│ A      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ A      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ A      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ A      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ A      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ A      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ A      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ A      ┆ 2024-01-08 00:00:00 ┆ 7    │
│ A      ┆ 2024-01-09 00:00:00 ┆ 8    │
│ A      ┆ 2024-01-10 00:00:00 ┆ 9    │
│ B      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ B      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ B      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ B      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ B      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ B      ┆ 2024-01-08 00:00:00 ┆ 7    │
└────────┴─────────────────────┴──────┘

df_helper内容：

shape: (2, 2)
┌────────┬───────────┐
│ symbol ┆ start_idx │
│ ---    ┆ ---       │
│ str    ┆ list[i64] │
╞════════╪═══════════╡
│ A      ┆ [0, 5]    │
│ B      ┆ [0, 3]    │
└────────┴───────────┘

需求说明 #

按symbol列分组，将df_orig拆分为长度为chunk_size（值为5）的块：

• A组从行索引0和5开始，各拆出一个长度为5的块
• B组从行索引0和3开始，各拆出一个长度为5的块
  最终将所有块按行拼接，新增split_idx列标识拆分来源（同一序号对应不同symbol的同批次拆分块），预期输出如下：

shape: (20, 4)
┌───────────┬────────┬─────────────────────┬──────┐
│ split_idx ┆ symbol ┆ date                ┆ data │
│ ---       ┆ ---    ┆ ---                 ┆ ---  │
│ i64       ┆ str    ┆ datetime[μs]        ┆ i64  │
╞═══════════╪════════╪═════════════════════╪══════╡
│ 0         ┆ A      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ 0         ┆ A      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ 0         ┆ A      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ 0         ┆ A      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 0         ┆ A      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 0         ┆ B      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ 0         ┆ B      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ 0         ┆ B      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ 0         ┆ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 0         ┆ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 1         ┆ A      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ 1         ┆ A      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ 1         ┆ A      ┆ 2024-01-08 00:00:00 ┆ 7    │
│ 1         ┆ A      ┆ 2024-01-09 00:00:00 ┆ 8    │
│ 1         ┆ A      ┆ 2024-01-10 00:00:00 ┆ 9    │
│ 1         ┆ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 1         ┆ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 1         ┆ B      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ 1         ┆ B      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ 1         ┆ B      ┆ 2024-01-08 00:00:00 ┆ 7    │
└───────────┴────────┴─────────────────────┴──────┘

注意：start_idx列的列表长度决定每个分组的块数，不同symbol的列表长度可能不同。

解决方案 #

# 给df_orig添加分组内的行索引并按symbol聚合
df_with_group_idx = df_orig.with_row_index("group_idx").group_by("symbol", maintain_order=True).agg(
    pl.all()
)

# 关联df_helper，展开start_idx并生成split_idx，再按split_idx聚合
df_combined = df_with_group_idx.join(df_helper, on="symbol").explode("start_idx").with_row_index("split_idx", offset=0).group_by("split_idx", maintain_order=True).agg(
    pl.all()
)

# 提取每个split批次对应的chunk并合并
result = df_combined.with_columns(
    pl.struct("group_idx", "symbol", "date", "data", "start_idx").map_elements(
        lambda x: pl.DataFrame({
            "symbol": [x["symbol"]]*chunk_size,
            "date": x["date"][x["start_idx"]:x["start_idx"]+chunk_size],
            "data": x["data"][x["start_idx"]:x["start_idx"]+chunk_size]
        }),
        return_dtype=pl.DataFrame
    ).alias("chunk")
).explode("chunk").unnest("chunk").sort("split_idx")

with pl.Config(tbl_rows=30):
    print(result)

步骤解释 #

1. 添加分组内索引：给原始数据添加group_idx列记录每个symbol分组内的行位置，再按symbol聚合保留所有列。
2. 关联并拆分起始索引：将聚合后的数据与辅助表关联，展开start_idx列表并生成split_idx作为拆分批次序号，再按批次聚合。
3. 截取数据块：通过map_elements对每个批次，根据起始索引截取指定长度的数据块，生成子DataFrame。
4. 合并结果：展开所有子DataFrame，拆分嵌套列后按批次序号排序，得到最终结果。

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