如何在Polars中高效处理DataFrame的多行切片（时间序列场景）

我有如下时间序列数据集（Rust + Polars）：

let series_one = Series::new(
    "a",
    (0..4).into_iter().map(|v| v as f64).collect::<Vec<_>>(),
);
let series_two = Series::new(
    "b",
    (4..8).into_iter().map(|v| v as f64).collect::<Vec<_>>(),
);
let series_three = Series::new(
    "c",
    (8..12).into_iter().map(|v| v as f64).collect::<Vec<_>>(),
);

let series_dates = Series::new(
    "date",
    (0..4)
        .into_iter()
        .map(|v| NaiveDate::default() + Duration::days(2 * v))
        .collect::<Vec<_>>(),
);
let df = DataFrame::new(vec![series_one, series_two, series_three, series_dates]).unwrap();

对应的DataFrame结构：

shape: (4, 4)
┌─────┬─────┬──────┬────────────┐
│ a   ┆ b   ┆ c    ┆ date       │
│ --- ┆ --- ┆ ---  ┆ ---        │
│ f64 ┆ f64 ┆ f64  ┆ date       │
╞═════╪═════╪══════╪════════════╡
│ 0.0 ┆ 4.0 ┆ 8.0  ┆ 1970-01-01 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1.0 ┆ 5.0 ┆ 9.0  ┆ 1970-01-02 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2.0 ┆ 6.0 ┆ 10.0 ┆ 1970-01-03 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3.0 ┆ 7.0 ┆ 11.0 ┆ 1970-01-04 │
└─────┴─────┴──────┴────────────┘

我需要对每一行，调用some_fn处理包含当前行及所有前置行的DataFrame切片，函数定义如下：

fn some_fn(_df: DataFrame) -> DataFrame {
    // 对输入的DataFrame切片执行操作，返回包含结果的新DataFrame
    DataFrame::new(vec![
        Series::new("a_result", vec![1.0, 2.0, 3.0, 4.0]),
        Series::new("b_result", vec![5.0, 6.0, 7.0, 8.0]),
        Series::new("c_result", vec![9.0, 10.0, 11.0, 12.0]),
    ])
    .unwrap()
}

我最初的实现方式是循环每一行，每次取前i+1行切片传入函数，再拼接结果：

let size = df.column("a").unwrap().len();
let results = (0..size)
    .into_iter()
    .map(|i| {
        let t = df.head((i + 1).into());
        some_fn(t)
    })
    .reduce(|acc, b| acc.vstack(&b).unwrap())
    .unwrap();

但这种方法效率极低，仅处理3000行就耗时约1ms（空函数基准测试，耗时主要来自切片操作）。请问如何利用Polars的特性实现高效处理？

1. 改用状态累积，避免重复切片与拼接 #

当前实现的核心问题是每次循环都要生成新的DataFrame切片，还要反复执行vstack拼接，带来大量不必要的内存分配和拷贝。可以改为逐行累积状态的方式：

// 先修改some_fn为接受引用参数，减少克隆开销
fn some_fn(_df: &DataFrame) -> DataFrame {
    // ... 原有逻辑不变
    DataFrame::new(vec![
        Series::new("a_result", vec![1.0, 2.0, 3.0, 4.0]),
        Series::new("b_result", vec![5.0, 6.0, 7.0, 8.0]),
        Series::new("c_result", vec![9.0, 10.0, 11.0, 12.0]),
    ])
    .unwrap()
}

// 累积处理逻辑
let mut accumulated_df = df.head(0).unwrap();
let mut results = Vec::with_capacity(df.height());

for row in df.iter_rows() {
    // 将当前行转为单行DataFrame并添加到累积集合
    let row_df = DataFrame::new(row.into_iter().map(|(name, val)| {
        Series::new(name, vec![val.clone()])
    }).collect()).unwrap();
    accumulated_df = accumulated_df.vstack(&row_df).unwrap();
    
    // 传入累积的DataFrame引用处理
    let res = some_fn(&accumulated_df);
    results.push(res);
}

// 最终拼接所有结果
let final_results = results.into_iter().reduce(|acc, b| acc.vstack(&b).unwrap()).unwrap();

这种方式只在每次循环中添加一行数据到累积集合，避免了重复生成完整切片，同时通过传递引用减少了DataFrame克隆的开销。

2. 优先使用Polars内置的窗口/累积函数（向量化优化） #

如果some_fn中的操作可以拆解为逐列的滚动累积计算（比如求和、均值、最大值等），直接使用Polars内置的向量化函数会比手动循环高效得多，因为Polars会利用底层的优化执行逻辑：

例如，如果a_result是a列从起始行到当前行的累积和，可以直接写：

let df_with_results = df
    .with_columns([
        col("a").cum_sum().alias("a_result"),
        col("b").cum_sum().alias("b_result"),
        col("c").cum_sum().alias("c_result"),
    ])
    .unwrap();

如果需要更复杂的滚动窗口逻辑，可以用rolling方法自定义窗口范围：

let df_with_results = df
    .with_columns([
        col("a").rolling_sum(None).alias("a_result"),
        col("b").rolling_sum(None).alias("b_result"),
        col("c").rolling_sum(None).alias("c_result"),
    ])
    .unwrap();

其中rolling_sum(None)表示窗口从数据集起始位置到当前行。

3. 大数据集用惰性求值API优化执行计划 #

如果处理的是超大数据集，推荐使用Polars的惰性求值API，它会自动优化执行计划，减少中间数据的内存占用：

let lazy_df = df.lazy();

let results = lazy_df
    .with_columns([
        col("a").cum_sum().alias("a_result"),
        col("b").cum_sum().alias("b_result"),
        col("c").cum_sum().alias("c_result"),
    ])
    .collect()
    .unwrap();

关键优化总结 #

• 避免频繁创建DataFrame切片和拼接，改用累积状态减少内存操作
• 尽可能将操作向量化，优先使用Polars内置的窗口/累积函数，避免手动循环
• 若必须处理完整切片，尽量传递DataFrame引用而非克隆整个对象
• 大数据集优先使用惰性求值API，让Polars自动优化执行流程

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