如何用chrono::NaiveDateTime范围过滤Rust Polars LazyFrame？

问题代码 #

原有代码尝试通过chrono::NaiveDateTime范围过滤LazyFrame，但因类型不兼容无法编译：

use polars::prelude::*;
use polars_lazy::prelude::*;

pub fn keep_range_lazy(
    df: &mut DataFrame,
    start: &NaiveDateTime,
    end: &NaiveDateTime,
) -> Result<(), PolarsError> {
    assert!(
        df.get_column_names().contains(&"timestamp"),
        "Dataframe does not contain timestamp column."
    );

    df = &mut df
        .lazy()
        .filter(
            col("timestamp")
                .dt()
                .datetime()
                .gt_eq(start)
                .and(col("timestamp").dt().datetime().lt(end)),
        )
        .collect()?;

    Ok(())
}

报错信息 #

error[E0277]: the trait bound `polars_lazy::dsl::Expr: std::convert::From<&chrono::NaiveDateTime>` is not satisfied
   --> src/utils.rs:32:24
    |
32  |                 .gt_eq(start)
    |                  ----- ^^^^^ the trait `std::convert::From<&chrono::NaiveDateTime>` is not implemented for `polars_lazy::dsl::Expr`
    |                  |
    |                  required by a bound introduced by this call
    |
 = help: the following other types implement trait `std::convert::From<T>`:
              <polars_lazy::dsl::Expr as std::convert::From<&str>>
              <polars_lazy::dsl::Expr as std::convert::From<bool>>
              <polars_lazy::dsl::Expr as std::convert::From<f32>>
              <polars_lazy::dsl::Expr as std::convert::From<f64>>
              <polars_lazy::dsl::Expr as std::convert::From<i32>>
              <polars_lazy::dsl::Expr as std::convert::From<i64>>
              <polars_lazy::dsl::Expr as std::convert::From<polars_lazy::dsl::AggExpr>>
              <polars_lazy::dsl::Expr as std::convert::From<u32>>
              <polars_lazy::dsl::Expr as std::convert::From<u64>>
 = note: required for `&chrono::NaiveDateTime` to implement `std::convert::Into<polars_lazy::dsl::Expr>`
note: required by a bound in `polars_plan::dsl::<impl polars_lazy::dsl::Expr>::gt_eq`
   --> /home/username/.cargo/registry/src/github.com-1ecc6299db9ec823/polars-plan-0.28.0/src/dsl/mod.rs:258:21
    |
258 |     pub fn gt_eq<E: Into<Expr>>(self, other: E) -> Expr {
    |                     ^^^^^^^^^^ required by this bound in `polars_plan::dsl::<impl Expr>::gt_eq`

解决方案 #

调整函数签名以符合Rust所有权规范，同时将chrono::NaiveDateTime转换为Polars兼容的字面量表达式，直接使用完整DateTime变量完成范围过滤：

use polars::prelude::*;
use polars_lazy::prelude::*;
use polars::datetime::PolarsNaiveDateTime;

pub fn get_range_lazy(df: DataFrame, start: &NaiveDateTime, end: &NaiveDateTime) -> Result<DataFrame, PolarsError> {
    assert!(
        df.get_column_names().contains(&"timestamp"),
        "Dataframe does not contain timestamp column."
    );

    // 将chrono时间类型转换为Polars内部兼容的时间类型
    let start_pl = PolarsNaiveDateTime::from(*start);
    let end_pl = PolarsNaiveDateTime::from(*end);

    let filtered_df = df
        .lazy()
        .filter(
            col("timestamp")
                .gt_eq(lit(start_pl))
                .and(col("timestamp").lt(lit(end_pl)))
        )
        .collect()?;

    Ok(filtered_df)
}

核心要点 #

• 通过PolarsNaiveDateTime::from完成跨库时间类型转换，保留完整时间信息
• 用lit()函数将转换后的时间包装为Lazy API所需的Expr字面量
• 无需拆分时分秒字段，直接通过完整DateTime变量进行范围比较，保留LazyFrame的性能优势
• 调整函数签名为返回DataFrame，避免可变引用带来的所有权问题

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