如何在Polars中用np.linspace计算列？有无Polars等效方法？

问题背景 #

现有如下Polars DataFrame：

import polars as pl

df = pl.DataFrame(
    data={
        "np_linspace_start": [0, 0, 0], 
        "np_linspace_stop": [8, 6, 7],
        "np_linspace_num": [5, 4, 4]
    }
)

输出结构：

shape: (3, 3)
┌───────────────────┬──────────────────┬─────────────────┐
│ np_linspace_start ┆ np_linspace_stop ┆ np_linspace_num │
│ ---               ┆ ---              ┆ ---             │
│ i64               ┆ i64              ┆ i64             │
╞═══════════════════╪══════════════════╪═════════════════╡
│ 0                 ┆ 8                ┆ 5               │
│ 0                 ┆ 6                ┆ 4               │
│ 0                 ┆ 7                ┆ 4               │
└───────────────────┴──────────────────┴─────────────────┘

需求是新增一列ls，存储对应行参数调用np.linspace后的结果（类型为np.array），尝试的代码如下：

df.with_columns(
    ls=np.linspace(
        start=pl.col("np_linspace_start"),
        stop=pl.col("np_linspace_stop"),
        num=pl.col("np_linspace_num")
    )
)

解决方案 #

Polars没有直接对应np.linspace的内置函数，但可以通过两种方式实现需求：

方法1：逐行调用np.linspace（适合小数据集） #

利用pl.map_elements逐行处理，直接调用numpy的linspace函数：

import numpy as np

result = df.with_columns(
    ls=pl.struct(["np_linspace_start", "np_linspace_stop", "np_linspace_num"])
          .map_elements(
              lambda row: np.linspace(row["np_linspace_start"], row["np_linspace_stop"], row["np_linspace_num"]),
              dtype=pl.Array(pl.Float64)  # 若需保留原生np.array类型，可去掉dtype参数
          )
)

print(result)

输出结果：

shape: (3, 4)
┌───────────────────┬──────────────────┬─────────────────┬───────────────────────────────────┐
│ np_linspace_start ┆ np_linspace_stop ┆ np_linspace_num ┆ ls                                │
│ ---               ┆ ---              ┆ ---             ┆ ---                               │
│ i64               ┆ i64              ┆ i64             ┆ array[f64]                        │
╞═══════════════════╪══════════════════╪═════════════════╪═══════════════════════════════════╡
│ 0                 ┆ 8                ┆ 5               ┆ [0.0, 2.0, 4.0, 6.0, 8.0]         │
│ 0                 ┆ 6                ┆ 4               ┆ [0.0, 2.0, 4.0, 6.0]              │
│ 0                 ┆ 7                ┆ 4               ┆ [0.0, 2.33333333, 4.66666667, 7.0]│
└───────────────────┴──────────────────┴─────────────────┴───────────────────────────────────┘

方法2：纯Polars表达式实现（适合大数据集） #

基于linspace的数学公式，用Polars原生表达式生成结果，避免依赖numpy，性能更优：

result = df.with_columns(
    ls=pl.arange(0, pl.col("np_linspace_num"), 1)
          .map_batches(
              lambda idx: (pl.col("np_linspace_start") + 
                          (pl.col("np_linspace_stop") - pl.col("np_linspace_start")) * idx / (pl.col("np_linspace_num") - 1))
          )
          .list()
)

如果需要将结果转为np.array，可以在list()后追加.map_elements(np.array)，但大数据集下不建议这么做，优先使用Polars原生的Array类型以保证性能。

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