如何在Polars中将枚举值序列拆分为索引与值两列（无需Python枚举）

用户尝试通过Polars创建枚举类型序列：

import polars as pl

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

后续创建DataFrame并尝试拆分操作失败：

date_table = pl.DataFrame(
    {"date": pl.Series(raw_dates, dtype=DateEnum)}
)
date_table.explode("date") # 执行失败

尝试直接转换类型也未成功：

date_table = pl.DataFrame({
    "date": pl.Series(raw_dates, dtype=DateEnum).cast(dtype=str),
    "index": pl.Series(raw_dates, dtype=DateEnum).cast(dtype=int),
})
date_table # 执行失败

用户期望生成如下结构的DataFrame（当前通过Python原生Enum实现）：

from typing import Enum

pyDateEnum = Enum("Dates", 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     │
└──────────┴───────┘

需要实现：不依赖Python原生Enum，仅用Polars达成该需求。

直接利用Polars枚举的内置属性和函数即可实现，无需原生Enum。

方法一：直接构造序列 #

import polars as pl

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

date_table = pl.DataFrame({
    "date": pl.Series(DateEnum.names(), dtype=str),
    "index": pl.Series(range(len(DateEnum.names())), dtype=int)
})

print(date_table)

方法二：Polars风格的列生成 #

import polars as pl

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

date_table = pl.DataFrame({"date": DateEnum.names()}) \
    .with_columns(index=pl.int_range(0, pl.count(), dtype=int))

print(date_table)

关键说明 #

• DateEnum.names()：Polars枚举对象的names()方法会直接返回所有枚举成员的名称列表，正好对应需求中的date列。
• 索引列：可以用Python原生range()生成从0开始的连续整数，或者用Polars的pl.int_range()函数在DataFrame中动态计算，后者更适合处理大规模数据集。

两种方法最终都会输出用户期望的结果：

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

内容的提问来源于stack exchange，提问作者bzm3r