Рваные данные в NumPy: object-массивы, np.vectorize, frompyfunc и быстрый цикл

Показываем, почему np.vectorize ломается с object-массивами NumPy на рваных данных, как исправить через otypes и np.frompyfunc, где быстрее простой цикл. Примеры.

Работать с «рваными» (jagged) данными в NumPy порой заманчиво, особенно когда хочется получить семантику массивов при списках разной длины. Распространённый приём — хранить списки Python в массиве с dtype=object и пытаться применять операции «векторизованным» способом. Это нередко ломается самым неожиданным образом. Ниже — минимальное, воспроизводимое руководство: что идёт не так, почему это происходит и что действительно работает.

Воспроизведение проблемы

Предположим, у нас есть объектный массив с 500 списками Python разной длины, и мы хотим добавить значение 100 к подмножеству строк разом.

import numpy as np
arr_jag = np.array([[t for t in range(np.random.randint(10))] for _ in range(500)], dtype=object)
targets = [0, 10, 20, 30, 40, 50]
v_add = np.vectorize(lambda seq: seq + [100])
# Raises: ValueError: setting an array element with a sequence
arr_jag[targets] = v_add(arr_jag[targets])

Код пытается «распространить» операцию уровня Python на выборку списков. В итоге он падает с ValueError.

Почему это не срабатывает

Проблема состоит из двух частей. Во‑первых, массивы NumPy с dtype=object хранят произвольные объекты Python; для списков нет поэлементной векторизованной арифметики. Во‑вторых, np.vectorize не привносит низкоуровневой векторизации; по сути это тонкая обёртка над циклом, которая многократно вызывает вашу Python‑функцию. Когда тип результата явно не задан для объектных данных, построение результирующего массива может не совпасть с вашей целью, что приводит к ошибке присваивания. Иными словами, если не указать np.vectorize возвращать элементы типа object, ваше присваивание в объектный массив конфликтует с попыткой NumPy сформировать обычный ndarray из последовательностей разной длины.

Рабочий вариант с np.vectorize

Можно сделать векторизованный вызов присваиваемым, принудив объектный тип возвращаемого значения. Тогда правая часть согласуется с левой частью выбора.

import numpy as np
arr_jag = np.array([[t for t in range(np.random.randint(10))] for _ in range(500)], dtype=object)
targets = [0, 10, 20, 30, 40, 50]
add_tail = lambda seq: seq + [100]
vec_obj = np.vectorize(add_tail, otypes=[object])
arr_jag[targets] = vec_obj(arr_jag[targets])

Это сохраняет задуманное поведение: каждому выбранному списку добавляется 100, а присваивание проходит успешно, поскольку результат — объектный массив.

Простой цикл быстрее

Несмотря на название, np.vectorize остаётся циклом на уровне Python с накладными расходами. Прямой проход по индексам часто быстрее, когда вы модифицируете объекты Python внутри объектного массива.

import numpy as np
arr_jag = np.array([[t for t in range(np.random.randint(10))] for _ in range(500)], dtype=object)
targets = [0, 10, 20, 30, 40, 50]
add_tail = lambda seq: seq + [100]
for pos in targets:
    arr_jag[pos] = add_tail(arr_jag[pos])

Эмпирические сравнения показывают, что явный цикл заметно быстрее np.vectorize для этого сценария. Использование np.frompyfunc для построения Python‑ufunc может быть быстрее, чем np.vectorize, хотя и здесь итерация выигрывает на той же задаче. Следующие замеры иллюстрируют относительную производительность на меньшем примере и выборке:

np.vectorize(..., otypes=[object]): около 19,4 μs на вызов; явный цикл for: около 2 μs на итерацию; np.frompyfunc: около 9,34 μs на вызов. В этом сценарии итерация быстрее.

Альтернативный API с np.frompyfunc

Если вам ближе сигнатура вызова в стиле ufunc, np.frompyfunc её предоставляет без изменения модели затрат. Функция возвращает объектный массив и в целом легче, чем np.vectorize.

import numpy as np
arr_jag = np.array([[t for t in range(np.random.randint(10))] for _ in range(500)], dtype=object)
targets = [0, 10, 20, 30, 40, 50]
add_tail = lambda seq: seq + [100]
apply_py = np.frompyfunc(add_tail, 1, 1)
arr_jag[targets] = apply_py(arr_jag[targets])

Это создаёт объектный массив, подходящий для присваивания обратно в arr_jag по выбранным индексам.

Почему это важно

NumPy особенно хорош с однородными числовыми массивами фиксированного размера. Массивы списков хранятся как объекты Python, и операции над ними выполняются в пространстве Python. Это означает, что обёртки вроде vectorize не дают ускорений на уровне C. Для «рваных» массивов и нерегулярных данных стоит использовать структуры, рассчитанные на такую форму. Например, представление таких данных в уплощённом виде плюс пары старт–конец может быть эффективно, а специализированные библиотеки лучше подходят для нерегулярных раскладок. Есть и разрежённые представления, перекликающиеся с идеей «массива объектов». В частности, формат scipy.sparse LIL внутренне использует два массива с dtype=object и удобен для итеративной сборки, хотя он не самый компактный и не лучший для вычислений; преобразования между разрежёнными форматами доступны без труда.

Итоги

Если вы храните списки Python внутри объектного массива NumPy и хотите добавлять элементы или преобразовывать выбранные записи, убедитесь, что любая «векторизованная» обёртка возвращает элементы типа object, либо просто используйте понятный цикл. Когда важна производительность, прямой обход нередко выигрывает в этом паттерне, а построение Python‑ufunc через np.frompyfunc может быть прагматичной альтернативой с массивным API. Для масштабных задач с нерегулярными данными пересмотрите контейнер: кодируйте «рваные» структуры в плоском виде со смещениями, используйте библиотеку, специально созданную для таких массивов, или прибегайте к разрежённому представлению, если оно соответствует вашей модели вычислений.