Кадр данных Polars: как эффективно агрегировать по множеству непересекающихся групп ⇐ Python

[Anonymous]

У меня есть кадр данных со столбцами x, y, c_1, c2, ..., c_K, где K несколько велико (K ≈ 1000 или 2000).
Каждый из столбцов c_i является логическим столбцом, и я хотел бы вычислить агрегацию f(x, y) по строки, где c_i имеет значение True. (Например, f(x,y) = x.sum() * y.sum().)
Один из способов сделать это:
ds.select([
f(pl.col("x").filter(pl.col(f"c_{i+1}"), pl.col("y").filter(pl.col(f"c_{i+1}"))
for i in range(K)
])

В моей задаче число K велико, и приведенный выше запрос кажется несколько неэффективным (фильтрация выполняется дважды).

• Какой рекомендуемый/самый эффективный/самый элегантный способ добиться этого?

< strong>РЕДАКТИРОВАТЬ.
Вот работоспособный пример (код внизу), а также некоторые тайминги, соответствующие ответу @Hericks ниже. TLDR: предложенный Метод 1 является лучшим на данный момент.





Время на стене




1
повторяется фильтры
409 мс


2
pl.concat
29,6 с (≈70 раз медленнее)


2*
pl.concat, ленивый
1,27 с (в 3 раза медленнее)


3
расплавить с помощью агг
1мин 17с


3*
растаяние с помощью агг, ленивый
1 минута 17 секунд (то же, что и 3)



import polars as pl
import polars.selectors as cs
import numpy as np
rng = np.random.default_rng()

def f(x,y):
return x.sum() * y.sum()

N = 2_000_000
K = 1000
dat = dict()
dat["x"] = np.random.randn(N)
dat["y"] = np.random.randn(N)
for i in range(K):
dat[f"c_{i+1}"] = rng.choice(2, N).astype(np.bool_)

tmpds = pl.DataFrame(dat)

## Method 1
tmpds.select([
f(
pl.col("x").filter(pl.col(f"c_{i+1}")),
pl.col("y").filter(pl.col(f"c_{i+1}")))
.alias(f"f_{i+1}") for i in range(K)
])

## Method 2
pl.concat([
tmpds.filter(pl.col(f"c_{i+1}")).select(f(pl.col("x"), pl.col("y")).alias(f"f_{i+1}"))
for i in range(K)
], how="horizontal")

## Method 2*
pl.concat([
tmpds.lazy().filter(pl.col(f"c_{i+1}")).select(f(pl.col("x"), pl.col("y")).alias(f"f_{i+1}")).collect()
for i in range(K)
], how="horizontal")

## Method 3
(
tmpds
.unpivot(on=cs.starts_with("c"), index=["x", "y"])
.filter("value")
.group_by("variable")
.agg(
f(pl.col("x"), pl.col("y"))
)
)

##Method 3*
(
tmpds
.lazy()
.unpivot(on=cs.starts_with("c"), index=["x", "y"])
.filter("value")
.group_by("variable", maintain_order=True)
.agg(
f(pl.col("x"), pl.col("y"))
)
.collect()
)


Подробнее здесь: https://stackoverflow.com/questions/790 ... int-groups