在 polars 中,我可以自己创建一个带有级别的分类类型吗?

In polars, can I create a categorical type with levels myself?

在Pandas中,我可以自己指定分类类型的级别:

MyCat = pd.CategoricalDtype(categories=['A','B','C'], ordered=True)
my_data = pd.Series(['A','A','B'], dtype=MyCat)

这意味着

  1. 我可以确保不同的列和集合使用相同的数据类型
  2. 我可以指定级别的顺序。

有没有办法用 Polars 做到这一点?我知道您可以使用字符串缓存功能以不同的方式实现 1),但是我很想知道是否可以直接指定我的 dtype/levels。我不知道有什么方法可以实现 2),但是我认为 Arrow 中的分类 dtypes 确实允许可选的排序,所以也许这是可能的?

不直接,但我们可以影响全局字符串缓存的填充方式。全局字符串缓存只是为每个添加的新类别增加一个计数器。

因此,如果我们从一个空缓存开始,然后按照我们认为重要的顺序执行 pre-fill,后面的类别将使用缓存的整数。

这是一个例子:


import string
import polars as pl

with pl.StringCache():
    # the first run will fill the global string cache counting from 0..25
    # for all 26 letters in the alphabet
    pl.Series(list(string.ascii_uppercase)).cast(pl.Categorical)
    
    
    # now the global string cache is populated with all categories
    # we cast the string columns
    df = (pl.DataFrame({
        "letters": ["A", "B", "D"],
        "more_letters": ["Z", "B", "J"]
    }).with_column(pl.col(pl.Utf8).cast(pl.Categorical))
      .with_column(pl.col(pl.Categorical).to_physical().suffix("_real_category"))
     )

print(df)

shape: (3, 4)
┌─────────┬──────────────┬───────────────────────┬────────────────────────────┐
│ letters ┆ more_letters ┆ letters_real_category ┆ more_letters_real_category │
│ ---     ┆ ---          ┆ ---                   ┆ ---                        │
│ cat     ┆ cat          ┆ u32                   ┆ u32                        │
╞═════════╪══════════════╪═══════════════════════╪════════════════════════════╡
│ A       ┆ Z            ┆ 0                     ┆ 25                         │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ B       ┆ B            ┆ 1                     ┆ 1                          │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ D       ┆ J            ┆ 3                     ┆ 9                          │
└─────────┴──────────────┴───────────────────────┴────────────────────────────┘