在Rust中获取Polars DataFrame datetime列的时间单位与时区

核心需求 #

现有包含time1 datetime列的Polars DataFrame，需传入Rust函数处理，要求在Rust中直接获取该列的时间单位和时区信息，替代当前Python侧提前提取的逻辑，完成转换后返回处理后的DataFrame。

解决方法 #

在Rust的Polars绑定中，DataType枚举的Datetime变体直接包含time_unit（TimeUnit类型）和time_zone（Option<TimeZone>类型）信息，通过匹配该变体即可提取所需参数。

完整Rust实现代码

use polars::prelude::*;
use pyo3::prelude::*;
use pyo3_polars::PyDataFrame;

#[pyfunction]
fn foo(pydf: PyDataFrame) -> PyResult<PyDataFrame> {
    // 将PyDataFrame转换为Polars原生DataFrame引用
    let df = pydf.as_ref();
    let schema = df.schema();

    // 获取time1列的数据类型，处理列不存在的情况
    let time1_type = schema.get("time1")
        .ok_or_else(|| PyErr::new::<PyValueError, _>("未找到列'time1'"))?;

    // 匹配Datetime类型，提取时间单位和时区
    let (time_unit, time_zone) = match time1_type {
        DataType::Datetime(unit, tz) => (unit, tz.clone()),
        _ => return Err(PyErr::new::<PyValueError, _>("列'time1'不是datetime类型")),
    };

    // 示例：将datetime转换为对应精度的时间戳（可替换为你的业务逻辑）
    let processed_df = df.with_column(
        col("time1")
            .dt()
            .timestamp(*time_unit)
            .alias("time2_epoch")
    )?;

    // 将处理后的DataFrame转换为PyDataFrame返回给Python
    Ok(PyDataFrame::from(processed_df))
}

// 导出Rust函数到Python
#[pymodule]
fn your_rust_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(foo, m)?)?;
    Ok(())
}

关键细节说明

• 时间单位提取：TimeUnit枚举包含Nanoseconds、Microseconds、Milliseconds三种变体，可直接用于Polars的dt().timestamp()方法生成对应精度的时间戳。
• 时区提取：time_zone是Option<TimeZone>类型，若列不含时区则为None，使用时需处理空值场景。
• 错误处理：加入了列不存在、类型不匹配的错误捕获，避免Panic。

简化后的Python调用代码 #

无需提前提取时区和单位，直接传入DataFrame即可：

from datetime import datetime
from zoneinfo import ZoneInfo
import polars as pl
from your_rust_module import foo

# 构建带时区的datetime列DataFrame
pydf = pl.DataFrame({'time1': datetime(2024, 1, 1, tzinfo=ZoneInfo("UTC"))})

# 调用Rust函数处理
rust_processed_df = foo(pydf)

# 从原DataFrame获取时区和单位，转回datetime列
time_schema = pydf.schema['time1']
final_df = rust_processed_df.with_columns(
    pl.from_epoch(pl.col('time2_epoch'), time_schema.time_unit).alias('time2')
).with_columns(
    pl.col('time2').dt.convert_time_zone(time_schema.time_zone)
)

可选优化：让Rust返回时区与单位 #

如果不想依赖原DataFrame获取元数据，可修改Rust函数返回处理后的DataFrame+时间单位字符串+时区字符串：

#[pyfunction]
fn foo_with_metadata(pydf: PyDataFrame) -> PyResult<(PyDataFrame, String, Option<String>)> {
    let df = pydf.as_ref();
    let schema = df.schema();

    let time1_type = schema.get("time1")
        .ok_or_else(|| PyErr::new::<PyValueError, _>("未找到列'time1'"))?;

    let (time_unit, time_zone) = match time1_type {
        DataType::Datetime(unit, tz) => (unit, tz.clone()),
        _ => return Err(PyErr::new::<PyValueError, _>("列'time1'不是datetime类型")),
    };

    let processed_df = df.with_column(
        col("time1")
            .dt()
            .timestamp(*time_unit)
            .alias("time2_epoch")
    )?;

    // 将TimeUnit转换为Python兼容的字符串（ns/us/ms）
    let time_unit_str = match time_unit {
        TimeUnit::Nanoseconds => "ns",
        TimeUnit::Microseconds => "us",
        TimeUnit::Milliseconds => "ms",
    }.to_string();

    // 时区转换为字符串（空时区则返回None）
    let time_zone_str = time_zone.map(|tz| tz.to_string());

    Ok((PyDataFrame::from(processed_df), time_unit_str, time_zone_str))
}

对应Python调用：

rust_df, time_unit, time_zone = foo_with_metadata(pydf)

final_df = rust_df.with_columns(
    pl.from_epoch(pl.col('time2_epoch'), time_unit).alias('time2')
).with_columns(
    pl.col('time2').dt.convert_time_zone(time_zone) if time_zone else pl.col('time2')
)

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