Если у меня есть фрейм данных polars и я хочу выполнять маскированные операции, в настоящее время я вижу два варианта:
# 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)
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())
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, когда речь идет о какой-либо сложности.
Есть ли способ применить маску более широко, чтобы охватить несколько последовательных операций?
Вы можете избежать шаблонной пластины, применив маску к своим операциям во вспомогательной функции.
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,
)
)
Это действительно красиво выглядит. Знаете ли вы, способен ли ленивый оптимизатор обнаружить повторяющуюся маску when и выполнить ее только один раз, или он повторит эту операцию для всех столбцов?
Вы можете использовать struct
с unnest
Ваши dfs не соответствовали между ленивостью и нетерпеливостью, поэтому я собираюсь сделать их обоих ленивыми.
df.join(df2, on='b') \
.with_columns(pl.when(mask).then(
pl.struct([
pl.col("a").sin().alias("new_1"),
pl.col("a").cos().alias("new_2"),
(pl.col("a") / pl.col("b").cast(pl.Float64())).alias("new_3")
]).alias('allcols'))).unnest('allcols') \
.with_columns([pl.when(mask).then(pl.col(x)).otherwise(None)
for x in df2.columns if x not in df]) \
.collect()
Я думаю, что суть вашего вопроса заключается в том, как написать when
then
с выходом нескольких столбцов, который покрывается первым with_columns
, а затем второй with_columns
охватывает поведение замены значения квази-полуобъединения.
Другой способ написать это - сначала создать список столбцов в df2, которые вы хотите подвергнуть маске, а затем поместить их в структуру. Некрасиво то, что вам нужно исключить эти столбцы, прежде чем вы сделаете unnest
df2_mask_cols=[x for x in df2.columns if x not in df.columns]
df.join(df2, on='b') \
.with_columns(pl.when(mask).then(
pl.struct([
pl.col("a").sin().alias("new_1"),
pl.col("a").cos().alias("new_2"),
(pl.col("a") / pl.col("b").cast(pl.Float64())).alias("new_3")
] + df2_mask_cols).alias('allcols'))) \
.select(pl.exclude(df2_mask_cols)) \
.unnest('allcols') \
.collect()
Удивительно, но этот подход оказался самым быстрым:
df.join(df2, on='b') \
.with_columns([
pl.col("a").sin().alias("new_1"),
pl.col("a").cos().alias("new_2"),
(pl.col("a") /pl.col("b").cast(pl.Float64())).alias("new_3")
]) \
.with_columns(pl.when(mask).then(pl.exclude(df.columns))).collect()
После выполнения некоторых тестов оказалось, что самым быстрым вариантом был тот, который вы отредактировали, где сначала выполняются вычисления, а маска используется только для очистки замаскированных строк постфактум. Однако я предполагаю, что производительность будет сильно различаться в зависимости от типов данных, свойств замаскированных строк и вычислений.
@DataWiz, это довольно удивительно ... думаю, я отредактирую его обратно.
Спасибо - я тоже думал о понимании списков, и вы объяснили это для меня. Есть идеи для присоединения?