Polars中group_by_dynamic自定义gmean返回list[f64]的问题

在使用Polars的group_by_dynamic()执行滚动计算时，自定义几何均值表达式.gmean()返回list[f64]类型，而非预期的f64；但在.lazy().select中使用该表达式可得到正确的f64结果。

正常返回f64的代码及输出 #

let df = self
            .lazy()
            .select([cols(apply_to).gmean()]
                )
            .collect()
            .expect(" Error implementing mean returns df");
            Ok(df)

输出：

┌──────────┬───────────┐
│ eer      ┆ pegy      │
│ ---      ┆ ---       │
│ f64      ┆ f64       │
╞══════════╪═══════════╡
│ 0.000909 ┆ -0.000442 │
└──────────┴───────────┘

返回list[f64]的group_by_dynamic代码及输出 #

let df: DataFrame = self
            .lazy()
            .group_by_dynamic(
                col("date"), 
                [], 
                DynamicGroupOptions {
                    every: Duration::parse(&window_len),
                    period: Duration::parse(&window_len),
                    offset: Duration::parse("0"),
                    ..Default::default()
                })
            .agg([cols(apply_to).gmean()])
            .collect()
            .expect("Error implementing rolling geometric mean returns df");
        Ok(df)

输出：

shape: (4, 3)
┌────────────┬─────────────┬─────────────┐
│ date       ┆ eer         ┆ pegy        │
│ ---        ┆ ---         ┆ ---         │
│ date       ┆ list[f64]   ┆ list[f64]   │
╞════════════╪═════════════╪═════════════╡
│ 2024-11-04 ┆ [0.001432]  ┆ [-0.001203] │
│ 2024-11-07 ┆ [0.002497]  ┆ [0.0003]    │
│ 2024-11-10 ┆ [0.001454]  ┆ [-0.001251] │
│ 2024-11-13 ┆ [-0.001033] ┆ [0.000366]  │
└────────────┴─────────────┴─────────────┘

自定义gmean表达式实现 #

fn geometric_mean(values: &Float64Chunked) -> f64 {
    let adjusted_values: Float64Chunked = values.apply(|opt_v| opt_v.map(|x| x + 1.0));
    let product: f64 = adjusted_values
        .into_iter()
        .filter_map(|opt| opt) // Remove None values
        .product(); // Compute the product of present values
    let count = adjusted_values.len() as f64;
    product.powf(1.0 / count)
}

fn gmean(series: &[Series]) -> PolarsResult<Series> {
    let _series = &series[0];
    let _chunk_array = _series.f64()?;
    let geo_mean= geometric_mean(_chunk_array) - 1.0;
    let new_chunk = Float64Chunked::from_vec(_series.name().clone(), vec![geo_mean]);
    Ok(new_chunk.into_series())
}

fn gmean_column(series: &Column) -> Result<Option<Column>,PolarsError> {
    let materialized_series_slice = std::slice::from_ref(series.as_materialized_series());
    Ok(Some(gmean(materialized_series_slice)?.into_column()))
}

pub trait GeoMean {
    fn gmean(self) -> Expr;
}

impl GeoMean for Expr {
    fn gmean(self) -> Expr {
        self.apply(|column| gmean_column(&column),GetOutput::from_type(DataType::Float64))
    }
}

问题解答 #

1. 为何.gmean()在group_by_dynamic中返回list[f64]？ #

Polars的agg方法默认支持多值聚合，会将每个分组的计算结果包装为列表。你的自定义gmean实现返回的是单元素Series，agg直接将这个Series作为列表存入结果列，最终形成list[f64]类型。

而在.select场景中，Polars处理的是全数据集计算，返回的单元素Series会直接展开为标量列，因此显示为f64类型。

2. 如何修改表达式使其返回f64？ #

调整自定义函数的实现逻辑，适配Polars聚合函数的要求（返回标量而非Series）：

// 重写为聚合函数，返回Scalar
fn gmean_agg(s: &Series) -> PolarsResult<Scalar> {
    let chunk = s.f64()?;
    let adjusted_values = chunk.apply(|opt_v| opt_v.map(|x| x + 1.0));
    let product: f64 = adjusted_values
        .into_iter()
        .filter_map(|opt| opt)
        .product();
    let count = adjusted_values.len() as f64;
    let geo_mean = product.powf(1.0 / count) - 1.0;
    Ok(Scalar::Float64(Some(geo_mean)))
}

// 更新trait实现，使用agg而非apply
impl GeoMean for Expr {
    fn gmean(self) -> Expr {
        self.agg(gmean_agg, GetOutput::from_type(DataType::Float64))
    }
}

核心改动：用.agg()替代.apply()定义聚合逻辑，直接返回标量值，这样group_by_dynamic().agg()会将结果解析为f64列。

3. 是否有方法提取列表中的值？ #

如果暂时不想修改自定义函数，可以直接提取列表中的唯一元素：

.agg([cols(apply_to).gmean()])
// 提取列表第一个元素并转为f64
.with_columns([
    col("eer").list().first().alias("eer"),
    col("pegy").list().first().alias("pegy")
])

或者用.get(0)实现相同效果：

.with_columns([
    col("eer").list().get(0).cast(DataType::Float64).alias("eer"),
    col("pegy").list().get(0).cast(DataType::Float64).alias("pegy")
])

该方法适用于每个列表仅含一个元素的场景，可快速将list[f64]转为f64类型。

内容的提问来源于stack exchange，提问作者Trevor Seibert