Оценка выражения внутри пользовательского класса в polars

На чтение 4 мин Опубликовано 11.11.2024

Содержание

Вопрос или проблема
Ответ или решение
Контекст задачи
Решение проблемы
Использование evaluate()
Объяснение кода
Заключение

Вопрос или проблема

Я пытаюсь расширить функциональность polars для работы с категориями Enum. Я следую этому руководству и этому разделу документации.

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 вызван с неподдерживаемым типом 'Expr' для параметра values, потому что cats является неоцененным выражением, а не списком или серией, как того ожидает pl.Enum(dtype=). Как мне оценить cats в фактический список/серию, чтобы предоставить новые категории для моего метода cast(pl.Enum)?


.


Ответ или решение
Вопрос, который вы подняли, касается расширения функциональности библиотеки Polars для работы с перечислениями (Enum). В вашем коде вы пытаетесь создать метод для реверса (обратного порядка) категорий в столбце с перечислением, но сталкиваетесь с ошибкой, связанной с тем, что выражение cats не может быть использовано как ожидаемый список для создания нового типа данных Enum. Давайте детально разберем, как решить эту проблему.
Контекст задачи
В вашей задаче вы расширяете функциональность библиотеки Polars, создавая класс, который регистрирует пространство имен с методами работы с перечислениями. Проблема возникает на этапе, когда необходимо извлечь категории из выражения Expr, чтобы преобразовать их в тип Enum.
Решение проблемы
Чтобы решить проблему с ошибкой TypeError: Series constructor called with unsupported type 'Expr' for the values parameter, необходимо убедиться, что категории, которые вы пытаетесь использовать для создания нового перечисления, являются оцененными значениями, а не неразрешенными выражениями. Ваша текущая реализация не позволяет получить значения категорий, так как вы пытаетесь напрямую использовать выражение.
Использование evaluate()
Одним из подходов к решению данной проблемы является использование метода .evaluate() для оценки категорий. Так как в Polyars вы работаете с ленивыми выражениями, вам нужно будет выполнить их, например, с использованием метода collect().
Пример исправленного кода может выглядеть так:
import polars as pl

# Исходный DataFrame с перечислением
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_expr = self._expr.cat.get_categories().collect()
        cats = cats_expr.to_list()  # Преобразуем в список
        tmp_sr = self._expr.cast(pl.Categorical)
        return tmp_sr.cast(dtype=pl.Enum(cats[::-1]))  # Реверсируем порядок категорий

# Применение метода rev()
result_df = (orig_df
    .with_columns(rev_series=pl.col("orig_series").fct.rev())
)
Объяснение кода


Извлечение категорий: Вместо использования get_categories() напрямую, мы вызываем .collect(), чтобы получить реальные значения категорий. Это преобразует выражение в результат — DataFrame, из которого мы можем получить категории в виде списка.


Обратный порядок категорий: Порядок категорий реверсируется с помощью cats[::-1], что позволяет создать новый Enum с обратным порядком.


Кастинг: Осуществляется преобразование tmp_sr обратно в Enum с новыми категориями.


Заключение
Данная проблема является распространенной при работе с ленивыми выражениями в библиотеках обработки данных, таких как Polars. Использование методов для оценки выражений и их преобразование в списки позволяет избежать таких проблем, как указана в вашем вопросе.
Если у вас возникнут дополнительные вопросы или потребуется дальнейшая помощь в расширении функциональности Polars, не стесняйтесь обращаться за поддержкой!