如何在Rust Polars LazyFrame列中获取下一个值并判断工作日连续

我认为在LazyFrame列中“向前查看”的操作在概念上存在矛盾，想请教最佳实现方式。我的需求是：为每个日期生成一个新列，存储判断结果(下一个工作日编号 - 当前工作日编号) == 1的布尔值。

以下是我目前的示例代码：

// PLEASE be aware to add the needed feature flags in your toml file

use polars::export::arrow::temporal_conversions::date32_to_date;
use polars::prelude::*;

fn main() -> Result<()> {
    let days = df!(
        "date_string" => &["1900-01-01", "1900-01-02", "1900-01-03", "1900-01-04", "1900-01-05",
        "1900-01-06", "1900-01-07", "1900-01-09", "1900-01-10"])?;

    let options = StrpTimeOptions {
        date_dtype: DataType::Date,   // the result column-datatype
        fmt: Some("%Y-%m-%d".into()), // the source format of the date-string
        strict: false,
        exact: true,
    };

    // convert date_string into dtype(date) and put into new column "date_type"
    // we convert the days DataFrame to a LazyFrame ...
    // because in my real-world example I am getting a LazyFrame
    let mut new_days = days.lazy().with_column(
        col("date_string")
            .alias("date_type")
            .str()
            .strptime(options),
    );

    // This is what I wanted to do ... but I get a string result .. need u32
    // let o = GetOutput::from_type(DataType::Date);
    // new_days = new_days.with_column(
    //     col("date_type")
    //         .alias("weekday_number")
    //         .map(|x| Ok(x.strftime("%w").unwrap()), o.clone()),
    // );

    // This is the convoluted workaround
    let o = GetOutput::from_type(DataType::Date);
    new_days = new_days.with_column(col("date_type").alias("weekday_number").map(
        |x| {
            Ok(x.date()
                .unwrap()
                .clone()
                .into_iter()
                .map(|opt_name: Option<i32>| {
                    opt_name.map(|datum: i32| {
                        // println!("{:?}", datum);
                        date32_to_date(datum)
                            .format("%w")
                            .to_string()
                            .parse::<u32>()
                            .unwrap()
                    })
                })
                .collect::<UInt32Chunked>()
                .into_series())
        },
        o,
    ));

    // Here is where my challenge is ..
    // I need to get the weekday_number of the following day to determine a condition
    // my pseudo code:
    // new_days = new_days.with_column(
    //     col("weekday_number")
    //         .alias("cold_day")
    //         .map(|x| Ok( (next_weekday_number - current_weekday_number) == 1 ), o.clone()),
    // );

    println!("{:?}", new_days.clone().collect());

    Ok(())
}

优化实现方案 #

LazyFrame完全支持“向前查看”操作，核心是用shift(-1)获取下一行列值，同时Polars内置的日期处理函数可以大幅简化工作日编号的计算：

1. 简化工作日编号计算：Polars的Date类型列直接提供weekday()方法，返回0（周日）到6（周六）的整数，无需手动转换字符串解析。
2. 实现向前查看：通过shift(-1)获取下一行的工作日编号，再计算差值是否为1，最后生成布尔列。

完整优化代码如下：

// 需在Cargo.toml中添加对应feature，比如"chrono", "lazy", "strings"
use polars::prelude::*;

fn main() -> Result<()> {
    let days = df!(
        "date_string" => &["1900-01-01", "1900-01-02", "1900-01-03", "1900-01-04", "1900-01-05",
        "1900-01-06", "1900-01-07", "1900-01-09", "1900-01-10"]
    )?;

    let options = StrpTimeOptions {
        date_dtype: DataType::Date,
        fmt: Some("%Y-%m-%d".into()),
        strict: false,
        exact: true,
    };

    let new_days = days.lazy()
        // 转换日期字符串为Date类型
        .with_column(
            col("date_string")
                .str()
                .strptime(options)
                .alias("date_type")
        )
        // 生成工作日编号：0=周日，1=周一...6=周六
        .with_column(
            col("date_type")
                .weekday()
                .cast(DataType::UInt32)
                .alias("weekday_number")
        )
        // 生成目标列：判断下一个工作日编号与当前的差值是否为1
        .with_column(
            (col("weekday_number").shift(-1) - col("weekday_number"))
                .eq(lit(1))
                .alias("cold_day")
        );

    // 执行计算并打印结果
    println!("{:?}", new_days.collect()?);

    Ok(())
}

关键说明 #

• weekday()：Polars内置方法，直接从Date类型列提取工作日编号，比手动转换高效且不易出错。
• shift(-1)：将列向上偏移一位，即每一行获取下一行的值，最后一行会得到Null（因为没有下一行）。
• 整个流程保持LazyFrame的懒执行特性，只有调用collect()时才会实际计算，符合生产环境中使用LazyFrame的需求。

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