【Polars】ユニーク数をカウントする時にNullを除外する

やりたいこと

groupbyでとあるカラムのユニーク数をカウントする際に、欠損値を除外したい。

結論

pl.col('hoge').n_unique() - pl.col('hoge').unique().null_count() で算出する。

import polars as pl

df = pl.DataFrame(
    {
        'foo': [1, 1, 2, 2, 2, 2]
	    , 'bar': [6, 7, None, 5, 8, None]
    }
)
print(df)
# shape: (6, 2)
# ┌─────┬──────┐
# │ foo ┆ bar  │
# │ --- ┆ ---  │
# │ i64 ┆ i64  │
# ╞═════╪══════╡
# │ 1   ┆ 6    │
# │ 1   ┆ 7    │
# │ 2   ┆ null │
# │ 2   ┆ 5    │
# │ 2   ┆ 8    │
# │ 2   ┆ null │
# └─────┴──────┘


agg = df.groupby('foo', maintain_order=True).agg(
    cnt=pl.col('bar').count() # nullも含めた件数が返ってくる
    , unique_cnt=pl.col('bar').n_unique() # nullも含めたユニーク件数が返ってくる
    , unique_cnt_non_null=pl.col('bar').n_unique() - pl.col('bar').unique().null_count()
)

print(agg)
# shape: (2, 4)
# ┌─────┬─────┬────────────┬─────────────────────┐
# │ foo ┆ cnt ┆ unique_cnt ┆ unique_cnt_non_null │
# │ --- ┆ --- ┆ ---        ┆ ---                 │
# │ i64 ┆ u32 ┆ u32        ┆ u32                 │
# ╞═════╪═════╪════════════╪═════════════════════╡
# │ 1   ┆ 2   ┆ 2          ┆ 2                   │
# │ 2   ┆ 4   ┆ 3          ┆ 2                   │
# └─────┴─────┴────────────┴─────────────────────┘

環境

import polars as pl

print(pl.__version__)
# ==> '0.18.3'