Digamos que eu tenha:
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],
}
Gostaria de adicionar uma coluna com uma média móvel com tamanho de janela 2, com min_periods=2
, sobre'id'
Em Polars, eu posso fazer:
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 │
└─────┴─────┴───────┴───────────────┘
Qual é o equivalente do DuckDB? Eu tentei
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')
mas pegue
┌─────────┬───────┬───────┬───────────────┐
│ 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 │
└─────────┴───────┴───────┴───────────────┘
Isso é muito próximo, mas o duckdb ainda calcula a média móvel quando há apenas um único valor na janela. Como posso replicar o min_periods=2
comportamento (padrão) do Polars?
Você pode usar
case
declaração ecount
:Observe que usei enquadramento de janela e linha nomeado aqui.