如何高效计算Polars DataFrame中值与上一次出现位置的距离？

需求说明 #

需要计算DataFrame中当前行到目标列（示例中为b列）对应值上一次出现位置的距离，逻辑如下：

• 若该值此前出现过，距离 = 当前行索引 - 上一次出现的索引 - 1
• 若为该值首次出现，距离为-1

输入示例 #

Rust代码定义的DataFrame：

let df_a = df![
    "a" => [1, 2, 2, 1, 4, 1],
    "b" => ["c","a", "b", "c", "c","a"]
].unwrap();

对应的表格：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i32 ┆ str │
╞═════╪═════╡
│ 1   ┆ c   │
│ 2   ┆ a   │
│ 2   ┆ b   │
│ 1   ┆ c   │
│ 4   ┆ c   │
│ 1   ┆ a   │
└─────┴─────┘

期望输出 #

┌─────┬─────┬────────┐
│ a   ┆ b   ┆ b_dist │
│ --- ┆ --- ┆ ---    │
│ i32 ┆ str ┆ i32    │
╞═════╪═════╪════════╡
│ 1   ┆ c   ┆ -1     │
│ 2   ┆ a   ┆ -1     │
│ 2   ┆ b   ┆ -1     │
│ 1   ┆ c   ┆ 2      │
│ 4   ┆ c   ┆ 0      │
│ 1   ┆ a   ┆ 3      │
└─────┴─────┴────────┘

高效实现方案 #

利用Polars的向量化窗口操作实现，完全避免循环，性能最优（基于Arrow引擎的批量处理）：

use polars::prelude::*;

fn main() {
    let df_a = df![
        "a" => [1, 2, 2, 1, 4, 1],
        "b" => ["c","a", "b", "c", "c","a"]
    ].unwrap();

    let result = df_a
        .with_row_index("idx", Some(0)) // 添加从0开始的连续行索引
        .with_column(
            col("idx")
                // 按b列分组后，偏移一行获取上一次出现的索引
                .sub(col("idx").shift(1).over([col("b")]))
                .sub(1) // 按需求计算距离
                .fill_null(lit(-1)) // 首次出现的行填充-1
                .alias("b_dist")
        )
        .drop("idx"); // 删除临时索引列

    println!("{}", result);
}

代码解释 #

1. 添加行索引：with_row_index("idx", Some(0))生成从0开始的连续索引列，满足问题中的索引计算逻辑。
2. 窗口偏移获取上一次索引：col("idx").shift(1).over([col("b")])是核心操作——按b列分组后，将每组的索引列向上偏移1行，这样每行就能拿到当前b值上一次出现的索引（首次出现的行会得到null）。
3. 计算距离：通过sub方法完成当前索引 - 上一次索引 - 1的计算。
4. 填充默认值：用fill_null(lit(-1))将首次出现行的null替换为需求中的-1。
5. 清理临时列：最后删除临时的idx列，得到目标结果。

这种方式充分利用Polars的批量处理能力，在大数据量场景下性能远优于手动遍历。

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