屏蔽 Polars 数据帧以进行复杂操作

Question

如果我有一个极坐标数据框并想要执行屏蔽操作，我目前看到两个选项：

\n

# create data\ndf = pl.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], schema = [\'a\', \'b\']).lazy()\n# create a second dataframe for added fun\ndf2 = pl.DataFrame([[8, 6, 7, 5], [15, 16, 17, 18]], schema=["b", "d"]).lazy()\n\n# define mask\nmask = pl.col(\'a\').is_between(2, 3)\n

\n

选项 1：创建过滤后的数据帧，执行操作并连接回原始数据帧

\n

masked_df = df.filter(mask)\nmasked_df = masked_df.with_columns(  # calculate some columns\n    [\n        pl.col("a").sin().alias("new_1"),\n        pl.col("a").cos().alias("new_2"),\n        (pl.col("a") / pl.col("b")).alias("new_3"),\n    ]\n).join(  # throw a join into the mix\n    df2, on="b", how="left"\n)\nres = df.join(masked_df, how="left", on=["a", "b"])\nprint(res.collect())\n

\n

选项 2：单独屏蔽每个操作

\n

res = df.with_columns(  # calculate some columns - we have to add `pl.when(mask).then()` to each column now\n    [\n        pl.when(mask).then(pl.col("a").sin()).alias("new_1"),\n        pl.when(mask).then(pl.col("a").cos()).alias("new_2"),\n        pl.when(mask).then(pl.col("a") / pl.col("b")).alias("new_3"),\n    ]\n).join(  # we have to construct a convoluted back-and-forth join to apply the mask to the join\n    df2.join(df.filter(mask), on="b", how="semi"), on="b", how="left"\n)\n\nprint(res.collect())\n

\n

输出：

\n

shape: (4, 6)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 a   \xe2\x94\x86 b   \xe2\x94\x86 new_1    \xe2\x94\x86 new_2     \xe2\x94\x86 new_3    \xe2\x94\x86 d    \xe2\x94\x82\n\xe2\x94\x82 --- \xe2\x94\x86 --- \xe2\x94\x86 ---      \xe2\x94\x86 ---       \xe2\x94\x86 ---      \xe2\x94\x86 ---  \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64 \xe2\x94\x86 f64      \xe2\x94\x86 f64       \xe2\x94\x86 f64      \xe2\x94\x86 i64  \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1   \xe2\x94\x86 5   \xe2\x94\x86 null     \xe2\x94\x86 null      \xe2\x94\x86 null     \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x82 2   \xe2\x94\x86 6   \xe2\x94\x86 0.909297 \xe2\x94\x86 -0.416147 \xe2\x94\x86 0.333333 \xe2\x94\x86 16   \xe2\x94\x82\n\xe2\x94\x82 3   \xe2\x94\x86 7   \xe2\x94\x86 0.14112  \xe2\x94\x86 -0.989992 \xe2\x94\x86 0.428571 \xe2\x94\x86 17   \xe2\x94\x82\n\xe2\x94\x82 4   \xe2\x94\x86 8   \xe2\x94\x86 null     \xe2\x94\x86 null      \xe2\x94\x86 null     \xe2\x94\x86 null \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

大多数时候，选项 2 会更快，但它变得相当冗长，并且当涉及任何复杂性时通常比选项 1 更难阅读。

\n

有没有办法更通用地应用掩码来覆盖多个后续操作？

\n

Answer 1

您可以通过在辅助函数中将掩码应用到操作中来避免样板代码。

def with_mask(operations: list[pl.Expr], mask) -> list[pl.Expr]:
    return [
        pl.when(mask).then(operation)
        for operation in operations
    ]

res = df.with_columns(
    with_mask(
        [
            pl.col("a").sin().alias("new_1"),
            pl.col("a").cos().alias("new_2"),
            pl.col("a") / pl.col("b").alias("new_3"),
        ],
        mask,
    )
)