Python + Polars：如何将enum类型列的值转换为整数（物理）表示？

Polars的Enum类型确实以整数作为物理存储，你可以直接通过内置方法获取对应的整数值，无需使用逐元素的map_elements操作——后者不仅代码繁琐，还因为是Python层面的逐元素处理，性能远不如内置矢量化方法。

优化后的代码示例 #

import numpy as np
import polars as pl

np.random.seed(556)

enum_vals = [
    "".join([chr(c_code) for c_code in np.random.randint(97, 123, 3)])
    for n in range(10)
]
enum_dtype = pl.Enum(pl.Series(enum_vals))

(
    pl.Series(
        "enum_vals",
        [enum_vals[x] for x in np.random.randint(0, len(enum_vals), 5)],
        dtype=enum_dtype,
    )
    .to_frame()
    .with_columns(
        # 直接调用to_physical()获取底层整数表示
        enum_repr=pl.col("enum_vals").to_physical()
    )
)

运行结果 #

shape: (5, 2)
┌───────────┬───────────┐
│ enum_vals ┆ enum_repr │
│ ---       ┆ ---       │
│ enum      ┆ u32       │
╞═══════════╪═══════════╡
│ loo       ┆ 8         │
│ sby       ┆ 5         │
│ cqm       ┆ 3         │
│ cbn       ┆ 2         │
│ vtk       ┆ 9         │
└───────────┴───────────┘

补充说明 #

• to_physical()默认返回UInt32类型（Enum的底层存储类型），如果需要转换成Int64或其他整数类型，可追加cast()操作：

  enum_repr=pl.col("enum_vals").to_physical().cast(pl.Int64())
  

• 该方法是底层矢量化实现，处理大数据集时的性能优势会非常明显。

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