主页 / user-1890061

Kevin's questions

Martin Hope
Kevin
Asked: 2024-09-27 09:17:40 +0800 CST
  • 8

我有一个包含列的数据框，其中xK有点大（K ≈ 1000 或 2000）。yc_1c2c_K

每列都是c_i布尔列，我想计算f(x, y)为 True 的行的聚合c_i。（例如，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很大，并且上述查询似乎效率低下（过滤进行了两次）。

  • 实现这一目标的推荐/最有效/最优雅的方法是什么？

编辑。

这是一个可运行的示例（代码在底部），以及一些与下面@Hericks 的答案相对应的时间。TLDR：建议的方法 1是目前最好的。

墙上时间
1 重复过滤 409毫秒
2 pl.concat 29.6秒（≈慢70倍）
2* pl.concat， 懒惰的 1.27 秒（慢 3 倍）
3 与骨料一起融化 1分17秒
3* 融化成 agg，懒惰 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()
)
python