如何在Polars中实现返回多特征的自定义函数？

我正在评估使用Polars进行时序特征提取，以复刻tsfresh等类库的功能。目前基础功能已正常运行：通过groupby_dynamic创建窗口，将tsfel的多数基础特征直接重写为Polars自定义函数（当前性能尚可，暂无需优化），示例代码如下：

def var_larger_than_std(x: pl.Series) -> bool:
    """
    判断方差是否大于标准差（等价于方差大于1）
    """
    y = x.var()
    return y > np.sqrt(y) 

q = (
    dataset.lazy() \
    .groupby_dynamic("ts_local", every="1d", by="id") \
    .agg([
        pl.col("value").count().alias("value__count"),
        var_larger_than_std(pl.col("value")).alias("value__var_larger_than_std"),
        ])
)

q.collect()

对于单输出的特征映射（如f(Series)->bool/数值），上述方式运行良好。但复杂特征（如FFT系数）通常会生成多个输出，在tsfel中这类结果以dict形式返回，并展开为DataFrame的独立列。现需在Polars中实现类似功能，示例需求如下：

自定义多输出函数：

from typing import Dict

def complex_feature(x: pl.Series) -> Dict[str, float]:
    return {"col1":0.0, "col2":1.0}

期望将函数返回的dict展开为独立列，输出格式如下：

┌────────────────────────────────────┬──────┬──────┐
│ ts_local                           ┆ col1 ┆ col2 │
│ --------                           ┆ ---- ┆ ---- │
│ timestamp[ns,tz=Australia/Sydney]  ┆ f64  ┆ f64  │
╞════════════════════════════════════╪══════╪══════╡
│ 2023-01-01                         ┆ 0.0  ┆ 1.0  │
└────────────────────────────────────┴──────┴──────┘

解决方案 #

方法1：利用Polars Struct包装后展开 #

这是最直观且符合Polars设计思路的方式，将多输出结果包装为pl.Struct，再通过unnest()展开为独立列。

方式A：修改函数返回Struct

def complex_feature(x: pl.Series) -> pl.Struct:
    return pl.struct({"col1": 0.0, "col2": 1.0})

q = (
    dataset.lazy()
    .groupby_dynamic("ts_local", every="1d", by="id")
    .agg([
        pl.col("value").map(complex_feature).alias("complex_features")
    ])
    .unnest("complex_features")  # 将Struct拆分为独立列
)
q.collect()

方式B：保留原dict返回，在查询中转换为Struct

如果不想修改原函数，可以在map中直接将dict转为Struct：

q = (
    dataset.lazy()
    .groupby_dynamic("ts_local", every="1d", by="id")
    .agg([
        pl.col("value").map(lambda s: pl.struct(complex_feature(s))).alias("complex_features")
    ])
    .unnest("complex_features")
)
q.collect()

方法2：返回元组后拆分列 #

如果函数返回元组，可以通过list.get()提取每个元素并命名为列，适合输出数量较少的场景：

def complex_feature(x: pl.Series) -> tuple[float, float]:
    return (0.0, 1.0)

q = (
    dataset.lazy()
    .groupby_dynamic("ts_local", every="1d", by="id")
    .agg([
        pl.col("value").map(complex_feature).alias("temp")
    ])
    .with_columns([
        pl.col("temp").list.get(0).alias("col1"),
        pl.col("temp").list.get(1).alias("col2")
    ])
    .drop("temp")
)
q.collect()

方法3：使用map_batches批量处理（适合高性能场景） #

如果特征计算可以批量优化，可让函数返回pl.DataFrame，利用map_batches处理分组后的批量数据，性能更优：

def complex_feature_batch(x: pl.Series) -> pl.DataFrame:
    # x为分组后的Series，返回的DataFrame需与分组数量匹配
    return pl.DataFrame(
        {"col1": [0.0]*len(x), "col2": [1.0]*len(x)}
    )

q = (
    dataset.lazy()
    .groupby_dynamic("ts_local", every="1d", by="id")
    .agg([
        pl.col("value").map_batches(complex_feature_batch)
    ])
)
q.collect()

推荐方案 #

优先选择方法1，Struct的方式既符合Polars的结构化数据模型，又能简洁地实现多输出列的展开，代码可读性和维护性最佳。

内容的提问来源于stack exchange，提问作者David Waterworth