Filtering selected columns based on column aggregate

user_1018861 0
Viewed 63

I wish to select only columns with fewer than 3 unique values. I can generate a boolean mask via pl.all().n_unique() < 3, but I don't know if I can use that mask via the polars API for this.

Currently, I am solving it via python. Is there a more idiomatic way?

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

Equivalent in pandas:

df_pandas = df.to_pandas()
mask = (df_pandas.nunique() < 3)
df_pandas.loc[:, mask]
python-polars
1 Answers

Edit: after some thinking, I discovered an even easier way to do this, one that doesn't rely on boolean masking at all.

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

Previous answer

One easy way is to use the compress function from Python's 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 allows us to apply a boolean mask to a list, which in this case is a list of column names.

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']
user_185598750
Related