根据列聚合过滤选定的列

我希望仅选择唯一值少于 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

相当于熊猫：

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

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

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    │
└──────┴──────┘

之前的回答

一种简单的方法是使用compress来自 Python 的函数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']