Возьмите N элементов окна в полярах ⇐ Python

[Anonymous]

Скажем, у меня есть это:

Код: Выделить всё

import numpy as np
import polars as pl

pl.Config(tbl_rows=20) # to show full output

df = (pl
.DataFrame(dict(
j=np.random.randint(10, 99, 20),
))
.with_row_index()
.select(
g=pl.col('index') // 4,
j='j',
)
)

Код: Выделить всё

shape: (20, 2)
┌─────┬─────┐
│ g   ┆ j   │
│ --- ┆ --- │
│ u32 ┆ i64 │
╞═════╪═════╡
│ 0   ┆ 95  │
│ 0   ┆ 80  │
│ 0   ┆ 51  │
│ 0   ┆ 68  │
│ 1   ┆ 71  │
│ 1   ┆ 92  │
│ 1   ┆ 44  │
│ 1   ┆ 97  │
│ 2   ┆ 36  │
│ 2   ┆ 64  │
│ 2   ┆ 70  │
│ 2   ┆ 80  │
│ 3   ┆ 75  │
│ 3   ┆ 69  │
│ 3   ┆ 54  │
│ 3   ┆ 16  │
│ 4   ┆ 88  │
│ 4   ┆ 89  │
│ 4   ┆ 97  │
│ 4   ┆ 37  │
└─────┴─────┘

Я хотел бы взять два верхних элемента в каждой группе, т.е. в итоге получить следующее:

Код: Выделить всё

 shape: (10, 2)
┌─────┬─────┐
│ g   ┆ j   │
│ --- ┆ --- │
│ u32 ┆ i64 │
╞═════╪═════╡
│ 0   ┆ 95  │
│ 0   ┆ 80  │
│ 1   ┆ 71  │
│ 1   ┆ 92  │
│ 2   ┆ 36  │
│ 2   ┆ 64  │
│ 3   ┆ 75  │
│ 3   ┆ 69  │
│ 4   ┆ 88  │
│ 4   ┆ 89  │
└─────┴─────┘

Я пробовал это:

Код: Выделить всё

dfj = (df
.select(
pl.all().head(2).over('g')
)
)
print(dfj)

но это приводит к такому исключению:

Код: Выделить всё

ComputeError: the length of the window expression did not match that of the group

Error originated in expression: 'col("g").slice(offset=0, length=2).over([col("g")])'

Я знаю, что могу снова добавить количество строк для каждой группы и фильтровать по этому значению, но мне интересно:

• Почему здесь не работает head
• Если есть лучшее решение (особенно без group_by)

Подробнее здесь: https://stackoverflow.com/questions/768 ... -in-polars