引言
在数据分析和预处理中，经常会遇到时间序列或有序数据中存在缺失值的情况。当数据需要按照某个类别进行分组，并在每个组内进行插值时，问题会变得更复杂。polars作为一个高性能的dataframe库，提供了强大的分组和数据操作能力。本教程将深入探讨如何在polars中，针对具有分类列的数据帧，在group_by的上下文中实现缺失值的插值，特别是当我们需要填充固定步长的序列时。
问题描述与示例数据
假设我们有一个Polars数据帧，包含三列：x（整数，可能存在缺失值但应按固定步长递增）、y（整数）和 z（字符串，表示类别）。我们的目标是根据列z进行分组，然后在每个组内，将x列填充为连续的、步长为5的整数序列，并对相应的y列进行线性插值。
原始数据帧示例：
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("原始数据帧:")
print(df)输出：
原始数据帧: shape: (7, 3) ┌─────┬─────┬─────┐ │ 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 │ └─────┴─────┴─────┘
期望的输出：
对于类别'A'，x值应从5到25，步长为5；对于类别'B'，x值应从10到30，步长为5。y值应根据x的变化进行插值。
┌─────┬─────┬─────┐ │ 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值的完整序列，然后通过左连接将其与原始数据合并，最后对缺失的y值进行插值。
# 示例：不带分组列的插值
single_group_data = dict(x=[10, 20, 30], y=[2, 4, 6])
single_df = pl.DataFrame(single_group_data)
step = 5
# 生成所有期望的x值
all_x_values = pl.DataFrame(dict(x=range(single_df["x"].min(), single_df["x"].max() + step, step)))
# 合并并插值
interpolated_single_df = (
all_x_values
.join(single_df, on="x", how="left")
.with_columns(pl.col("y").interpolate())
)
print("\n非分组场景下的插值结果:")
print(interpolated_single_df)输出：
非分组场景下的插值结果: shape: (5, 2) ┌─────┬─────┐ │ x ┆ y │ │ --- ┆ --- │ │ i64 ┆ i64 │ ╞═════╪═════╡ │ 10 ┆ 2 │ │ 15 ┆ 3 │ │ 20 ┆ 4 │ │ 25 ┆ 5 │ │ 30 ┆ 6 │ └─────┴─────┘
这种方法在单个序列中效果良好，但直接将其应用于group_by上下文会遇到挑战，因为我们需要为每个组独立生成x的完整序列。
Polars分组插值解决方案
为了在group_by上下文中实现插值，我们需要扩展上述思路，为每个分组动态地生成其完整的x序列。
步骤一：构建分组内完整序列
首先，我们需要为每个类别z确定其x值的最小和最大范围，并基于此生成一个包含所有期望x值的“上采样”数据帧。
# 步骤一：构建分组内完整序列
upsampled_df = (
df
.group_by("z")
.agg(
# 为每个组生成从min(x)到max(x)（包含）的步长为5的整数序列
pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x_range")
)
.explode("x_range") # 展开列表，使每个x值成为一行
.rename({"x_range": "x"}) # 重命名列为x，以便后续连接
)
print("\n步骤一：构建的上采样数据帧:")
print(upsampled_df)解释：
- df.group_by("z"): 按类别z对数据进行分组。
- .agg(...): 对每个分组执行聚合操作。
- pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x_range"): 这是核心。对于每个组，我们获取其x列的最小值 (pl.col("x").min()) 和最大值 (pl.col("x").max())。然后，pl.int_range函数会生成一个从最小值到最大值（包含最大值，因此需要+5）的整数序列，步长为5。这个序列被命名为x_range。
- .explode("x_range"): agg操作会为每个组生成一个包含x序列的列表。explode操作将这个列表展开，使得列表中的每个元素都成为新的一行，并复制相应的z值。
- .rename({"x_range": "x"}): 将新生成的x序列列重命名为x，以便与原始数据帧进行连接。
输出：
步骤一：构建的上采样数据帧: shape: (10, 2) ┌─────┬─────┐ │ z ┆ x │ │ --- ┆ --- │ │ str ┆ i64 │ ╞═════╪═════╡ │ A ┆ 5 │ │ A ┆ 10 │ │ A ┆ 15 │ │ A ┆ 20 │ │ A ┆ 25 │ │ B ┆ 10 │ │ B ┆ 15 │ │ B ┆ 20 │ │ B ┆ 25 │ │ B ┆ 30 │ └─────┴─────┘
步骤二：执行左连接与缺失值插值
现在我们有了一个包含所有期望x值和对应z类别的新数据帧。接下来，我们将这个上采样数据帧与原始数据帧进行左连接，然后对y列的缺失值进行插值。
# 步骤二：执行左连接与缺失值插值
final_interpolated_df = (
upsampled_df
.join(
df,
on=["x", "z"], # 在x和z两列上进行连接
how="left"
)
.with_columns(
pl.col("y").interpolate() # 对y列进行插值
)
.sort(["z", "x"]) # 可选：为了更好的可读性，对结果进行排序
)
print("\n最终插值结果:")
print(final_interpolated_df)解释：
- upsampled_df.join(df, on=["x", "z"], how="left"): 使用left_join操作，将上采样数据帧作为左表，原始数据帧作为右表。连接键是x和z。这意味着对于upsampled_df中的每一行x和z组合，如果df中有匹配的行，则会合并df中的y值；如果没有匹配，y值将为null。
- .with_columns(pl.col("y").interpolate()): interpolate()方法会填充y列中的所有null值。Polars默认使用线性插值。
- .sort(["z", "x"]): 为了使输出结果与期望的输出一致且更具可读性，我们按z和x列进行排序。
最终输出：
最终插值结果: 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 │ └─────┴─────┴─────┘
可以看到，y列的数据类型已变为f64，这是插值操作的常见结果，因为插值可能产生浮点数。
注意事项与最佳实践
-
LazyFrame的应用： 原始问题中提到希望使用pl.LazyFrame。上述解决方案完全兼容LazyFrame。只需将初始的pl.DataFrame替换为pl.LazyFrame，并在最后调用.collect()即可。对于大规模数据集，使用LazyFrame可以显著提高性能，因为它允许Polars进行查询优化和惰性计算。
# 示例 LazyFrame 应用 lazy_df = pl.LazyFrame(data) lazy_upsampled_df = ( lazy_df .group_by("z") .agg( pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x_range") ) .explode("x_range") .rename({"x_range": "x"}) ) lazy_final_interpolated_df = ( lazy_upsampled_df .join( lazy_df, on=["x", "z"], how="left" ) .with_columns( pl.col("y").interpolate() ) .sort(["z", "x"]) .collect() # 惰性计算的最后一步 ) print("\nLazyFrame 模式下的插值结果:") print(lazy_final_interpolated_df) - 数据类型转换： interpolate()操作通常会将整数列转换为浮点数类型（例如i64变为f64），以适应插值可能产生的非整数结果。如果需要整数类型，可能需要后续进行四舍五入或类型转换。
- 结果排序： group_by操作默认不保证输出顺序。为了确保结果的可预测性和可读性，尤其是在插值后，建议显式地使用.sort()方法对结果进行排序。
- 插值方法： Polars的interpolate()方法默认执行线性插值。对于更复杂的插值需求（如样条插值），可能需要结合其他库或自定义函数。
- 性能考量： explode操作会增加数据帧的行数，可能会消耗更多内存。然而，对于插值这类场景，这是构建完整序列的必要步骤。结合LazyFrame可以缓解内存压力，通过优化执行计划来提高效率。
总结
本教程详细展示了在Polars中如何有效地在group_by上下文下实现缺失值插值。核心思想是首先为每个分组生成一个包含所有预期序列值的“上采样”数据帧，然后通过左连接将其与原始数据合并，最后利用Polars强大的interpolate()功能填充缺失值。这种方法不仅解决了复杂的插值问题，而且通过结合LazyFrame，能够高效处理大规模数据集，是Polars数据处理能力的一个典型应用。
