如何在Polars中生成List[f64]浮点序列并修复边界问题

问题背景 #

需要基于Polars的List[f64]类型列"a"，生成新的List[f64]类型列"b"，要求每行的"b"序列是从"a"列表的最小值到最大值、间隔0.5且包含首尾的连续序列。原解决方案在"a"列表仅含唯一值的场景下失效，且需要保证性能。

核心解决方案 #

通过数值缩放避开Polars arange()仅支持整数步长的限制，同时自然处理唯一值边界场景，代码如下：

import polars as pl

df = df.with_columns(
    pl.col("a").list.eval(
        pl.arange(
            # 将最小值乘以2转为整数，作为序列起始
            (pl.element().min() * 2).cast(pl.Int64),
            # 将最大值乘以2加1转为整数，保证包含最大值
            (pl.element().max() * 2).cast(pl.Int64) + 1,
            # 整数步长1，对应原间隔0.5
            1
        ).cast(pl.Float64) / 2,  # 缩放回原数值范围
        parallel=True  # 开启并行计算保证性能
    ).alias("b")
)

方案优势 #

• 边界场景处理：当"a"列表仅含唯一值时，min等于max，生成的序列仅包含该唯一值，符合预期
• 性能高效：利用list.eval的并行计算能力，避免冗余的append/unique/sort操作，计算逻辑简洁直接
• 绕过步长限制：通过将数值乘以2转为整数序列，间接实现0.5的间隔需求，完美适配Polars的arange特性

测试验证 #

针对提供的测试数据，运行上述代码可得到与预期完全一致的结果：

# 测试数据
test_df = pl.DataFrame([
    pl.Series('a', [[4.0, 5.0, 3.0, 7.0, 0.0, 1.0, 6.0, 2.0], [2.0, 4.0, 3.0, 0.0, 1.0], [1.0, 2.0, 3.0, 0.0, 4.0], [1.0, 3.0, 2.0, 0.0], [1.0, 0.0]], dtype=pl.List(pl.Float64)),        
    pl.Series('b', [[0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0], [0.0, 0.5, 1.0]], dtype=pl.List(pl.Float64))
])

# 生成预测列
result_df = test_df.with_columns(
    pl.col("a").list.eval(
        pl.arange(
            (pl.element().min() * 2).cast(pl.Int64),
            (pl.element().max() * 2).cast(pl.Int64) + 1,
            1
        ).cast(pl.Float64) / 2,
        parallel=True
    ).alias("b_pred")
)

# 验证结果一致性
print(result_df.select(pl.col("b") == pl.col("b_pred")).all())

执行后输出true，说明预测列与预期列完全匹配。

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