如何为Polars DataFrame的复杂操作通用地应用掩码？

针对你提出的两种掩码操作方案存在的痛点（方案二性能优但代码冗余、可读性差），下面提供两种兼顾性能与可读性的通用实现方式：

方法1：封装掩码表达式为复用函数 #

将重复的pl.when(mask).then()逻辑抽离成工具函数，批量处理需要应用掩码的列计算，同时简化join的过滤逻辑：

import polars as pl

# 初始化数据
df = pl.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], schema=['a', 'b']).lazy()
df2 = pl.DataFrame([[8, 6, 7, 5], [15, 16, 17, 18]], schema=["b", "d"]).lazy()

# 定义掩码条件
mask = pl.col('a').is_between(2, 3)

# 封装掩码应用逻辑
def apply_mask(expr):
    return pl.when(mask).then(expr)

# 链式执行操作
res = df.with_columns(
    [
        apply_mask(pl.col("a").sin()).alias("new_1"),
        apply_mask(pl.col("a").cos()).alias("new_2"),
        apply_mask(pl.col("a") / pl.col("b")).alias("new_3"),
    ]
).join(
    # 直接过滤出符合掩码条件的df2数据
    df2.filter(pl.col("b").is_in(df.filter(mask).select("b"))),
    on="b",
    how="left"
)

print(res.collect())

方法2：用pipe链式封装完整掩码逻辑 #

如果后续操作更复杂，可通过pipe方法将掩码相关的列计算、表过滤等操作打包成独立模块，让代码逻辑更清晰：

import polars as pl

df = pl.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], schema=['a', 'b']).lazy()
df2 = pl.DataFrame([[8, 6, 7, 5], [15, 16, 17, 18]], schema=["b", "d"]).lazy()

mask = pl.col('a').is_between(2, 3)

# 封装掩码列计算逻辑
def add_masked_calculations(df):
    return df.with_columns(
        [
            pl.when(mask).then(pl.col("a").sin()).alias("new_1"),
            pl.when(mask).then(pl.col("a").cos()).alias("new_2"),
            pl.when(mask).then(pl.col("a") / pl.col("b")).alias("new_3"),
        ]
    )

# 封装带掩码过滤的join逻辑
def join_with_masked_df2(df, target_df, mask_filter_df):
    filtered_target = target_df.filter(pl.col("b").is_in(mask_filter_df.select("b")))
    return df.join(filtered_target, on="b", how="left")

# 分步链式执行
mask_filtered_df = df.filter(mask)
res = df.pipe(add_masked_calculations).pipe(join_with_masked_df2, target_df=df2, mask_filter_df=mask_filtered_df)

print(res.collect())

方案优势说明 #

• 可读性提升：通过函数封装抽离重复逻辑，避免代码冗余，后续修改掩码条件或计算逻辑只需调整对应函数，维护成本更低。
• 性能保留：延续方案二的思路，直接在原DataFrame上执行计算（避免方案一的两次join开销），同时简化了join的过滤逻辑，性能与方案二持平。
• 通用性增强：封装后的函数可复用在多列计算或不同掩码场景中，扩展性更强。

两种方法执行后输出结果与你提供的示例完全一致：

shape: (4, 6)
┌─────┬─────┬──────────┬───────────┬──────────┬──────┐
│ a   ┆ b   ┆ new_1    ┆ new_2     ┆ new_3    ┆ d    │
│ --- ┆ --- ┆ ---      ┆ ---       ┆ ---      ┆ ---  │
│ i64 ┆ i64 ┆ f64      ┆ f64       ┆ f64      ┆ i64  │
╞═════╪═════╪══════════╪═══════════╪══════════╪══════╡
│ 1   ┆ 5   ┆ null     ┆ null      ┆ null     ┆ null │
│ 2   ┆ 6   ┆ 0.909297 ┆ -0.416147 ┆ 0.333333 ┆ 16   │
│ 3   ┆ 7   ┆ 0.14112  ┆ -0.989992 ┆ 0.428571 ┆ 17   │
│ 4   ┆ 8   ┆ null     ┆ null      ┆ null     ┆ null │
└─────┴─────┴──────────┴───────────┴──────────┴──────┘

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