Как перевести операции DataFrame pandas в Polars в Python?

Я пытаюсь преобразовать некоторые операции DataFrame pandas в Polars в Python, но сталкиваюсь с трудностями, особенно с операциями по строкам и сравнениями по элементам. Вот код панд, с которым я работаю:

df_a = pd.DataFrame({
    "feature1": [1, 2, 3],
    "feature2": [7, 8, 9],
})

df_b = pd.DataFrame({
    "feature1": [3, 8, 2],
    "feature2": [7, 4, 9],
})

if selection_mode == 'option1':
    max_values = df_a.max(axis=1)
    selected_features = df_a.eq(max_values, axis=0)
    final_result = selected_features.mul(df_b).sum(axis=1) / selected_features.sum(axis=1)

elif selection_mode == 'option2':
    above_avg = df_a.ge(df_a.mean(axis=1), axis=0)
    combined_df = above_avg.mul(df_a).mul(df_b)
    sum_combined = combined_df.sum(axis=1)
    sum_above_avg = above_avg.mul(df_a).sum(axis=1)
    final_result = sum_combined / sum_above_avg

Будем очень признательны за любые рекомендации по переводу этого кода pandas в Polars!

🤔 А знаете ли вы, что...
С Python можно создавать ботов для социальных сетей и мессенджеров.


2
52
2

Ответы:

Вариант 1 должен быть эквивалентен маскировке (с использованием concat и When ) и mean_horizontal:

(pl.concat([dfb, dfa.with_columns(pl.all().eq(dfa.max_horizontal()))
                    .select(pl.all().name.suffix('_mask'))],
           how='horizontal')
   .select(pl.when(pl.col(f'{c}_mask')).then(pl.col(c)) for c in dfb.columns)
   .mean_horizontal()
)

Выход:

Series: 'mean' [f64]
[
    7.0
    4.0
    9.0
]

Вариант 2 немного сложнее, поскольку вам нужно выполнить два вычисления:

(pl.concat([dfa.select(pl.all().name.suffix('_a')),
            dfb.select(pl.all().name.suffix('_b')),
            dfa.with_columns(pl.all().ge(dfa.mean_horizontal()))
               .select(pl.all().name.suffix('_mask'))],
           how='horizontal')
   .select(pl.sum_horizontal(pl.when(pl.col(f'{c}_mask')).then(pl.col(f'{c}_a')*pl.col(f'{c}_b')).alias(c) for c in dfb.columns)
             .truediv(pl.sum_horizontal(pl.when(pl.col(f'{c}_mask')).then(pl.col(f'{c}_a')).alias(c) for c in dfb.columns))
             .alias('final')
          )['final']
)

Выход:

Series: 'final' [f64]
[
    7.0
    4.0
    9.0
]

Решено

В Polars выделены горизонтальные функции для «построчных» операций.

df_a.max_horizontal()
shape: (3,)
Series: 'max' [i64]
[
    7
    8
    9
]

Для DataFrames Polars «транслирует» операцию по всем столбцам, если правая часть представляет собой серию.

df_a == df_a.max_horizontal() # df_a.select(pl.all() == pl.Series([7, 8, 9]))
shape: (3, 2)
┌──────────┬──────────┐
│ feature1 ┆ feature2 │
│ ---      ┆ ---      │
│ bool     ┆ bool     │
╞══════════╪══════════╡
│ false    ┆ true     │
│ false    ┆ true     │
│ false    ┆ true     │
└──────────┴──────────┘

Вариант №1

max_values = df_a.max_horizontal()
selected_features = df_a == max_values

final_result = (
    (selected_features * df_b).sum_horizontal() / selected_features.sum_horizontal()
)

Вариант №2

above_avg = df_a >= df_a.mean_horizontal()
combined_df = above_avg * df_a * df_b
sum_combined = combined_df.sum_horizontal()
sum_above_avg = (above_avg * df_a).sum_horizontal()

final_result = sum_combined / sum_above_avg