Polars DataFrame行转列表并高效收集数据的优化方法咨询

我有一个Polars DataFrame：

df = pl.DataFrame(
    {
        "t_left": [0.0, 1.0, 2.0, 3.0],
        "t_right": [1.0, 2.0, 3.0, 4.0],
        "counts": [1, 2, 3, 4],
    }
)

需要将每行转换为对应数量的区间内数值，最终收集所有值传入matplotlib.hist这类函数。

目前手动实现的方式在大数据集下速度极慢：

times = []
for t_left, t_right, counts in df.rows():
    times.extend(np.linspace(t_left, t_right, counts + 1)[1:])

补充的大数据集可复现示例：

import polars as pl
import numpy as np

size = 1000000
df = pl.DataFrame(
    {
        "t_left": np.random.rand(size),
        "t_right": np.random.rand(size) + 1,
        "counts": [1] * size,
    }
)

times = []
for t_left, t_right, counts in df.rows():
    times.extend(np.linspace(t_left, t_right, counts + 1)[1:])

作为Python和Polars新手，想知道有没有更高效的实现方式。

高效实现方案 #

用Polars的向量化操作替代Python循环，能大幅提升速度，核心是利用批量生成索引+广播计算的思路：

通用场景（counts为任意正整数） #

import polars as pl

result = (
    df
    .with_row_index()
    # 按counts生成对应数量的索引，再展开行
    .explode(pl.int_range(0, pl.col("counts")).alias("idx"))
    # 计算每个位置的数值，和原linspace[1:]逻辑一致
    .with_columns(
        (pl.col("t_left") + (pl.col("t_right") - pl.col("t_left")) * (pl.col("idx") + 1) / pl.col("counts"))
        .alias("time")
    )
    .select("time")
    .to_numpy()
    .flatten()
)

特殊场景（counts全为1） #

如果你的counts列全部是1（如补充示例），可以直接简化逻辑，效率更高：

# 和原代码逻辑一致，取t_right（因为linspace(t_left,t_right,2)[1:]就是t_right）
result = df.select("t_right").to_numpy().flatten()

# 若需要区间中点也可以用：
# result = df.with_columns((pl.col("t_left") + pl.col("t_right"))/2).select("t_left").to_numpy().flatten()

效率提升原因 #

• Polars的向量化操作基于Rust实现，完全避开了Python循环的性能开销
• explode+int_range是批量生成重复行，比逐个调用np.linspace高效得多
• 所有计算在Polars内部完成，减少了跨库数据转换的额外消耗

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