每个窗口内的行数

Question

我有一些表达式，稍后我将在窗口函数内或不使用窗口函数时对其进行评估。这通常工作正常。有

pl.col("x").max()

——稍后添加

.over("y")

。有

pl.arange(0, pl.count())

——稍后添加

.over("y")

。一种不起作用的表达方式是

pl.count()

。

如果您尝试打开窗口

pl.count()

，Polars 错误：

import polars as pl

df = pl.DataFrame(dict(x=[1,1,0,0], y=[1,2,3,4]))
expression = pl.count()

df.with_columns([expression.over("x").alias("z")])
# exceptions.ComputeError: Cannot apply a window function, did not find a root column. This is likely due to a syntax error in this expression: count()

有没有可以处理窗口化的

count

版本？我知道我可以做到

pl.col("x").count().over("x")

，但是我必须提前知道将存在哪些列，并且表达式和窗口列来自代码中完全不同的部分。

Answer 1

目前

window

表达式的限制是它们需要有根列。这意味着在

Dataframe

上下文中启动表达式的列。

pl.count()

是一个不引用任何列的表达式，并且对任何列都是通用的。

我们可以轻松规避此限制，因为知道

column.len() == count()

。所以我们可以只取

first()

中的

DataFrame

列。

df = pl.DataFrame(dict(x=[1,1,0,0], y=[1,2,3,4]))
expression = pl.first().len()

df.with_columns([expression.over("x").alias("z")])

shape: (4, 3)
┌─────┬─────┬─────┐
│ x   ┆ y   ┆ z   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ u32 │
╞═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 2   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 1   ┆ 2   ┆ 2   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 0   ┆ 3   ┆ 2   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 0   ┆ 4   ┆ 2   │
└─────┴─────┴─────┘

Answer 2

升级至 Polars >=0.14。从该版本开始，原始问题中的行为无需修改即可开始工作。

import polars as pl

df = pl.DataFrame(dict(x=[1,1,0,0], y=[1,2,3,4]))
expression = pl.count()

df.with_columns([expression.over("x").alias("z")])
# shape: (4, 3)
# ┌─────┬─────┬─────┐
# │ x   ┆ y   ┆ z   │
# │ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ u32 │
# ╞═════╪═════╪═════╡
# │ 1   ┆ 1   ┆ 2   │
# │ 1   ┆ 2   ┆ 2   │
# │ 0   ┆ 3   ┆ 2   │
# │ 0   ┆ 4   ┆ 2   │
# └─────┴─────┴─────┘

Answer 3

不确定您是否需要它成为一个窗口，但根据您给出的示例，一个简单的

group_by

就足够了

import polars as pl

df = pl.DataFrame(dict(x=[1,1,0,0], y=[1,2,3,4]))
expression = pl.len()

print(df.group_by(pl.col("x")).agg(expression))
# print(df.select(expression])
# ^this would work on the whole dataframe