假设我有:
data = {
'id': ['a', 'a', 'a', 'b', 'b', 'b', 'b'],
'd': [1,2,3,0,1,2,3],
'sales': [5,1,3,4,1,2,3],
}
我想添加一个带有滚动平均值的列,窗口大小为 2 min_periods=2
,'id'
在 Polars 中,我可以执行以下操作:
import polars as pl
df = pl.DataFrame(data)
df.with_columns(sales_rolling = pl.col('sales').rolling_mean(2).over('id'))
shape: (7, 4)
┌─────┬─────┬───────┬───────────────┐
│ id ┆ d ┆ sales ┆ sales_rolling │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ f64 │
╞═════╪═════╪═══════╪═══════════════╡
│ a ┆ 1 ┆ 5 ┆ null │
│ a ┆ 2 ┆ 1 ┆ 3.0 │
│ a ┆ 3 ┆ 3 ┆ 2.0 │
│ b ┆ 0 ┆ 4 ┆ null │
│ b ┆ 1 ┆ 1 ┆ 2.5 │
│ b ┆ 2 ┆ 2 ┆ 1.5 │
│ b ┆ 3 ┆ 3 ┆ 2.5 │
└─────┴─────┴───────┴───────────────┘
DuckDB 的对应产品是什么?我试过
import duckdb
duckdb.sql("""
select
*,
mean(sales) over (
partition by id
order by d
range between 1 preceding and 0 following
) as sales_rolling
from df
""").sort('id', 'd')
但得到
┌─────────┬───────┬───────┬───────────────┐
│ id │ d │ sales │ sales_rolling │
│ varchar │ int64 │ int64 │ double │
├─────────┼───────┼───────┼───────────────┤
│ a │ 1 │ 5 │ 5.0 │
│ a │ 2 │ 1 │ 3.0 │
│ a │ 3 │ 3 │ 2.0 │
│ b │ 0 │ 4 │ 4.0 │
│ b │ 1 │ 1 │ 2.5 │
│ b │ 2 │ 2 │ 1.5 │
│ b │ 3 │ 3 │ 2.5 │
└─────────┴───────┴───────┴───────────────┘
这非常接近,但当窗口中只有一个值时,duckdb 仍会计算滚动平均值。我如何复制min_periods=2
Polars 的(默认)行为?
您可以使用
case
语句和count
:注意我在这里使用了命名窗口和行框架。