Типизация UDF map_batches в Polars v1.32+: Array, List и return_dtype

Разбираем типизацию UDF map_batches в Polars v1.32+: как задать return_dtype для массивов int8, когда выбирать Array с фиксированным размером, а когда List.

Как корректно типизировать UDF map_batches, возвращающие массивы, в Polars v1.32+

Проблема

При применении пользовательской функции к датафрейму Polars через map_batches функция возвращает массив типа int8. Начиная с Polars v1.32, аргумент return_dtype стал обязательным для map_batches. Наивная попытка вроде return_dtype=pl.Array[pl.Int8] не срабатывает. До этого изменения схема уже содержала информацию о форме, например:

Schema([(...),
        ('signals', Array(Int8, shape=(9,)))])

Значения столбца signals формируются UDF, возвращающей массив numpy; эту функцию можно применять по группам.

Как воспроизвести проблему в коде

Ниже показан типичный пример. UDF возвращает numpy-массив с dtype int8, а map_batches указывают вернуть тип массива через синтаксис с квадратными скобками — и это падает:

import polars as pl
# возвращает numpy-массив с dtype int8
def build_vector(batch):
    ...
# некорректная типизация возвращаемого значения
result = frame.with_columns(
    pl.col("signals").map_batches(build_vector, return_dtype=pl.Array[pl.Int8])
)

Почему так происходит

Polars различает два вложенных типа: Array и List. Array — это типизированный массив фиксированной длины, в котором нужно указать и внутренний тип, и точный размер. List — это типизированный список переменной длины. Начиная с Polars v1.32, для map_batches требуется явный и полностью заданный return_dtype; для Array это означает, что нужно передать и тип элементов, и фиксированный размер. Запись pl.Array[pl.Int8] не задаёт фиксированную длину и потому недопустима. Подмена на List — не эквивалент Array; если там ожидается Array, а приходит List, Polars выбросит ошибку:

polars.exceptions.InvalidOperationError: expected Array type, got: list[i8]

Решение

Если UDF возвращает массив фиксированной длины, явно укажите размер через polars.datatypes.Array. Синтаксис: pl.Array(inner_dtype, size). Для массивов из трёх элементов int8 это будет так:

import polars as pl
# возвращает numpy-массив с dtype int8 фиксированной длины
def build_vector(batch):
    ...
out_type = pl.Array(pl.Int8, 3)
fixed = frame.with_columns(
    pl.col("signals").map_batches(build_vector, return_dtype=out_type)
)

Если фиксированная длина в данных равна девяти, укажите size=9, чтобы соответствовать схеме вида Array(Int8, shape=(9,)).

Когда длина заранее неизвестна, типизируйте результат как List из int8 вместо Array:

import polars as pl
# возвращает numpy-массив с dtype int8 переменной длины
def build_vector(batch):
    ...
list_type = pl.List(pl.Int8)
variable = frame.with_columns(
    pl.col("signals").map_batches(build_vector, return_dtype=list_type)
)

Помните, что List и Array — разные типы. Если ниже по конвейеру ожидается Array, а приходит List, Polars сообщит об этом так, как показано выше.

Зачем это важно

Начиная с v1.32, Polars требует явно указывать тип результата для map_batches, а вложенные типы вроде Array должны быть полностью определены. Это убирает неоднозначности в схемах и делает поведение UDF предсказуемым, особенно когда функции возвращают numpy-массивы и когда вычисления выполняются по группам.

Выводы

Если ваша UDF выдаёт массивы int8 фиксированной длины, задайте return_dtype как pl.Array(pl.Int8, size) с нужным размером. Если длина не фиксирована, используйте pl.List(pl.Int8) и убедитесь, что последующие выражения ожидают именно List, а не Array. После обновления проверьте итоговую схему, чтобы она соответствовала требованиям последующей логики.