评估极坐标中自定义类内的表达式

Question

我正在尝试扩展

polars

的功能来操作 Enum 的类别。我正在遵循 this 指南和 this 文档部分

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())
    )

这会导致

TypeError: Series constructor called with unsupported type 'Expr' for the 

values

 parameter

出错，因为

cats

是一个未计算的表达式，而不是

pl.Enum(dtype=)

所期望的列表或系列。如何将

cats

评估到实际列表/系列中，以为我的

cast(pl.Enum)

方法提供新类别？

Answer 1

您可以使用

.map_batches()

@pl.api.register_expr_namespace('fct')
class CustomEnumMethodsCollection:
    def __init__(self, expr: pl.Expr):
        self._expr = expr
    
    def rev(self) -> pl.Expr:
        return self._expr.map_batches(lambda s:
            s.cast(pl.Enum(s.cat.get_categories().reverse()))
        )

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']))})

df.with_columns(rev_series=pl.col('orig_series').fct.rev()).schema

Schema([('idx', Int64),
        ('orig_series', Enum(categories=['Alpha', 'Beta', 'Gamma', 'Omega'])),
        ('rev_series', Enum(categories=['Omega', 'Gamma', 'Beta', 'Alpha']))])