group_by mean with error pyo3_runtime.PanicException: implementation error, cannot get ref Float64 from Float32

[mr-morj]

Checks

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest version of Polars.

Reproducible example

import os
os.environ["POLARS_VERBOSE"] = "1"

import polars as pl
from datetime import date, timedelta
import random
from typing import List, Callable


class SalesLagAggregatedFeature:
    name = "temp_feature"
    keys_list = ["date", "lagerId"]
    _fill_default_value = 0
    _col_name_to_lag = "cleaned_sales"
    _lag_index = 31
    _agg_func_name = "mean"
    _agg_func_map = {"sum": pl.sum, "mean": pl.mean, "median": pl.median}

    def _get_filtered_keys(self, skeleton_df: pl.LazyFrame) -> pl.LazyFrame:
        return skeleton_df.unique(subset=self.keys_list).select(self.keys_list)

    def _calc_lag(
        self, x: pl.LazyFrame, key_list_to_lag: List[str], col_to_lag: str
    ) -> pl.LazyFrame:
        lag_df = x.with_columns(
            pl.col("date").dt.offset_by(f"{self._lag_index}d"),
            pl.col(col_to_lag).alias(self.name),
        ).select(key_list_to_lag + [self.name])
        return pl.concat([x, lag_df], how="align")

    def _get_agg_func(self) -> Callable[[str], pl.Expr]:
        agg_func = self._agg_func_map.get(self._agg_func_name)
        if agg_func is None:
            raise ValueError(f"Unsupported aggregation name: {self._agg_func_name}")
        return agg_func

    def _agg_operation(self, df: pl.LazyFrame, col_to_agg: str) -> pl.LazyFrame:
        agg_func = self._get_agg_func()
        agg_df = df.group_by(self.keys_list).agg(agg_func(col_to_agg))
        return agg_df.sort(self.keys_list)

    def generate(self, skeleton: pl.LazyFrame, sales: pl.LazyFrame) -> None:
        
        lag_df = self._agg_operation(
            df=sales, col_to_agg=self._col_name_to_lag
        )
        lag_df = self._calc_lag(
            lag_df, key_list_to_lag=self.keys_list, col_to_lag=self._col_name_to_lag
        )

        unique_skeleton_df = self._get_filtered_keys(skeleton)
        self.generated_data = unique_skeleton_df.join(
            lag_df, how="left", on=self.keys_list
        )
        self.generated_data = self.generated_data.select(self.keys_list + [self.name])

    
    def fill(self) -> None:
        self.generated_data = self.generated_data.with_columns(
            pl.col(self.name).fill_null(self._fill_default_value)
        )


n_rows = 1623
start = date(2025, 1, 19)
end = date(2025, 2, 23)
date_pool = [start + timedelta(days=i) for i in range((end - start).days + 1)]

filial_pool = list(range(2000, 5000))
lager_pool  = [976320, 976322]
sales_pool  = [float(i) for i in range(6)]

data = {
    "date":          [random.choice(date_pool)   for _ in range(n_rows)],
    "filialId":      [random.choice(filial_pool) for _ in range(n_rows)],
    "lagerId":       [random.choice(lager_pool)  for _ in range(n_rows)],
    "cleaned_sales": [random.choice(sales_pool)  for _ in range(n_rows)],
}

schema = [
    ("date",          pl.Date),
    ("filialId",      pl.Int32),
    ("lagerId",       pl.Int32),
    ("cleaned_sales", pl.Float32),
]

df_sales = pl.DataFrame(data, schema=schema).lazy()
# print(df_sales.collect())


date_start = date(2025, 4, 19)
date_end   = date(2025, 5, 18)
date_range = [
    date_start + timedelta(days=i)
    for i in range((date_end - date_start).days + 1)
]
n_dates = len(date_range)

n_pairs = 169_380 // n_dates

filial_pool = list(range(2000, 5000))
lager_pool  = list(range(400_000, 500_000))

pairs = [
    (random.choice(filial_pool), random.choice(lager_pool))
    for _ in range(n_pairs)
]

rows = {
    "filialId": [],
    "lagerId":  [],
    "date":     [],
}
for filial, lager in pairs:
    for d in date_range:
        rows["filialId"].append(filial)
        rows["lagerId"].append(lager)
        rows["date"].append(d)

schema = [
    ("filialId", pl.Int32),
    ("lagerId",  pl.Int32),
    ("date",     pl.Date),
]

df_skeleton = pl.DataFrame(rows, schema=schema).lazy()
# print(df_skeleton.collect())

feature_instance = SalesLagAggregatedFeature()
feature_instance.generate(skeleton=df_skeleton, sales=df_sales)
# OK
# print(feature_instance.generated_data.collect())

feature_instance.fill()
#  ERROR
res = feature_instance.generated_data.collect()

Log output

found multiple sources; run comm_subplan_elim
join parallel: false
join parallel: false
estimated unique values: 434
run PARTITIONED HASH AGGREGATION
CACHE SET: cache id: 0
CACHE HIT: cache id: 0
FULL join dataframes finished
LEFT join dataframes finished

thread '<unnamed>' panicked at crates/polars-core/src/series/mod.rs:1017:13:
implementation error, cannot get ref Float64 from Float32
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Traceback (most recent call last):
  File "/home/jovyan/L3IBaselineForecastService/temp.py", line 133, in <module>
    res = feature_instance.generated_data.collect()
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/polars/lazyframe/frame.py", line 2056, in collect
    return wrap_df(ldf.collect(callback))
                   ^^^^^^^^^^^^^^^^^^^^^
pyo3_runtime.PanicException: implementation error, cannot get ref Float64 from Float32

Issue description

I am performing some actions on lazyframe. At the input for the target field, I have data type f32. Then after aggregation group_by+mean I get f64. At the same time, intermediate collect() are executed and I can print dataframes. At the end I execute fill_null(0) and I get an error during the collection after that: pyo3_runtime.PanicException: implementation error, cannot get ref Float64 from Float32.

If I make some other subset of data, then after aggregation the data type f32 is preserved and no error occurs.

Expected behavior

Saved type f32 after group_by+mean

Installed versions

--------Version info---------
Polars:              1.20.0
Index type:          UInt32
Platform:            Linux-6.1.131-143.221.amzn2023.x86_64-x86_64-with-glibc2.35
Python:              3.11.11 | packaged by conda-forge | (main, Dec  5 2024, 14:17:24) [GCC 13.3.0]
LTS CPU:             False

----Optional dependencies----
Azure CLI            <not installed>
adbc_driver_manager  <not installed>
altair               <not installed>
azure.identity       <not installed>
boto3                1.36.16
cloudpickle          <not installed>
connectorx           <not installed>
deltalake            0.25.4
fastexcel            <not installed>
fsspec               <not installed>
gevent               <not installed>
google.auth          2.38.0
great_tables         <not installed>
matplotlib           3.10.0
nest_asyncio         1.6.0
numpy                2.1.2
openpyxl             <not installed>
pandas               2.2.3
pyarrow              17.0.0
pydantic             2.9.2
pyiceberg            <not installed>
sqlalchemy           2.0.40
torch                <not installed>
xlsx2csv             <not installed>
xlsxwriter           <not installed>
None

[ritchie46]

Can you make a reproducible example? We cannot do anything if we cannot reproduce it.

[mr-morj]

@ritchie46 Sorry, I updated it.

[cmdlineluser]

Attempted minimal repro.

import polars as pl

N = 1000

df1 = pl.select(
    x=pl.repeat(1, N // 2).append(pl.repeat(2, N // 2)).shuffle(),
    y=pl.lit(3.0, pl.Float32),
).lazy()

df2 = pl.select(x=pl.repeat(4, N)).lazy()

(
    df2.join(df1.group_by("x").mean().with_columns(z="y"), how="left", on="x")
    .with_columns(pl.col("z").fill_null(0))
    .collect()
)
# PanicException: implementation error, cannot get ref Float64 from Float32

If I lower N it doesn't error, so N may need to be increased depending on your system.