Polars 中按组使用另一 DataFrame 的最大值截断数值-Python教程-PHP中文网

Polars 中按组使用另一 DataFrame 的最大值截断数值

本文介绍如何在 polars 中基于分组聚合结果（如 df1 中每组的最大 index）对另一 dataframe（df2）的对应列进行上限截断（clip），确保 df2 中各组的值不超过 df1 同组的最大值。

本文介绍如何在 polars 中基于分组聚合结果（如 df1 中每组的最大 index）对另一 dataframe（df2）的对应列进行上限截断（clip），确保 df2 中各组的值不超过 df1 同组的最大值。

在数据处理中，常需根据参考数据集（如 df1）的统计量（例如每组最大值）来约束目标数据集（如 df2）的取值范围。Polars 提供了高效、声明式的表达方式实现该需求，核心是结合 clip() 与分组聚合逻辑。

基础方案：两 DataFrame 组别完全一致

当 df1 和 df2 的 group 列包含完全相同的组（无缺失、无额外组）时，可直接利用 over() 窗口函数计算 df1 中各组最大值，并作为 clip() 的 upper_bound：

import polars as pl

df1 = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B", "B"],
    "index": [1, 3, 5, 1, 3, 8],
})

df2 = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B", "B"],
    "index": [3, 4, 7, 2, 7, 10],
})

out = df2.with_columns(
    pl.col("index").clip(
        upper_bound=df1.select(pl.col("index").max().over("group"))["index"]
    )
)

print(out)

输出符合预期：

shape: (6, 2)
┌───────┬───────┐
│ group ┆ index │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ A     ┆ 3     │
│ A     ┆ 4     │
│ A     ┆ 5     │
│ B     ┆ 2     │
│ B     ┆ 7     │
│ B     ┆ 8     │
└───────┴───────┘

✅ 原理说明：

df1.select(pl.col("index").max().over("group")) 返回一个单列 DataFrame，其行数与 df1 相同，每行值为对应 group 在 df1 中的 index 最大值；
由于 df1 与 df2 组结构一致且顺序无关（over 保证按 group 对齐逻辑），提取该列后可直接作为广播式上界传入 clip()。

健壮方案：组别可能不完全匹配（推荐用于生产）

若 df1 和 df2 的 group 存在差异（如 df2 包含 df1 中不存在的组，或反之），应改用显式 join 对齐，避免因隐式对齐导致错误或空值传播：

# 示例：df2 多出一个 'B' 行，且 df1 中 B 组最大 index 为 7（非 8）
df1 = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B", "B"],
    "index": [1, 3, 5, 1, 3, 7],  # B 组 max = 7
})

df2 = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B", "B", "B"],  # 多一行 B
    "index": [3, 4, 7, 2, 7, 8, 9],
})

# 步骤：1. 计算 df1 每组最大值 → 2. 与 df2 左连接 → 3. 截断
group_max = df1.group_by("group").agg(pl.col("index").max().alias("max_index"))
out = df2.join(group_max, on="group", how="left").with_columns(
    pl.col("index").clip(upper_bound="max_index")
).select("group", "index")  # 清理辅助列

print(out)

输出：

shape: (7, 2)
┌───────┬───────┐
│ group ┆ index │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ A     ┆ 3     │
│ A     ┆ 4     │
│ A     ┆ 5     │
│ B     ┆ 2     │
│ B     ┆ 7     │
│ B     ┆ 7     │
│ B     ┆ 7     │
└───────┴───────┘

⚠️ 注意事项：

若 df2 中某 group 在 df1 中不存在，join 后 max_index 将为 null，clip() 会保留原值（或报错，取决于 Polars 版本）。如需严格限制，建议提前校验或用 fill_null() 设默认上界；
clip() 是向量化操作，性能优异，优于 map_elements 或循环；
所有操作均为惰性求值友好，可无缝嵌入 lazy() 流程。