Geostatistical Modeling Spatial Statistics with Python

Chunked Raster Processing with Dask-GeoPandas

TL;DR: Open the raster with rioxarray.open_rasterio(path, chunks={"x": 2048, "y": 2048}) for lazy array algebra, or read only the pixels you need with rasterio.windows.from_bounds(*bounds, transform=src.transform). Load polygons into a dask_geopandas.GeoDataFrame, spatial_shuffle() them into compact partitions, then map_partitions() a zonal-statistics function that reads the matching raster window per partition. Write results with a windowed rasterio writer or to partitioned GeoParquet so nothing larger than one tile ever lives in RAM.

Why This Matters

Nationwide land-cover mosaics, sub-metre DEMs, and multi-band satellite composites routinely run to tens of gigabytes, far beyond the working memory of a typical analysis node. A single src.read() on such a file allocates one contiguous array and crashes the interpreter before any statistic is computed. The fix is to never hold the whole grid at once: read bounded windows, back arrays with lazy Dask blocks, and stream results out tile by tile. This page turns the principles from the Memory-Efficient Processing guide into a concrete raster-plus-vector workflow, and it sits inside the wider Python Workflows for Spatial Modeling & Regression stack where these rasters usually become covariates for kriging or spatial regression.

The pairing that makes this tractable is rioxarray (Dask-backed raster arrays) with dask-geopandas (partitioned vector frames). Together they let you express a zonal-statistics or point-sampling job as a graph of independent per-tile tasks, each of which touches a bounded slice of pixels and a bounded slice of features. The scheduler runs them with a memory ceiling you control, and the result is written incrementally rather than assembled in one buffer.

Environment

bash
pip install \
  rioxarray==0.15.5 \
  rasterio==1.3.10 \
  xarray==2024.2.0 \
  dask==2024.3.1 \
  dask-geopandas==0.3.1 \
  geopandas==0.14.4 \
  pyogrio==0.7.2 \
  numpy==1.26.4 \
  shapely==2.0.4
python
import numpy as np
import geopandas as gpd
import dask_geopandas as dgpd
import rioxarray
import rasterio
from rasterio.windows import from_bounds, Window
from rasterio.features import geometry_mask
from dask.diagnostics import ProgressBar

gpd.options.io_engine = "pyogrio"

Step-by-Step Implementation

Step 1 — Open the Raster as Lazy Dask Chunks

rioxarray.open_rasterio with a chunks argument returns a DataArray whose values are a Dask array. No pixels are read until a reduction forces it. Aligning the chunk size to a multiple of the file’s internal block size keeps each read on whole blocks.

python
def inspect_blocking(path: str) -> tuple[int, int]:
    """Return the internal tile size of a GeoTIFF (block width, height)."""
    with rasterio.open(path) as src:
        bw, bh = src.block_shapes[0][1], src.block_shapes[0][0]
        print(f"CRS={src.crs}  size={src.width}x{src.height}  blocks={bw}x{bh}")
        return bw, bh

bw, bh = inspect_blocking("dem_national_1m.tif")

# Chunk in multiples of the internal block size (here 512 -> 2048)
dem = rioxarray.open_rasterio(
    "dem_national_1m.tif",
    chunks={"x": bw * 4, "y": bh * 4},
    masked=True,          # NoData becomes NaN
).squeeze("band", drop=True)

print(dem.data)           # dask.array — note chunksize in the repr, nothing loaded yet

Array algebra on dem stays lazy. A slope or hillshade expression builds a task graph; only .compute() or a windowed write materialises blocks, one chunk at a time.

Step 2 — Read Only the Pixels You Need with a Window

When the goal is values under specific features rather than whole-grid math, skip Dask array algebra entirely and read a tight window per feature bounding box. from_bounds maps geographic coordinates to a pixel window against the raster transform.

python
def read_window_for_bounds(src, bounds):
    """Read the pixel window intersecting `bounds`, clipped to the raster."""
    win = from_bounds(*bounds, transform=src.transform)
    col_off = max(int(np.floor(win.col_off)), 0)
    row_off = max(int(np.floor(win.row_off)), 0)
    col_end = min(int(np.ceil(win.col_off + win.width)), src.width)
    row_end = min(int(np.ceil(win.row_off + win.height)), src.height)
    if col_end <= col_off or row_end <= row_off:
        return None, None
    clipped = Window(col_off, row_off, col_end - col_off, row_end - row_off)
    data = src.read(1, window=clipped, masked=True)
    return data, src.window_transform(clipped)

The returned transform is essential: it is the affine mapping for the clipped sub-array, so any subsequent rasterisation of polygons lines up pixel-for-pixel.

Step 3 — Partition the Vector Layer

Load the zones into a dask_geopandas frame and spatial_shuffle them so each partition is a spatially compact tile rather than a random scatter of features. Compact partitions mean each partition’s bounding box covers a small raster window.

python
zones = gpd.read_parquet("watersheds.parquet")          # must share the raster CRS
assert zones.crs is not None

dz = dgpd.from_geopandas(zones, npartitions=64)
dz = dz.spatial_shuffle(by="hilbert", npartitions=64)   # Hilbert-curve locality
print(dz.npartitions, "partitions")

spatial_shuffle reorders features along a Hilbert curve, so features that are near in space land in the same partition. That locality is what bounds the raster window each task must read in the next step.

Step 4 — Compute Zonal Statistics Per Partition

map_partitions applies a plain-GeoPandas function to one partition at a time. Inside it, open the raster, read the single window covering the whole partition, then rasterise and reduce each polygon against that in-memory tile.

python
def zonal_stats_partition(part: gpd.GeoDataFrame, raster_path: str) -> gpd.GeoDataFrame:
    """Mean and count of raster pixels within each polygon of one partition."""
    out = part.copy()
    means = np.full(len(part), np.nan, dtype="float32")
    counts = np.zeros(len(part), dtype="int32")
    if len(part) == 0:
        out["zonal_mean"], out["pixel_count"] = means, counts
        return out

    with rasterio.open(raster_path) as src:
        tile, tile_transform = read_window_for_bounds(src, part.total_bounds)
        if tile is not None:
            for i, geom in enumerate(part.geometry.values):
                mask = geometry_mask(
                    [geom], out_shape=tile.shape,
                    transform=tile_transform, invert=True,
                )
                pixels = tile[mask & ~tile.mask]
                if pixels.size:
                    means[i] = pixels.mean()
                    counts[i] = pixels.size
    out["zonal_mean"] = means
    out["pixel_count"] = counts
    return out

meta = dz._meta.assign(zonal_mean="float32", pixel_count="int32")
result = dz.map_partitions(
    zonal_stats_partition, "dem_national_1m.tif", meta=meta,
)

with ProgressBar():
    stats = result.compute()      # scheduler runs partitions with bounded memory
print(stats[["zonal_mean", "pixel_count"]].describe())

Peak memory is one partition of polygons plus one raster window, regardless of how large the full mosaic is. Raising npartitions shrinks both.

Step 5 — Write Tiled Outputs

Never assemble a full output grid in memory. For a derived raster, write block by block through a windowed writer; for tabular zonal results, write partitioned GeoParquet.

python
# 5a. Tabular results -> partitioned GeoParquet (one file per partition)
result.to_parquet("zonal_stats.parquet")

# 5b. Derived raster -> stream blocks with a windowed writer
slope = dem.rio.reproject(dem.rio.crs)      # placeholder for any lazy expression
profile = {
    "driver": "GTiff", "height": dem.rio.height, "width": dem.rio.width,
    "count": 1, "dtype": "float32", "crs": dem.rio.crs,
    "transform": dem.rio.transform(), "tiled": True,
    "blockxsize": 512, "blockysize": 512, "compress": "deflate",
}
with rasterio.open("slope_tiled.tif", "w", **profile) as dst:
    for _, window in dst.block_windows(1):
        y0, y1 = window.row_off, window.row_off + window.height
        x0, x1 = window.col_off, window.col_off + window.width
        block = slope.isel(y=slice(y0, y1), x=slice(x0, x1)).values
        dst.write(block.astype("float32"), 1, window=window)

Interpreting the Output

stats["zonal_mean"] holds the mean pixel value per polygon and pixel_count the number of valid (non-NoData) pixels reduced. A pixel_count of zero flags a polygon that fell in a NoData gap or, more often, a CRS or extent mismatch — sanity-check those rows first. Compare stats["pixel_count"].sum() against the polygon areas divided by pixel area: a large shortfall means windows are being clipped by the raster edge, i.e. your zones extend beyond raster coverage.

For the derived raster, open slope_tiled.tif and confirm it reports internal tiling (gdalinfo shows Block=512x512). Tiled, compressed output is what makes the next consumer of this file able to window-read it cheaply, closing the loop.

Critical Best Practices

Match Chunk Size to Internal Blocks

A GeoTIFF stored in 512-pixel blocks read with a 700-pixel chunk forces every read to decompress overlapping blocks twice. Always chunk in integer multiples of src.block_shapes so each Dask block maps onto whole file blocks. If a raster is stripped rather than tiled, rewrite it once with tiled=True, blockxsize=512, blockysize=512 before heavy processing.

Keep the Vector CRS Equal to the Raster CRS

from_bounds and geometry_mask both assume feature coordinates are in the raster’s pixel-defining CRS. Reproject the vector layer to match the raster, never the raster to match the vector, which would trigger a full resample. Systematic CRS harmonisation is covered in the Reprojecting CRS for Accurate Distance Calculations workflow.

Size Partitions, Not Just Chunks

Two knobs control memory here: raster chunk size and vector partition count. A partition whose bounding box spans the whole study area defeats windowing because its window is the entire raster. spatial_shuffle plus a generous npartitions keeps each partition’s footprint small. Aim for partitions that read a window under ~256 MB.

Prefer map_partitions over Row-wise apply

Opening the raster once per partition amortises the file-open cost across thousands of features. A per-row apply that re-opens the dataset for every polygon is orders of magnitude slower and thrashes the GDAL block cache. Read the window once, reduce many polygons against it.

Watch the GDAL Block Cache

GDAL keeps a global block cache (default 5% of RAM). With many parallel workers each opening rasters, this can balloon. Cap it with the GDAL_CACHEMAX environment variable (e.g. 256 for 256 MB) so worker memory stays predictable under the Dask scheduler.

Troubleshooting

Symptom Likely cause Fix
MemoryError on open_rasterio chunks omitted, so the array is a single NumPy buffer Pass chunks={"x": ..., "y": ...} to force Dask backing
Zonal means all NaN Vector CRS differs from raster CRS Reproject the GeoDataFrame to src.crs before partitioning
pixel_count far below polygon area Window clipped at raster edge; zones exceed coverage Confirm raster covers the full AOI; drop or flag out-of-bounds zones
Reads far slower than expected Chunk size not aligned to internal blocks, or stripped GeoTIFF Rewrite raster tiled=True and chunk in block multiples
Worker memory climbs across tasks GDAL block cache accumulating Set GDAL_CACHEMAX=256 and lower worker count
spatial_shuffle errors on missing index Geometry column not set or empty partitions Ensure a valid active geometry and non-empty input before shuffling

Next Steps

With covariates summarised per zone you can feed them into a modelling stage; see the parent Memory-Efficient Processing guide for sparse-weight and lazy-evaluation patterns, and pair this with Using GeoParquet for Large Spatial Datasets to persist the tiled tabular results with predicate pushdown.

Related:

← Back to Memory-Efficient Processing

Related

Reducing Memory Bottlenecks in Geospatial Workflows Using GeoParquet for Large Spatial Datasets