如何使用另一DataFrame对Polars DataFrame进行掩码操作？

场景说明 #

存在两个结构匹配的Polars DataFrame：

• 原始数据表df：time列为datetime类型，其余列为数值类型，数值列数量不固定
• 掩码表df_mask：和df列名完全一致，time列取值和df完全匹配，其余列为布尔类型，标记对应位置数值是否保留

要求掩码值为false的位置替换为null，true的位置保留原数值。
示例原始数据：

import polars as pl
df = pl.from_repr("""
┌─────────────────────┬─────────┬─────────┐
│ time                ┆ 1       ┆ 2       │
│ ---                 ┆ ---     ┆ ---     │
│ datetime[μs]        ┆ f64     ┆ f64     │
╞═════════════════════╪═════════╪═════════╡
│ 2021-10-02 00:05:00 ┆ 2.9048  ┆ 2.8849  │
│ 2021-10-02 00:10:00 ┆ 48224.0 ┆ 48068.0 │
└─────────────────────┴─────────┴─────────┘
""")

df_mask = pl.from_repr("""
┌─────────────────────┬───────┬───────┐
│ time                ┆ 1     ┆ 2     │
│ ---                 ┆ ---   ┆ ---   │
│ datetime[μs]        ┆ bool  ┆ bool  │
╞═════════════════════╪═══════╪═══════╡
│ 2021-10-02 00:05:00 ┆ false ┆ false │
│ 2021-10-02 00:10:00 ┆ true  ┆ true  │
└─────────────────────┴───────┴───────┘
""")

实现代码 #

场景1：两个表行顺序完全一致 #

如果能保证两个表的time列顺序、取值完全对应，不需要对齐行，直接批量处理所有非时间列即可，自动适配任意数量的数值列：

# 自动提取所有数值列（排除time主键列）
value_cols = [col for col in df.columns if col != "time"]

result = df.with_columns(
    # where方法会保留条件为true的值，其余位置自动置为null
    pl.col(c).where(df_mask[c]) for c in value_cols
)

场景2：无法保证两个表行顺序一致 #

如果两个表的行可能存在错位，先按time主键做join对齐，再执行掩码逻辑，避免行匹配错误：

value_cols = [col for col in df.columns if col != "time"]
# join对齐行，给掩码列加临时后缀避免列名冲突
df_joined = df.join(df_mask, on="time", suffix="_mask")

result = df_joined.with_columns(
    pl.col(c).where(pl.col(f"{c}_mask")).alias(c) for c in value_cols
# 最后选回原表的列结构，丢弃临时生成的掩码列
).select(df.columns)

执行后即可得到符合要求的结果，掩码为false的位置全部替换为null，true位置保留原数值，不需要手动指定任何数值列名。

注：示例输出中2021-10-02 00:10:00行的1列值标注为2.8849属于笔误，按逻辑该位置掩码为true应保留原值48224.0，上述代码完全遵循掩码保留原值的逻辑实现。

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