Polars自定义类中评估表达式实现Enum类别反转的问题

你想扩展Polars的表达式命名空间来操作Enum类型，写了自定义的fct.rev()方法，目的是反转Enum类别的字符串并转换为新的Enum类型，但运行时触发以下错误：

TypeError: Series构造函数的values参数不支持'Expr'类型

问题出在self._expr.cat.get_categories()返回的是未评估的表达式（Expr），而pl.Enum()需要传入实际的类别列表或Series，根本没法直接接收Expr类型。

你的代码示例：

orig_df = pl.DataFrame({
    'idx': pl.int_range(5, eager=True),
    'orig_series': pl.Series(['Alpha', 'Omega', 'Alpha', 'Beta', 'Gamma'], 
                dtype=pl.Enum(['Alpha', 'Beta', 'Gamma', 'Omega']))})

@pl.api.register_expr_namespace('fct')
class CustomEnumMethodsCollection:
    def __init__(self, expr: pl.Expr):
        self._expr = expr
    
    def rev(self) -> pl.Expr:
        cats = self._expr.cat.get_categories()
        tmp_sr = self._expr.cast(pl.Categorical)
        return tmp_sr.cast(dtype=pl.Enum(cats.str.reverse()))

(orig_df
    .with_columns(rev_series=pl.col("orig_series").fct.rev())
    )

解决办法 #

Polars的表达式是延迟执行的，在自定义表达式方法里没法直接拿到未评估Expr对应的实际数据（比如Enum的类别）。要实现需求，有两种可行思路：

方法1：用map_batches批量处理

修改自定义方法，通过map_batches对每个批次的Series单独处理，这样就能拿到实际的类别列表：

@pl.api.register_expr_namespace('fct')
class CustomEnumMethodsCollection:
    def __init__(self, expr: pl.Expr):
        self._expr = expr
    
    def rev(self) -> pl.Expr:
        def _rev_enum(s: pl.Series) -> pl.Series:
            # 获取当前批次Series的Enum类别
            orig_cats = s.cat.get_categories().to_list()
            # 反转每个类别字符串
            reversed_cats = [cat[::-1] for cat in orig_cats]
            # 创建新的Enum类型
            new_dtype = pl.Enum(reversed_cats)
            # 把原值反转后转成新Enum
            reversed_vals = s.cast(pl.Utf8).str.reverse()
            return reversed_vals.cast(new_dtype)
        
        return self._expr.map_batches(_rev_enum)

方法2：提前获取类别再构造表达式

如果你的数据类别是固定的，也可以先提前拿到原列的类别，手动构造反转后的Enum类型，再做转换：

# 提前获取原Enum的类别列表
orig_cats = orig_df['orig_series'].cat.get_categories().to_list()
# 反转每个类别字符串
reversed_cats = [cat[::-1] for cat in orig_cats]
# 创建新的Enum dtype
new_enum_dtype = pl.Enum(reversed_cats)

# 转换列
result_df = orig_df.with_columns(
    rev_series=pl.col("orig_series").cast(pl.Utf8).str.reverse().cast(new_enum_dtype)
)

补充说明 #

• 方法1更通用，适合延迟执行的场景，能处理动态变化的类别；
• 方法2执行效率更高，适合类别固定、不需要动态处理的场景；
• 两种方法都避开了直接把未评估的Expr传给pl.Enum()的问题。

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