Маскирование кадра данных Polars для сложных операций ⇐ Python

[Anonymous]

Если у меня есть полярный Dataframe и я хочу выполнять маскированные операции, в настоящее время я вижу два варианта:

Код: Выделить всё

# create data
df = pl.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], schema = ['a', 'b']).lazy()
# create a second dataframe for added fun
df2 = pl.DataFrame([[8, 6, 7, 5], [15, 16, 17, 18]], schema=["b", "d"]).lazy()

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

Вариант 1: создать отфильтрованный фрейм данных, выполнить операции и вернуться к исходному фрейму данных

Код: Выделить всё

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

Вариант 2: маскировать каждую операцию индивидуально

Код: Выделить всё

res = df.with_columns(  # calculate some columns - we have to add `pl.when(mask).then()` to each column now
[
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(  # we have to construct a convoluted back-and-forth join to apply the mask to the join
df2.join(df.filter(mask), on="b", how="semi"), on="b", how="left"
)

print(res.collect())

Вывод:

Код: Выделить всё

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 │
└─────┴─────┴──────────┴───────────┴──────────┴──────┘

В большинстве случаев вариант 2 будет быстрее, но он становится довольно многословным и, как правило, его труднее читать, чем вариант 1, когда речь идет о какой-либо сложности.
Есть ли способ применить маску в более общем смысле, чтобы охватить несколько последующих операций?

Подробнее здесь: https://stackoverflow.com/questions/754 ... operations