如何利用Polars中另一DataFrame的分组最大值截断数值上限

可以通过分组聚合获取最大值 + 关联DataFrame + 取最小值截断这三步实现需求：

步骤1：提取每个分组的最大index值 #

先从df1中按group分组，计算每个组的index最大值：

group_max = df1.group_by("group").agg(pl.max("index").alias("max_index"))

得到的结果如下：

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

步骤2：关联并截断df2的index #

将df2和上述group_max按group关联，然后对每行的index和对应分组的max_index取最小值，即可完成截断：

result = df2.join(group_max, on="group").with_columns(
    pl.min([pl.col("index"), pl.col("max_index")]).alias("index")
).drop("max_index")

最终结果 #

执行后得到的result与预期完全一致：

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

可选：窗口函数简化写法 #

如果不想单独创建group_max，可以用窗口函数直接在df2中引用df1的分组最大值，不过数据量大时性能略低于join方案：

result = df2.with_columns(
    pl.min([
        pl.col("index"),
        pl.col("group").map_elements(lambda g: df1.filter(pl.col("group") == g)["index"].max())
    ]).alias("index")
)

内容的提问来源于stack exchange，提问作者Kevin Li