如何基于pl.Expr编写多列筛选的自定义Polars表达式？

原函数还原 #

先明确你提到的两个函数的典型实现：

函数f（硬编码列名） #

import polars as pl

def f(df: pl.DataFrame) -> pl.DataFrame:
    return df.filter(pl.col("C") > 7).select(["A", "B"])

函数g（Series解耦但操作繁琐） #

def g(s_a: pl.Series, s_b: pl.Series, s_c: pl.Series) -> pl.DataFrame:
    mask = s_c > 7
    # 手动筛选每个Series再拼接，步骤冗余
    return pl.DataFrame({
        "A": s_a.filter(mask),
        "B": s_b.filter(mask)
    })

核心结论：pl.Expr完全适用于该场景 #

用pl.Expr实现的函数既能彻底解耦列名，又能避免g中的繁琐操作，同时保留Polars的向量化执行效率。

函数h的实现（基于pl.Expr） #

这里提供两种实用写法：

写法1：返回表达式列表（推荐） #

直接返回筛选后的目标列表达式，调用时无需额外处理：

def h(col_a: pl.Expr, col_b: pl.Expr, col_c: pl.Expr) -> list[pl.Expr]:
    # 自动保留原列名，无需硬编码
    return [
        col_a.filter(col_c > 7).alias(col_a.meta.output_name()),
        col_b.filter(col_c > 7).alias(col_b.meta.output_name())
    ]

调用示例

# 测试数据
df = pl.DataFrame({
    "A": [1, 2, 3, 8, 9],
    "B": [4, 5, 6, 10, 11],
    "C": [5, 6, 8, 9, 6]
})

# 调用h得到与f、g一致的结果
result = df.select(h(pl.col("A"), pl.col("B"), pl.col("C")))
print(result)
# 输出：
# shape: (2, 2)
# ┌─────┬──────┐
# │ A   ┆ B    │
# │ --- ┆ ---  │
# │ i64 ┆ i64  │
# ╞═════╪══════╡
# │ 3   ┆ 6    │
# │ 8   ┆ 10   │
# └─────┴──────┘

写法2：返回结构体表达式 #

适合需要将结果打包成单个结构体再展开的复杂场景：

def h(col_a: pl.Expr, col_b: pl.Expr, col_c: pl.Expr) -> pl.Expr:
    return pl.struct(
        col_a.alias(col_a.meta.output_name()),
        col_b.alias(col_b.meta.output_name())
    ).filter(col_c > 7)

调用示例

result = df.select(h(pl.col("A"), pl.col("B"), pl.col("C"))).unnest("struct")

为什么pl.Expr更合适？ #

• 彻底解耦列名：调用时可传入任意Expr（比如pl.col("X")、pl.col("Y")*2等），无需在函数内部硬编码列名，灵活性拉满。
• 消除冗余操作：无需手动筛选每个Series再拼接，所有逻辑由Polars表达式引擎自动处理，代码更简洁。
• 保持高效执行：所有操作都在Polars的查询引擎中完成，比手动处理Series的效率更高，尤其适合大数据集。

内容的提问来源于stack exchange，提问作者Benjamin Trendelkamp-Schroer