[Python] Polars maximum length reached when rolling with list aggregation

I'm working with time series data of daily historic demands for a lot of stores/products with a dataframe size of around 27.000.000 rows and 82.000 individual time series (specified by id_store and id_product). For a previous dataset of way smaller size I built this moving average logic as a 2-step forecast baseline (first looking at the same weekdays and then fill the gaps by looking at the whole time series): 1st rolling operation:

baseline_forecast = (df.sort("target_date").rolling(
index_column="target_date",
period="52w",
offset="-53w",
group_by=["id_store", "id_product", "day_of_week"],
closed="both",
)
# aggregate into list and calculate mean of last 4 observations
.agg(
pl.col("value").alias("baseline_forecast"),
)
.with_columns(
pl.col("baseline_forecast").list.drop_nulls(),
)
.with_columns(
# in case of only one data point per weekday this would lead to
# eternal usage of this one value without regarding the current
# general sales (not weekday specific) -> in this case use fallback naive
pl.when(pl.col("baseline_forecast").list.len() >= 2)
.then(
pl.col("baseline_forecast")
.list.slice(offset=-4)
.list.mean()
.fill_nan(None)
)
.otherwise(None)
.alias("baseline_forecast"),
)
)


2nd baseline as a fallback without respecting weekdays → panics

baseline_fallback = (df.sort("target_date")
.rolling(
index_column="target_date",
period="52w",
offset="-53w", # for the start of the period
group_by=["id_store", "id_product"],
closed="both",
)
# aggregate into list and calculate mean of last 4 observations
.agg(
pl.col("value").alias("baseline_forecast_fallback"),
)
.with_columns(
pl.col("baseline_forecast_fallback")
.list.drop_nulls()
.list.slice(offset=-15)
.list.mean()
.fill_nan(None) # polars considers mean of empty lists as NaN
.alias("baseline_forecast_fallback"),
).select("id_store", "id_product", "target_date", "baseline_forecast_fallback")


So this works totally fine for my smaller dataset but for the current dataset the 2nd baseline fallback solution runs in an pyo3_runtime.PanicException: Polars' maximum length reached. Consider installing 'polars-u64-idx'..

Installing polars-u64-idx removes the panic, but the resulting DataFrame still has the same ≈27 M rows, which makes me wonder what internal buffer actually overflows in the regular build.

I would like to understand why this is happening and if my guess is correct. In the second rolling operation I'm looking back a year (if available) for each of my 27 million rows, so my guess is that polars internally builds up an huge amount of data of approximately 27.000.000 x ~300 days size (store closing days subtracted) and thus is crashing - am I correct?

Continue reading...