Как исправить рост памяти в Datashader при dask‑xarray

Почему Datashader при шейдинге dask‑xarray вызывает compute() и раздувает память. Разбор причины, воспроизведение, и патч, устраняющий проблему в 0.16.1.

Отрисовка массивов огромного размера в Datashader на первый взгляд проста: разбить данные на чанки с помощью Dask, преобразовать их в растр на небольшом холсте и поручить остальное планировщику. На практике же легко получить всплески потребления памяти и затянутые прогоны. Типичная ситуация — dask‑поддерживаемый xarray порядка 150k × 90k с чанками 8192 × 8192 на Windows‑VM с 100 ГБ ОЗУ и 16 ядрами: уже запуск построения графика стремительно подводит RAM к пределу. Трассировка стека показывает вызов compute(), что и объясняет раздувание памяти. Ниже — что происходит и как это исправить.

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

Минимальный пример ниже создает крупный dask.array, оборачивает его в xarray и передает через Canvas.raster из Datashader, а затем в shade.

# импорт
import numpy as np
import dask.array as dk
import datashader as dz
from datashader import transfer_functions as tfun
import xarray as xa

# создаем большой массив на базе dask
grid = dk.random.random((100000, 100000), chunks=(1000, 1000))

# оборачиваем в xarray DataArray
xr_view = xa.DataArray(
    grid,
    dims=["u", "v"],
    coords={"u": np.arange(100000), "v": np.arange(100000)},
    name="sample_vals"
)

# пробуем отрисовать
tfun.shade(dz.Canvas(plot_height=300, plot_width=300).raster(xr_view))

Что на самом деле происходит и почему растет потребление памяти

Ключевая деталь видна в трассировке: путь шейдинга в Datashader вызывает внутреннюю функцию, которая при обнаружении массива на базе dask делает data = data.compute(). Эта строка принудительно материализует dask‑массив в памяти. При массивах порядка 10¹⁰ элементов это сразу объясняет устойчивый рост RAM и высокую загрузку CPU: на этапе шейдинга выполнение перестает идти по чанкам и сводится к вычислению всего массива разом.

Хотя dask поддерживает по‑чанковое исполнение, показанный в трассировке путь шейдинга обходит ленивые вычисления явным вызовом compute(). Наблюдаемое поведение напрямую следует из этого вызова.

Исправление: попробуйте патч Datashader

Доступно точечное исправление. Примените изменение из pull request по адресу https://github.com/holoviz/datashader/pull/1448 и проверьте, устраняет ли оно проблему в вашей среде. Есть сообщения, что при использовании последнего релиза Datashader после применения изменения может возникать другая ошибка, связанная с декоратором @ngjit, тогда как Datashader 0.16.1 вместе с этим патчем устраняет исходную проблему.

Использование того же кода после исправления

Ваш код не меняется; корректировка вносится в библиотеку. Можно продолжать рендерить с Datashader как и раньше:

# импорт
import numpy as np
import dask.array as dk
import datashader as dz
from datashader import transfer_functions as tfun
import xarray as xa

# создаем массив на базе dask
grid = dk.random.random((100000, 100000), chunks=(1000, 1000))

# подготавливаем для xarray/Datashader
xr_view = xa.DataArray(
    grid,
    dims=["u", "v"],
    coords={"u": np.arange(100000), "v": np.arange(100000)},
    name="sample_vals"
)

# отрисовываем
tfun.shade(dz.Canvas(plot_height=300, plot_width=300).raster(xr_view))

Если на последней версии вы сталкиваетесь с ошибкой, связанной с декоратором @ngjit после применения изменения, учтите, что такое поведение описано. Тот же патч, примененный к Datashader 0.16.1, подтвержденно устраняет исходную проблему с памятью.

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

Конвейеры визуализации на больших данных опираются на ленивые вычисления и разбиение на чанки, чтобы удерживать расход памяти в разумных пределах. Когда слой библиотеки вызывает compute() для объекта на базе dask, этот подход ломается и приводит к полной материализации. Понимание, когда это происходит, критично для прогнозирования ресурсоемкости и предотвращения сюрпризов на продакшн‑машинах или в ограниченных средах.

Выводы

Если при шейдинге dask‑поддерживаемого xarray Datashader резко наращивает потребление памяти, а в трассировке виден вызов compute(), приведите свою среду в соответствие с исправлением из упомянутого pull request. Если на последнем релизе Datashader после применения изменения появляется ошибка, связанная с @ngjit, попробуйте Datashader 0.16.1 с тем же патчем, где отмечено, что проблема решается, и следите за обсуждением в pull request для обновлений. Так вы сохраните преимущества данных на базе dask в своем процессе визуализации без неожиданных всплесков памяти.