如何在 group_by 聚合上下文中使用 polars.lit

小能豆

如何在 group_by 聚合上下文中使用 polars.lit

py

基本上，我想计算一个常数的幂的总和。

举例来说：

import polars as pl

c = 2

df = pl.DataFrame({"a": [1, 2, 3]})

df.select(pl.lit(c).pow(pl.col("a")).sum())
shape: (1, 1)
┌─────────┐
│ literal │
│ ---     │
│ i32     │
╞═════════╡
│ 14      │
└─────────┘

但是，当我尝试在 groupby agg 上下文中执行同样的事情时，出现错误：

import polars as pl

c = 2

df = pl.DataFrame({"a": [1, 2, 3, 1, 2, 3], "b": [1, 1, 1, 2, 2, 2]})

df.group_by("b").agg(pl.lit(c).pow(pl.col("a")).sum())
# Error originated in expression: '2i32.pow([col("a")])'

我可以让它工作的一个（临时）方法是预先在 df 上附加一个列，例如

df = df.with_columns(pl.lit(c).alias("c"))

但它并不干净，因为原始 df 中可能有一个名为“c”（或我给出的任何别名）的现有列，这可能会导致列名冲突。

我确信一定有更好更干净的方法来实现这一点。但是怎么做呢？

阅读 19

2025-01-04

小能豆

更新：旧的行为是一个错误。文字现在可以按预期广播。

df.group_by("b").agg(pl.lit(c).pow(pl.col("a")).sum())
shape: (2, 2)
┌─────┬─────────┐
│ b   ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ i32     │
╞═════╪═════════╡
│ 1   ┆ 14      │
│ 2   ┆ 14      │
└─────┴─────────┘

看起来您可能需要.repeat_by使长度匹配。

df.group_by("b").agg(
    pl.lit(c).repeat_by("a").pow(pl.col("a")).sum()
)
shape: (2, 2)
┌─────┬─────────┐
│ b   ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ f64     │
╞═════╪═════════╡
│ 1   ┆ 14.0    │
│ 2   ┆ 14.0    │
└─────┴─────────┘

2025-01-04

小能豆

感谢更新和分享新的解决方案！从您提供的代码来看，新的行为修复了之前的问题，并且.repeat_by()提供了一种优雅的方法来处理长度匹配的问题。

核心逻辑解释

问题来源： 原始操作中，广播行为在某些情况下未按预期工作，导致错误结果或意外行为。
.repeat_by()的作用： 它确保了常量值（如`pl.lit(c)）被正确重复，以匹配目标列的长度，从而避免了广播不一致的问题。

更新后的代码行为

使用.repeat_by("a")对常量c进行扩展，然后与列进行幂次运算并求和：

df.group_by("b").agg(
    pl.lit(c).repeat_by("a").pow(pl.col("a")).sum()
)

输出结果如预期地对分组后的b列求得了正确的总和。

输出解读

输出表明，每个组中的值被正确计算并累加：

┌─────┬─────────┐
│ b   ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ f64     │
╞═════╪═════════╡
│ 1   ┆ 14.0    │
│ 2   ┆ 14.0    │
└─────┴─────────┘

这是因为c的值被正确扩展和广播，并与列a一起进行了幂运算。

总结

这个解决方案展示了如何利用Polars的.repeat_by()方法优雅地解决分组聚合中的广播问题。如果您有更多类似的操作场景，这种方法或许能够成为标准的处理方式！ 🚀

2025-01-04