Python Polars：如何为所有列应用聚合函数并传入额外列参数？

错误原因分析 #

你遇到的 TypeError: 'Expr' object is not subscribable 是因为在自定义聚合函数中尝试用Python索引方式（t[v.arg_max()]）操作Polars表达式对象。Polars的表达式系统不支持这种语法，必须使用其内置的函数来实现逻辑。此外，原函数中还存在一个bug：smin_t 错误地使用了 v.arg_max() 而非 v.arg_min()。

高效解决方案 #

针对8k个传感器列的场景，我们可以利用Polars原生聚合表达式实现高效处理，避免使用性能较低的map_elements。以下是具体实现：

import polars as pl

# 假设df是通过scan_parquet创建的lazy dataframe
df = pl.scan_parquet("your_data.parquet")

# 获取所有传感器列（排除region和time）
sensor_cols = [col for col in df.columns if col not in ["region", "time"]]

# 生成聚合表达式：为每个传感器列计算最大值、最小值及对应时间
agg_exprs = []
for col in sensor_cols:
    agg_exprs.extend([
        pl.col(col).max().alias(f"{col}_max"),
        pl.col(col).min().alias(f"{col}_min"),
        # 取最大值对应的首个时间（若有多个最大值，按出现顺序取第一个）
        pl.col("time").sort_by(pl.col(col), descending=True).first().alias(f"{col}_max_time"),
        # 取最小值对应的首个时间
        pl.col("time").sort_by(pl.col(col), ascending=True).first().alias(f"{col}_min_time")
    ])

# 执行分组聚合（lazy模式下需collect触发计算）
result = df.group_by("region").agg(agg_exprs).collect()
print(result)

关键优化点 #

• 原生Polars表达式：所有操作均在Polars的优化引擎中执行，比Python自定义函数快一个数量级，适合处理8k列的大规模数据。
• 批量生成表达式：通过遍历传感器列批量生成聚合逻辑，避免手动编写重复代码。
• 排序取时间：使用sort_by替代过滤操作，更高效地获取极值对应的时间戳。

示例输出（基于你提供的测试数据） #

┌────────┬─────────┬─────────┬──────────────┬──────────────┬─────────┬─────────┬──────────────┬──────────────┬─────────┬─────────┬──────────────┬──────────────┐
│ region ┆ sen1_max┆ sen1_min┆ sen1_max_time┆ sen1_min_time┆ sen2_max┆ sen2_min┆ sen2_max_time┆ sen2_min_time┆ sen3_max┆ sen3_min┆ sen3_max_time┆ sen3_min_time│
│ ---    ┆ ---     ┆ ---     ┆ ---          ┆ ---          ┆ ---     ┆ ---     ┆ ---          ┆ ---          ┆ ---     ┆ ---     ┆ ---          ┆ ---          │
│ str    ┆ f64     ┆ f64     ┆ i64          ┆ i64          ┆ f64     ┆ f64     ┆ i64          ┆ i64          ┆ f64     ┆ f64     ┆ i64          ┆ i64          │
╞════════╪═════════╪═════════╪══════════════╪══════════════╪═════════╪═════════╪══════════════╪══════════════╪═════════╪═════════╪══════════════╪══════════════╡
│ us     ┆ 13.0    ┆ 10.0    ┆ 4            ┆ 1            ┆ 14.0    ┆ 10.0    ┆ 2            ┆ 3            ┆ 19.0    ┆ 12.0    ┆5             ┆1             │
│ uk     ┆13.7     ┆10.0     ┆6             ┆1             ┆14.0     ┆10.0     ┆2             ┆3             ┆21.9     ┆12.0     ┆7             ┆3             │
└────────┴─────────┴─────────┴──────────────┴──────────────┴─────────┴─────────┴──────────────┴──────────────┴─────────┴─────────┴──────────────┴──────────────┘

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