Multi-threaded compression in rasterio?

Question

For writing large raster files (e.g. satellite image GeoTIFFs) to disk, the throughput rate tends to be limited by the compression algorithm rather than by the raw disk IO. (Benchmarks: 1, 2.) This is good news because data compression (computation) is far easier to parallelise than file IO generally is. Can we take advantage of this in python?

When using the rasterio python library, is there any way to use multithreading to write a large file more quickly? (For example, is there a way to set GDAL_NUM_THREADS=ALL_CPUS?)

Answer 1

Yes, although this seems not to be mentioned in the documentation.

If you include num_threads=8 or num_threads='all_cpus' as an argument to rasterio.open then multithreading will be enabled (for writing of compressed data).

Example:

import numpy as np, rasterio
size = 16384
chunk = 512
with rasterio.open('test.tiff', 'w', driver='GTiff', nodata=0,
                   width=size, height=size, count=1, dtype=np.float32,
                   tiled=True, blockxsize=256, blockysize=256, 
                   compress='lzw', num_threads='all_cpus') as dst: 
    for i in range(0, size, chunk):
        for j in range(0, size, chunk):
            data = np.random.random((chunk, chunk)).astype(np.float32)
            dst.write(data, window=rasterio.windows.Window(i, j, chunk, chunk), indexes=1)

Alternatively, you can temporarily set the global configuration:

with rasterio.Env(GDAL_NUM_THREADS='ALL_CPUs'):
    ...