引言
在数据分析中，我们经常会遇到时间序列或有序数据中存在缺失值的情况。polars作为一款高性能的数据处理库，提供了强大的插值功能。然而，当需要在分组（groupby）的上下文中进行插值，并且插值目标不仅仅是填充现有缺失值，而是要补齐分组内按特定步长缺失的整个序列时，常规的groupby().agg(pl.col("y").interpolate())方法可能无法满足需求。例如，如果希望在每个分组内，x列的值以固定步长（如5）连续递增，而原始数据中这些中间值可能完全缺失，我们需要一种方法来“创造”这些缺失的x值及其对应的插值y值。
本文将深入探讨如何在Polars中实现这种复杂的分组内序列填充与插值，尤其适用于处理大规模数据集。
问题场景描述
假设我们有一个Polars DataFrame，包含三列：x（整数，可能存在缺失的步长连续值）、y（整数，与x相关）和z（字符串，类别列）。我们的目标是：
- 按z列进行分组。
- 在每个z组内，确保x列的值是连续的，并且以固定的步长（例如5）递增，从该组的x最小值到最大值。
- 对于新生成的或原始数据中x值对应的y值，如果缺失，则进行线性插值。
以下是示例数据：
import polars as pl
# 原始DataFrame
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)输出：
原始DataFrame: 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 │ └─────┴─────┴─────┘
我们期望的输出是这样的，其中x值在每个z组内以步长5填充，并且y值也相应插值：
shape: (10, 3) ┌─────┬─────┬─────┐ │ 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 │ └─────┴─────┴─────┘
解决方案：生成完整序列与左连接插值
解决此问题的核心思路是：
- 为每个分组生成完整的x值序列：根据每个分组的x最小值和最大值，以及预设的步长，生成该分组内所有可能x值的完整序列。
- 创建扩展的DataFrame：将所有分组生成的完整x序列组合成一个临时的DataFrame，其中包含z和所有预期的x值。
- 执行左连接：将这个扩展的DataFrame与原始DataFrame进行左连接，连接键是z和x。这将导致那些在原始数据中缺失的x值在合并后的DataFrame中对应y值为null。
- 应用插值：对连接后生成的DataFrame中的y列应用interpolate()方法，填充所有null值。
步骤一：生成每个分组的完整x值序列
首先，我们需要为每个z分组生成其对应的完整x值范围。这可以通过group_by结合agg和pl.int_range实现。pl.int_range(start, end, step)可以生成一个整数序列。
# 步骤一：为每个分组生成所有预期的x值
# 使用group_by("z")按类别分组
# 使用agg聚合操作，对每个组：
# pl.col("x").min() 获取当前组x的最小值
# pl.col("x").max() 获取当前组x的最大值
# pl.int_range() 生成从最小值到最大值（包含）的步长为5的整数序列
# .alias("x") 将生成的序列命名为"x"
# explode("x") 将列表形式的"x"列展开成多行，每行对应一个x值
upsampled_df = (
df
.group_by("z")
.agg(
pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x")
)
.explode("x")
)
print("\n生成完整x值序列的DataFrame (upsampled_df):")
print(upsampled_df)输出：
生成完整x值序列的DataFrame (upsampled_df): 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 │ └─────┴─────┘
可以看到，upsampled_df现在包含了每个z组内所有期望的x值，包括原始数据中缺失的x值。
步骤二：执行左连接并应用插值
接下来，我们将upsampled_df与原始df进行左连接。左连接会保留upsampled_df中的所有行，并根据z和x匹配原始df中的数据。如果upsampled_df中的某个z-x组合在原始df中不存在，则y列将填充为null。最后，我们对y列执行interpolate()操作。
# 步骤二：左连接原始数据并插值
# 使用upsampled_df作为左表，df作为右表
# on=["x", "z"] 指定连接键为x和z
# how="left" 执行左连接，保留upsampled_df的所有行
# with_columns(pl.col("y").interpolate()) 对连接结果中的y列进行插值
final_df = (
upsampled_df
.join(
df,
on=["x", "z"],
how="left"
)
.with_columns(
pl.col("y").interpolate()
)
)
print("\n最终插值结果DataFrame:")
print(final_df)输出：
最终插值结果DataFrame: 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 │ └─────┴─────┴─────┘
可以看到，最终结果完美地满足了我们的需求。在A组中，x=15被填充，y值被插值为3.0（10:2和20:4之间）。在B组中，x=15和x=25被填充，y值也相应插值。
完整代码示例
import polars as pl
# 原始DataFrame
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)
# 步骤一：为每个分组生成所有预期的x值
upsampled_df = (
df
.group_by("z")
.agg(
pl.int_range(pl.col("x").min(), pl.col("x").max() + 5, step=5).alias("x")
)
.explode("x")
)
print("\n生成完整x值序列的DataFrame (upsampled_df):")
print(upsampled_df)
# 步骤二：左连接原始数据并插值
final_df = (
upsampled_df
.join(
df,
on=["x", "z"],
how="left"
)
.with_columns(
pl.col("y").interpolate()
)
)
print("\n最终插值结果DataFrame:")
print(final_df)注意事项与性能优化
-
Polars LazyFrame：对于大型数据集，建议使用pl.LazyFrame来延迟计算，以获得更好的性能和内存效率。上述代码中的所有操作都可以无缝地应用于LazyFrame。只需将pl.DataFrame(data)替换为pl.LazyFrame(data)，并在最后调用.collect()即可。
# 示例 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") ) .explode("x") ) lazy_final_df = ( lazy_upsampled_df .join( lazy_df, on=["x", "z"], how="left" ) .with_columns( pl.col("y").interpolate() ) ).collect() # 触发计算 print("\nLazyFrame 最终插值结果DataFrame:") print(lazy_final_df) -
排序问题：在某些Polars版本中，group_by操作可能不保证输出的顺序。如果最终结果的行顺序对您很重要（例如，希望每个z组内的x值严格递增），可以在最终DataFrame上应用sort(["z", "x"])。
# 对最终结果进行排序 final_df_sorted = final_df.sort(["z", "x"]) print("\n排序后的最终插值结果DataFrame:") print(final_df_sorted) interpolate()方法：pl.col("y").interpolate()默认执行线性插值。Polars的interpolate方法非常高效，并且能够处理非均匀间隔的数据点。它会根据已有的非null值，推断出null位置的值。
步长和范围：在pl.int_range(pl.col("x").min(), pl.col("x").max() + step, step=step)中，+ step是为了确保max()值本身也能被包含在生成的序列中，即使max()不是min()开始的步长的精确倍数。请根据您的具体需求调整步长（step）参数。
总结
本文提供了一种在Polars中进行复杂分组内数据插值的强大策略。通过巧妙地结合group_by生成完整序列、explode展开数据以及left_join合并原始数据，我们能够有效地填充分组内缺失的序列值，并利用interpolate()方法对关联列进行精确插值。这种方法不仅功能强大，而且在Polars的优化执行引擎下，能够高效处理大规模数据集，是数据预处理中非常有用的技巧。
