Polars Python：分组上下文下的x、y列插值实现（支持LazyFrame）

针对你的需求，我们可以通过分组计算边界→生成完整序列→关联原数据→组内插值的流程实现，全程使用pl.LazyFrame处理，适配大数据量场景。

完整代码实现 #

import polars as pl

# 构造示例LazyFrame
data = {
    "x": [5, 10, 20, 25, 10, 20, 30],
    "y": [1, 2, 4, 5, 2, 4, 6],
    "z": ["A", "A", "A", "A", "B", "B", "B"]
}
lf = pl.LazyFrame(data)

step = 5

result = (
    lf
    # 1. 按z分组，获取每组x的最小/最大值
    .group_by("z")
    .agg(
        pl.col("x").min().alias("min_x"),
        pl.col("x").max().alias("max_x")
    )
    # 2. 生成每组完整的x序列（步长5），并拆分为行
    .with_columns(
        pl.int_ranges(pl.col("min_x"), pl.col("max_x") + step, step).alias("x")
    )
    .explode("x")
    # 3. 关联原数据，匹配对应y值
    .join(lf, on=["z", "x"], how="left")
    # 4. 按z分组对y进行线性插值，转回整数类型
    .group_by("z")
    .agg(
        pl.col("x"),
        pl.col("y").interpolate().cast(pl.Int64)
    )
    # 5. 展开序列并排序，对齐预期输出
    .explode(["x", "y"])
    .sort(["z", "x"])
    .collect()
)

print(result)

代码步骤说明 #

1. 分组计算边界：按分类列z聚合，得到每个组内x的最小、最大值，确定补全序列的范围。
2. 生成完整x序列：使用Polars内置的pl.int_ranges函数（兼容LazyFrame）生成步长为5的连续整数序列，再通过explode将序列拆分为单独的行。
3. 关联原数据：以z和x为联合键左连接原数据，将已有y值匹配到完整序列，缺失位置留空为null。
4. 组内插值：再次按z分组，对y列执行线性插值，最后将插值后的浮点结果转回Int64类型，和原数据类型保持一致。
5. 整理输出：展开聚合后的序列，按z和x排序，得到和预期一致的结果。

输出结果 #

运行代码后将得到：

┌─────┬─────┬─────┐
│ z   ┆ x   ┆ y   │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ A   ┆ 5   ┆ 1   │
│ A   ┆ 10  ┆ 2   │
│ A   ┆ 15  ┆ 3   │
│ A   ┆ 20  ┆ 4   │
│ A   ┆ 25  ┆ 5   │
│ B   ┆ 10  ┆ 2   │
│ B   ┆ 15  ┆ 3   │
│ B   ┆ 20  ┆ 4   │
│ B   ┆ 25  ┆ 5   │
│ B   ┆ 30  ┆ 6   │
└─────┴─────┴─────┘

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