主页 / coding / 问题 / 78842466
Accepted
Andi
Asked: 2024-08-07 15:52:59 +0800 CST

使用 groupby 将 Polars DataFrame 拆分为多个块

  • 772

考虑以下pl.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, 4]]})

chunk_size = 5

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


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    │
└────────┴─────────────────────┴──────┘

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

现在，我需要将数据框拆分为两个长度为 5 的块（chunk_size），并按symbol列分组。 列start_idx指示每个组中块的起始行。 也就是说，组 A 将被拆分为两个长度为 5 的块，从第 0 行和第 5 行开始，而组 B 的块从第 0 行和第 3 行开始。 最后，所有块都需要连接起来axis=0，从而新列split_idx指示拆分来自何处。

以下是我所寻找的内容：

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对于每一行来说长度可能都是可变的。每个列表的长度决定了每个组的块数。

python-polars

2 个回答

  1. Best Answer

    这是一个完全符合极坐标表达式 API 的解决方案。

    主要思想是将辅助数据框预处理为symbol、split_idx和的数据框row_idx。这里，row_idx是符号和拆分索引定义的组内行的索引。

    它可以作为一个“骨架”，我们可以（在向中添加这样的行索引后df_orig）轻松地使用它与进行左合并df_orig。

    pl.Config().set_tbl_rows(-1)

def preprocess_helper(df_helper: pl.DataFrame) -> pl.DataFrame:
    return (
        df_helper
        .explode("start_idx")
        .with_columns(
            pl.int_range(pl.len()).over("symbol").alias("split_idx"),
            pl.int_ranges(pl.col("start_idx"), pl.col("start_idx") + chunk_size).alias("row_idx"),
        )
        .explode("row_idx")
    )

(
    preprocess_helper(df_helper)
    .join(
        df_orig.with_columns(pl.int_range(pl.len()).over("symbol").alias("row_idx")),
        on=["symbol", "row_idx"],
        how="left",
    )
    .drop("row_idx", "start_idx")
    .sort("split_idx", "symbol")
)

    注意：如果输出的确切列/顺序无关紧要，则可以省略最后的pl.DataFrame.drop/ 。pl.DataFrame.sort

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

  2. 我无法提出一个完全本机的解决方案，因为polars.Expr.slice它似乎不支持每行不同的值。

    以下是使用混合 Python 代码实现此目的的一种方法：

    首先，df_orig按符号分组并将其放入字典中。

    然后，迭代df_helper并提取切片并将它们放入列表中。

    最后，将数据框连接成单个数据框并对其进行排序。

    df_orig_by_symbol = {
    k: v for ((k,), v) in df_orig.group_by("symbol", maintain_order=True)
}

dfs = []

for symbol, start_idxs in df_helper.select("symbol", "start_idx").iter_rows():
    for split_idx, start_idx in enumerate(start_idxs):
        dfs.append(
            df_orig_by_symbol[symbol]
            .slice(start_idx, chunk_size)
            .with_columns(split_idx=split_idx)
        )

final_df = pl.concat(dfs).sort("split_idx", "symbol")

    输出：

    shape: (20, 4)
┌────────┬─────────────────────┬──────┬───────────┐
│ symbol ┆ date                ┆ data ┆ split_idx │
│ ---    ┆ ---                 ┆ ---  ┆ ---       │
│ str    ┆ datetime[μs]        ┆ i64  ┆ i32       │
╞════════╪═════════════════════╪══════╪═══════════╡
│ 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    ┆ 0         │
│ 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    ┆ 1         │
└────────┴─────────────────────┴──────┴───────────┘

    完整代码：

    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

df_orig_by_symbol = {
    k: v for ((k,), v) in df_orig.group_by("symbol", maintain_order=True)
}

dfs = []

for symbol, start_idxs in df_helper.select("symbol", "start_idx").iter_rows():
    for split_idx, start_idx in enumerate(start_idxs):
        dfs.append(
            df_orig_by_symbol[symbol]
            .slice(start_idx, chunk_size)
            .with_columns(split_idx=split_idx)
        )

final_df = pl.concat(dfs).sort("split_idx", "symbol")

with pl.Config(tbl_rows=100):
    print(final_df)
    • 0

相关问题