Polars：将枚举值序列“分解”为枚举索引和枚举值的两列

Question

import polars as pl

dates = ["2024 Jan", "2024 Feb"]
DateEnum = pl.Enum(dates)

date_table = pl.DataFrame(
    {"date": pl.Series(raw_dates, dtype=DateEnum)}
)
date_table.explode("date") # fails

    {
        "date": pl.Series(raw_dates, dtype=plDateEnum).cast(dtype=str),
        "index": pl.Series(raw_dates, dtype=plDateEnum).cast(dtype=int),
    }
)
date_table # fails

我希望看到这样的数据框：

from typing import Enum

pyDateEnum = Enum("Dates", raw_dates)
date_table = pl.DataFrame({
  "date": pl.Series("dates", [x.name for x in pyDateEnum], dtype=str),
  "index": pl.Series("indices", [x.value - 1 for x in pyDateEnum], dtype=int)
})
date_table

shape: (2, 2)
┌──────────┬───────┐
│ date     ┆ index │
│ ---      ┆ ---   │
│ str      ┆ i64   │
╞══════════╪═══════╡
│ 2024 Jan ┆ 0     │
│ 2024 Feb ┆ 1     │
└──────────┴───────┘

如何使用 Polars 来完成此操作，而不求助于 Python 枚举？

Answer 1

我不确定是否有更原生的方法可以做到这一点，但你总是可以

cast

：

date_table.with_columns(
    name = pl.col.date.cast(str),
    value = pl.col.date.cast(pl.UInt32)
)

┌──────────┬──────────┬───────┐
│ date     ┆ name     ┆ value │
│ ---      ┆ ---      ┆ ---   │
│ enum     ┆ str      ┆ u32   │
╞══════════╪══════════╪═══════╡
│ 2024 Jan ┆ 2024 Jan ┆ 0     │
│ 2024 Feb ┆ 2024 Feb ┆ 1     │
└──────────┴──────────┴───────┘