1. 问题背景与挑战
在数据分析中，我们经常会遇到时间序列或有序数据中存在缺失值的情况，需要通过插值来填充这些空缺。当数据还包含一个或多个分类维度时，插值通常需要在每个分类组内独立进行。例如，一个数据集包含 x（有序数值，可能缺失）、y（待插值数值）和 z（分类标签）三列。我们的目标是针对 z 的每个类别，确保 x 列按照固定的步长（例如5）连续递增，并根据已有的 y 值对缺失的 y 值进行插值。
一个典型的场景是，原始数据可能只记录了特定 x 值的 y 值：
┌─────┬─────┬─────┐ │ x ┆ y ┆ z │ │ --- ┆ --- ┆ --- │ │ i64 ┆ i64 ┆ str │ ╞═════╪═════╪═════╡ │ 5 ┆ 1 ┆ A │ │ 10 ┆ 2 ┆ A │ │ 20 ┆ 4 ┆ A │ │ 25 ┆ 5 ┆ A │ │ 10 ┆ 2 ┆ B │ │ 20 ┆ 4 ┆ B │ │ 30 ┆ 6 ┆ B │ └─────┴─────┴─────┘
我们期望的输出是，在每个 z 分组内，x 值以5为步长填充所有缺失的中间值，并且对应的 y 值也被插值：
┌─────┬─────┬─────┐ │ x ┆ y ┆ z │ │ --- ┆ --- ┆ --- │ │ i64 ┆ i64 ┆ str │ ╞═════╪═════╪═════╡ │ 5 ┆ 1 ┆ A │ │ 10 ┆ 2 ┆ A │ │ 15 ┆ 3 ┆ A │ │ 20 ┆ 4 ┆ A │ │ 25 ┆ 5 ┆ A │ │ 10 ┆ 2 ┆ B │ │ 15 ┆ 3 ┆ B │ │ 20 ┆ 4 ┆ B │ │ 25 ┆ 5 ┆ B │ │ 30 ┆ 6 ┆ B │ └─────┴─────┴─────┘
对于不含分类列 z 的简单情况，可以通过生成完整的 x 值范围，然后与原始数据进行左连接，最后应用 interpolate() 方法来解决。然而，如何在 group_by 上下文中实现这一复杂逻辑，是 Polars 用户常遇到的挑战。
2. 单组数据插值回顾
在深入探讨分组插值之前，我们先回顾一下对单个组（即没有分类列 z）进行插值的方法。这种方法是理解分组插值的基础。
假设我们只有 x 和 y 两列，并且希望 x 值以5为步长进行填充，并插值 y：
import polars as pl
# 原始数据
data = dict(x=[10, 20, 30], y=[2, 4, 6])
df = pl.DataFrame(data)
# 生成所有 x 值的完整范围
step = 5
df_full_x = pl.DataFrame(dict(x=range(df["x"].min(), df["x"].max() + step, step)))
# 合并并插值
result_single_group = (
df_full_x
.join(df, on="x", how="left")
.with_columns(pl.col("y").interpolate())
)
print("单组插值结果:")
print(result_single_group)输出结果如下：
┌─────┬─────┐ │ x ┆ y │ │ --- ┆ --- │ │ i64 ┆ i64 │ ╞═════╪═════╡ │ 10 ┆ 2 │ │ 15 ┆ 3 │ │ 20 ┆ 4 │ │ 25 ┆ 5 │ │ 30 ┆ 6 │ └─────┴─────┘
这种方法的核心是创建一个包含所有期望 x 值的“模板”DataFrame，然后通过左连接将原始数据合并到这个模板上。连接后，那些在原始数据中缺失的 x 值对应的 y 列将变为 null，这时就可以使用 pl.col("y").interpolate() 来填充这些 null 值。
3. Polars 中分组内数据插值的实现
要将上述单组插值逻辑扩展到 group_by 上下文，关键在于如何为每个分组独立地生成其完整的 x 值范围。Polars 的 group_by 和 agg 操作提供了实现这一目标的高效途径。
3.1 准备示例数据
首先，我们定义原始的 Polars DataFrame：
import polars as pl
# 示例数据
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"]
}
df = pl.DataFrame(data)
print("原始 DataFrame:")
print(df)3.2 步骤一：为每个分组生成完整的 x 值范围
我们需要为每个 z 分组，根据其 x 列的最小值和最大值，生成一个以指定步长（例如5）递增的完整 x 值序列。
# 步骤一：为每个分组生成完整的 x 值范围
upsampled_template = (
df
.group_by("z")
.agg(
# 为每个组生成从 min(x) 到 max(x)+step 的整数范围
# 并将其命名为 "x"
pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x")
)
.explode("x") # 将列表形式的 x 展开为多行
)
print("\n生成的上采样模板 DataFrame:")
print(upsampled_template)解释：
- df.group_by("z"): 按照分类列 z 进行分组。
- agg(...): 对每个分组执行聚合操作。
- pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x"): 这是核心。对于每个分组，它计算 x 列的最小值和最大值，然后生成一个从最小值到最大值（包含最大值，因此 max() + 5）的整数序列，步长为5。alias("x") 确保新列名为 x。
- .explode("x"): agg 操作的结果会是每个 z 组对应一个包含 x 值列表的行。explode("x") 将这些列表展开，使得每个 x 值都成为独立的一行，同时保留其所属的 z 分组信息。
经过此步骤，我们得到了一个包含所有 z 分组和其各自完整 x 值序列的 DataFrame，这将作为我们后续连接的左侧 DataFrame。
3.3 步骤二：执行左连接并插值
现在，我们将这个 upsampled_template DataFrame 与原始 DataFrame df 进行左连接。连接键是 x 和 z。连接后，那些在 upsampled_template 中存在但在原始 df 中不存在的 (x, z) 组合，其 y 列将变为 null。最后，我们对 y 列应用 interpolate() 方法来填充这些 null 值。
# 步骤二：执行左连接并插值
final_result = (
upsampled_template
.join(
df,
on=["x", "z"], # 基于 x 和 z 两列进行连接
how="left" # 左连接，保留 upsampled_template 中的所有行
)
.with_columns(
pl.col("y").interpolate() # 对 y 列进行插值
)
# 可选：如果需要特定排序，例如按 z, x 排序
.sort(["z", "x"])
)
print("\n最终分组内插值结果:")
print(final_result)输出：
最终分组内插值结果: shape: (10, 3) ┌─────┬─────┬─────┐ │ z ┆ x ┆ y │ │ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ f64 │ ╞═════╪═════╪═════╡ │ A ┆ 5 ┆ 1.0 │ │ A ┆ 10 ┆ 2.0 │ │ A ┆ 15 ┆ 3.0 │ │ A ┆ 20 ┆ 4.0 │ │ A ┆ 25 ┆ 5.0 │ │ B ┆ 10 ┆ 2.0 │ │ B ┆ 15 ┆ 3.0 │ │ B ┆ 20 ┆ 4.0 │ │ B ┆ 25 ┆ 5.0 │ │ B ┆ 30 ┆ 6.0 │ └─────┴─────┴─────┘
可以看到，最终结果完美地实现了在每个 z 分组内对 x 值进行上采样，并对 y 值进行插值。
4. 注意事项与最佳实践
-
LazyFrame 的使用: 原始问题中提到数据集可能非常大，希望使用 pl.LazyFrame。上述解决方案完全兼容 LazyFrame。只需在开始时将 pl.DataFrame 替换为 pl.LazyFrame，并在最后调用 .collect() 即可。例如：
# 使用 LazyFrame df_lazy = pl.LazyFrame(data) upsampled_template_lazy = ( df_lazy .group_by("z") .agg( pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x") ) .explode("x") ) final_result_lazy = ( upsampled_template_lazy .join( df_lazy, on=["x", "z"], how="left" ) .with_columns( pl.col("y").interpolate() ) .sort(["z", "x"]) .collect() # 在 LazyFrame 链的末尾调用 collect() 获取结果 ) print("\nLazyFrame 分组内插值结果:") print(final_result_lazy)使用 LazyFrame 可以在处理大规模数据时获得更好的性能，因为它允许 Polars 优化查询计划并延迟计算。
排序的重要性: 在执行插值之前，确保数据在每个分组内是按照 x 列排序的至关重要。join 操作通常不会改变数据的内部顺序，但为了确保 interpolate() 方法的正确性，以及最终输出的可读性，建议在最终结果上使用 .sort(["z", "x"]) 来明确排序。Polars 的 interpolate 方法默认假定数据是排序的。
插值方法: pl.col("y").interpolate() 默认执行线性插值。如果需要更复杂的插值方法（例如，前后填充 fill_null(strategy="forward") 或 fill_null(strategy="backward")），可以根据具体需求调整。但对于这种数值序列的填充，线性插值通常是合适的。
性能考量: Polars 的 group_by 和 join 操作都是高度优化的。通过 explode 生成完整的 x 范围，然后进行连接，这种模式在 Polars 中是处理此类问题的惯用且高效的方式。
5. 总结
本文提供了一种在 Polars 中对分组数据进行插值和上采样的通用且高效的方法。通过结合 group_by、agg、explode 和 join 操作，我们能够为每个分组独立地生成完整的数值序列，并通过左连接引入缺失值，最终利用 interpolate() 方法精确填充这些缺失值。这种方法不仅适用于小规模数据集，通过结合 LazyFrame 也能轻松扩展到大规模数据场景，充分体现了 Polars 在数据处理方面的灵活性和高性能。
