Маскировка фрейма данных polars для сложных операций

Если у меня есть фрейм данных 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)

Вариант 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, когда речь идет о какой-либо сложности.

Есть ли способ применить маску более широко, чтобы охватить несколько последовательных операций?

Почему в Python есть оператор "pass"?
Почему в Python есть оператор "pass"?
Оператор pass в Python - это простая концепция, которую могут быстро освоить даже новички без опыта программирования.
Некоторые методы, о которых вы не знали, что они существуют в Python
Некоторые методы, о которых вы не знали, что они существуют в Python
Python - самый известный и самый простой в изучении язык в наши дни. Имея широкий спектр применения в области машинного обучения, Data Science,...
Основы Python Часть I
Основы Python Часть I
Вы когда-нибудь задумывались, почему в программах на Python вы видите приведенный ниже код?
LeetCode - 1579. Удаление максимального числа ребер для сохранения полной проходимости графа
LeetCode - 1579. Удаление максимального числа ребер для сохранения полной проходимости графа
Алиса и Боб имеют неориентированный граф из n узлов и трех типов ребер:
Оптимизация кода с помощью тернарного оператора Python
Оптимизация кода с помощью тернарного оператора Python
И последнее, что мы хотели бы показать вам, прежде чем двигаться дальше, это
Советы по эффективной веб-разработке с помощью Python
Советы по эффективной веб-разработке с помощью Python
Как веб-разработчик, Python может стать мощным инструментом для создания эффективных и масштабируемых веб-приложений.
1
0
101
2
Перейти к ответу Данный вопрос помечен как решенный

Ответы 2

Вы можете избежать шаблонной пластины, применив маску к своим операциям во вспомогательной функции.


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,
    )
)

Спасибо - я тоже думал о понимании списков, и вы объяснили это для меня. Есть идеи для присоединения?

DataWiz 16.02.2023 20:39

Это действительно красиво выглядит. Знаете ли вы, способен ли ленивый оптимизатор обнаружить повторяющуюся маску when и выполнить ее только один раз, или он повторит эту операцию для всех столбцов?

Dean MacGregor 16.02.2023 21:31
Ответ принят как подходящий

Вы можете использовать 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()

Я думаю, что суть вашего вопроса заключается в том, как написать whenthen с выходом нескольких столбцов, который покрывается первым 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 17.02.2023 13:13

@DataWiz, это довольно удивительно ... думаю, я отредактирую его обратно.

Dean MacGregor 17.02.2023 13:15

Другие вопросы по теме