根据列聚合过滤选定的列

Question

我希望仅选择唯一值少于 3 个的列。我可以通过

pl.all().n_unique() < 3

生成布尔掩码，但我不知道是否可以通过 Polars API 使用该掩码。

目前我正在通过python解决这个问题。有更惯用的方法吗？

import polars as pl, pandas as pd
df = pl.DataFrame({"col1":[1,1,2], "col2":[1,2,3], "col3":[3,3,3]})
# target is:
# df_few_unique = pl.DataFrame({"col1":[1,1,2], "col3":[3,3,3]})

# my attempt:
mask = df.select(pl.all().n_unique() < 3).to_numpy()[0]
cols = [col for col, m in zip(df.columns, mask) if m]
df_few_unique = df.select(cols)
df_few_unique

相当于 pandas：

df_pandas = df.to_pandas()
mask = (df_pandas.nunique() < 3)
df_pandas.loc[:, mask]

Answer 1

编辑：经过一番思考，我发现了一种更简单的方法来做到这一点，一种根本不依赖于布尔掩码的方法。

pl.select(
    [s for s in df
     if s.n_unique() < 3]
)

shape: (3, 2)
┌──────┬──────┐
│ col1 ┆ col3 │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 3    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 1    ┆ 3    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 2    ┆ 3    │
└──────┴──────┘

之前的回答

一种简单的方法是使用 Python 的

compress

中的

itertools

函数。

from itertools import compress
df.select(compress(df.columns, df.select(pl.all().n_unique() < 3).row(0)))

>>> df.select(compress(df.columns, df.select(pl.all().n_unique() < 3).row(0)))
shape: (3, 2)
┌──────┬──────┐
│ col1 ┆ col3 │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 3    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 1    ┆ 3    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 2    ┆ 3    │
└──────┴──────┘

compress

允许我们将布尔掩码应用于列表，在本例中是列名称列表。

list(compress(df.columns, df.select(pl.all().n_unique() < 3).row(0)))

>>> list(compress(df.columns, df.select(pl.all().n_unique() < 3).row(0)))
['col1', 'col3']

Answer 2

所选答案虽然语法清晰，但效率极低。您可以做得更好大约 400 倍（如果您的机器上有多个过滤器和更多核心，效果会更好）

让我们首先包含至少两个过滤器，而不仅仅是一个

问题：仅选择唯一值数量在 1 到 200 之间的列

TL； DR - 理想的解决方案 -

df.lazy().select((pl.all().n_unique().is_between(1,200)))

继续阅读了解详情

要考虑的事情是，无论如何您都需要传递数据。所以，阅读它是第一步

那么，如果你这样做

pl.select(
    [s for s in df
     if s.n_unique() < 3]
)

您正在按顺序计算过滤器并将它们保存在内存中 理想的解决方案是在利用 Polars 查询规划优化的同时并行完成所有工作。

让我们做一些基准测试。我用的是4核机器。并行性将进一步减少具有更多内核的机器上的时间

设置数据框：

import polars as pl
import numpy as np
df = pl.DataFrame({f'a_{i}':np.random.choice(['a','b','c','d'], 100000) for i in range(10000)})

这将占用大约 20 GiB RAM。所以，如果你想复制的话要小心

1. 所选方案（htop确认该方案仅使用一个核心）

%%timeit
_df = pl.select(
    [s for s in df
     if s.n_unique() < 3]
)
output:
8.77 s ± 59 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

现在让我们尝试并行运行过滤器（htop 确认）

%%timeit
_df = df.select((pl.all().n_unique() < 200) & (pl.all().n_unique() > 1))
output:
7.73 s ± 114 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

但是略有改进，我们在上面的两个

.n_unique()

调用中计算每个过滤器两次。让我们使用

in_between

来完成一个（并行执行 - htop 确认）

%%timeit
_df = df.select((pl.all().n_unique().is_between(1,200)))
output:
3.95 s ± 47.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

现在让我们尝试看看使用查询执行优化的理想解决方案：

首先，激活 df 的惰性语义

df_lazy = df.lazy()

现在，尝试以上两种解决方案

%%timeit
_df = df_lazy.select((pl.all().n_unique() < 200) & (pl.all().n_unique() > 1))
output:
38.4 ms ± 720 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

• 请注意，上述解决方案不会计算过滤器两次，即使代码结构建议如此。在惰性模式下，查询计划优化可确保相同的事情不会计算两次

但是，使用

in_between

直接 Rust 实现会稍微快一点

%%timeit
_df = df_lazy.select((pl.all().n_unique().is_between(1,200)))
output:
24 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

如您所见，理想的解决方案比公认的解决方案快了近 400 倍