Rechunker: The missing link for chunked array analytics
By: Ryan Abernathey
Tags: Python, Data Science, Distributed Systems, Geospatial, Big Data

by Ryan Abernathey and Tom Augspurger

TLDR: this post describes a new python library called rechunker, which performs efficient on-disk rechunking of chunked array storage formats. Rechunker allows you to write code like this:

from rechunker import rechunk

target_chunks = (100, 10, 1)

max_mem = "2GB"

plan = rechunk(source_array, target_chunks, max_mem,

"target_store.zarr",

"temp_store.zarr")

plan.execute()

…and have the operation parallelized over any number of Dask workers.

Motivation

Chunked arrays are a key part of the modern scientific software stack in fields such as geospatial analytics and bioinformatics. Chunked arrays take a large multidimensional array dataset, such as an image captured over many timesteps, and split it up into many “chunks” — smaller arrays which can comfortably fit in memory. These chunks can form the basis of parallel algorithms that can make data science workflows go a lot faster.

Example of a chunked array, as represented by Dask.

Chunked arrays are implemented in both parallel computing frameworks — such as Dask and NumpyWren — and as an on-disk storage format. Some storage formats that support chunked arrays include HDF5, TileDB, Zarr, and Cloud Optimized Geotiff. When these chunked array storage formats are paired with the above computing frameworks, excellent scaling performance can be achieved.

However, chunked arrays workflows can fail hard when the chunks are not aligned with the desired analysis method. A great example can be found in this post from a user on the Pangeo forum:

Geospatial satellite data is often produced as a global map once per day, creating a natural chunk structure (e.g. one file per day). But what happens if you want to do a timeseries analysis at each point in space? This analysis can’t be parallelized over chunks. Many array-based workflows get suck on similar problems.

One existing solution is to use Dask’s rechunk function to create a new chunk structure lazily, on the fly, in memory. This works great for some problems. For others, particularly those involving a full rechunk (every source chunk goes into every target chunk), Dask’s algorithm can run out of memory, or produce an unmanageably large number of tasks. (More details can be found in the post linked above.)

To address this problem, we created a new package that aims to solve this specific problem in an optimal way: rechunker.

The Rechunker Algorithm

Rechunker takes an input chunked array (or group of arrays) stored in a persistent storage device (such as a filesystem or a cloud storage bucket) and writes out an array (or group of arrays) with the same data, but different chunking scheme, to a new location. Along the way, it may create a temporary, intermediate copy of the array in persistent storage. The reliance on persistent storage is a key difference between Rechunker and Dask’s rechunk function.

Figuring out the most efficient way to do this was a fun computer science problem to solve. Via our Discourse forum, many people contributed to the discussion and shared different ideas they had implemented in the past. We identified a couple of key requirements for Rechunker’s algorithm:

Respect memory limits. Rechunker’s algorithm guarantees that worker processes will not exceed a user-specified memory threshold.

Rechunker’s algorithm guarantees that worker processes will not exceed a user-specified memory threshold. Minimize the number of required tasks. Specifically, for N source chunks and M target chunks, the number of tasks is always less than N + M.

Specifically, for N source chunks and M target chunks, the number of tasks is always less than N + M. Be embarrassingly parallel. The task graph should be as simple as possible, to make it easy to execute using different task scheduling frameworks. This also means avoiding write locks, which are complex to manage, and inter-worker communication.

These considerations led to the creation of an algorithm we call Push-Pull-Consolidated.