Polars数据帧分组插值方法详解

来到golang学习网的大家，相信都是编程学习爱好者，希望在这里学习文章相关编程知识。下面本篇文章就来带大家聊聊《Polars数据帧分组插值教程》，介绍一下，希望对大家的知识积累有所帮助，助力实战开发！

引言

在数据分析和预处理中，经常会遇到时间序列或有序数据中存在缺失值的情况。当数据需要按照某个类别进行分组，并在每个组内进行插值时，问题会变得更复杂。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)

解释：

1. df.group_by("z"): 按类别z对数据进行分组。
2. .agg(...): 对每个分组执行聚合操作。
3. 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。
4. .explode("x_range"): agg操作会为每个组生成一个包含x序列的列表。explode操作将这个列表展开，使得列表中的每个元素都成为新的一行，并复制相应的z值。
5. .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)

解释：

1. upsampled_df.join(df, on=["x", "z"], how="left"): 使用left_join操作，将上采样数据帧作为左表，原始数据帧作为右表。连接键是x和z。这意味着对于upsampled_df中的每一行x和z组合，如果df中有匹配的行，则会合并df中的y值；如果没有匹配，y值将为null。
2. .with_columns(pl.col("y").interpolate()): interpolate()方法会填充y列中的所有null值。Polars默认使用线性插值。
3. .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数据处理能力的一个典型应用。

今天带大家了解了的相关知识，希望对你有所帮助；关于文章的技术知识我们会一点点深入介绍，欢迎大家关注golang学习网公众号，一起学习编程~