Разделите фрейм данных Polars на несколько частей с помощью groupby ⇐ Python

[Anonymous]

Рассмотрим следующие 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

), сгруппированные по столбцу символа. Столбец start_idx указывает строки, с которых начинается фрагмент в каждой группе. То есть группа A будет разделена на два фрагмента длиной 5, начиная с строк 0 и 5, а фрагменты группы B начинаются со строк 0 и 3.
Наконец, все фрагменты необходимо объединить по оси = 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 может иметь переменную длину для каждой отдельной строки. Длина каждого списка определяет количество фрагментов для каждой группы.

Подробнее здесь: https://stackoverflow.com/questions/788 ... th-groupby