如何在Python Polars中更地道地实现跨DataFrame时间序列插值

问题背景 #

有两个Polars DataFrame：df_ref包含参考时间轴，df是另一个不同频率的时间序列数据，需要将df的数值列插值到df_ref的时间轴上。

示例代码 #

import polars as pl

# 参考时间轴DataFrame
df_ref = pl.DataFrame({"dt": ["2022-12-14T14:00:01.000", "2022-12-14T14:00:02.000",
                              "2022-12-14T14:00:03.000", "2022-12-14T14:00:04.000",
                              "2022-12-14T14:00:05.000", "2022-12-14T14:00:06.000"]})
df_ref = df_ref.with_columns(pl.col("dt").str.to_datetime())

# 需要被插值的时间序列DataFrame
df = pl.DataFrame({
        "dt": ["2022-12-14T14:00:01.500", "2022-12-14T14:00:03.500", "2022-12-14T14:00:05.500"],
        "v1": [1.5, 3.5, 5.5]})
df = df.with_columns(pl.col("dt").str.to_datetime())

失败尝试 #

直接左连接后调用interpolate()因时间键完全不匹配，结果全为null：

print(df_ref.join(df, on="dt", how="left").interpolate())

输出：

shape: (6, 2)
┌─────────────────────┬──────┐
│ dt                  ┆ v1   │
│ ---                 ┆ ---  │
│ datetime[μs]        ┆ f64  │
╞═════════════════════╪══════╡
│ 2022-12-14 14:00:01 ┆ null │
│ 2022-12-14 14:00:02 ┆ null │
│ 2022-12-14 14:00:03 ┆ null │
│ 2022-12-14 14:00:04 ┆ null │
│ 2022-12-14 14:00:05 ┆ null │
│ 2022-12-14 14:00:06 ┆ null │
└─────────────────────┴──────┘

用户已通过SciPy的interp1d实现需求，但希望找到Polars原生方案以获取更好性能。

Polars原生实现方案 #

Polars目前无直接的时间轴插值API，但可通过join_asof匹配每个参考时间的前后数据点，再手动计算线性插值，完全基于原生函数实现，无需依赖第三方库。

方法1：单列插值实现 #

# 准备用于匹配后一个时间点的数据集
df_next = df.rename({"dt": "dt_next", "v1": "v1_next"})

# 匹配每个参考时间的前一个/后一个数据点
df_joined = df_ref.join_asof(df, on="dt", direction="backward")
df_joined = df_joined.join_asof(df_next, left_on="dt", right_on="dt_next", direction="forward")

# 计算线性插值（含首尾外推逻辑）
df_result = df_joined.with_columns(
    # 时间转为时间戳数值，方便计算
    t_prev = pl.col("dt").dt.timestamp(),
    t_next = pl.col("dt_next").dt.timestamp(),
    t_ref = pl.col("dt").dt.timestamp()
).with_columns(
    # 计算斜率
    slope = (pl.col("v1_next") - pl.col("v1")) / (pl.col("t_next") - pl.col("t_prev"))
).with_columns(
    pl.when(pl.col("t_ref") < pl.col("t_prev"))
    # 参考时间早于第一个数据点，外推
    .then(pl.col("v1") + pl.col("slope") * (pl.col("t_ref") - pl.col("t_prev")))
    .when(pl.col("t_ref") > pl.col("t_next"))
    # 参考时间晚于最后一个数据点，外推
    .then(pl.col("v1_next") + pl.col("slope") * (pl.col("t_ref") - pl.col("t_next")))
    # 中间区间线性插值
    .otherwise(pl.col("v1") + pl.col("slope") * (pl.col("t_ref") - pl.col("t_prev")))
    .alias("v1_interp")
).select("dt", "v1_interp")

print(df_result)

输出结果与SciPy方案完全一致：

shape: (6, 2)
┌─────────────────────┬───────────┐
│ dt                  ┆ v1_interp │
│ ---                 ┆ ---       │
│ datetime[μs]        ┆ f64       │
╞═════════════════════╪═══════════╡
│ 2022-12-14 14:00:01 ┆ 1.0       │
│ 2022-12-14 14:00:02 ┆ 2.0       │
│ 2022-12-14 14:00:03 ┆ 3.0       │
│ 2022-12-14 14:00:04 ┆ 4.0       │
│ 2022-12-14 14:00:05 ┆ 5.0       │
│ 2022-12-14 14:00:06 ┆ 6.0       │
└─────────────────────┴───────────┘

方法2：批量处理多列 #

如果有多个数值列需要插值，可封装为复用函数：

def interpolate_column(df_ref: pl.DataFrame, df_data: pl.DataFrame, col_name: str) -> pl.DataFrame:
    df_next = df_data.rename({"dt": "dt_next", col_name: f"{col_name}_next"})
    df_joined = df_ref.join_asof(df_data, on="dt", direction="backward")
    df_joined = df_joined.join_asof(df_next, left_on="dt", right_on="dt_next", direction="forward")
    
    return df_joined.with_columns(
        t_prev = pl.col("dt").dt.timestamp(),
        t_next = pl.col("dt_next").dt.timestamp(),
        t_ref = pl.col("dt").dt.timestamp()
    ).with_columns(
        slope = (pl.col(f"{col_name}_next") - pl.col(col_name)) / (pl.col("t_next") - pl.col("t_prev"))
    ).with_columns(
        pl.when(pl.col("t_ref") < pl.col("t_prev"))
        .then(pl.col(col_name) + pl.col("slope") * (pl.col("t_ref") - pl.col("t_prev")))
        .when(pl.col("t_ref") > pl.col("t_next"))
        .then(pl.col(f"{col_name}_next") + pl.col("slope") * (pl.col("t_ref") - pl.col("t_next")))
        .otherwise(pl.col(col_name) + pl.col("slope") * (pl.col("t_ref") - pl.col("t_prev")))
        .alias(f"{col_name}_interp")
    ).select("dt", f"{col_name}_interp")

# 调用函数处理v1列
df_result = interpolate_column(df_ref, df, "v1")

方案优势 #

• 完全基于Polars原生函数，避免跨语言调用开销
• 底层由Rust实现，大规模数据场景下性能优于SciPy的Python实现
• 支持批量处理多列，代码复用性强

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