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.

Chunk edges falling on block edges, and not Two grids each show 36 internal GeoTIFF blocks of 512 pixels. On the left, chunk boundaries drawn every 2048 pixels coincide with block boundaries, so every block belongs to exactly one chunk. On the right, chunk boundaries every 700 pixels fall inside blocks; only the four blocks in the intact rows and columns lie wholly inside one chunk, and the other 32 are decompressed twice. Chunk size against the GeoTIFF’s internal 512-pixel blocks chunks={'x': 2048, 'y': 2048} 2048 = 4 × 512 → chunk edges land on block edges one chunk 16 whole blocks 36 of 36 blocks read whole, once each chunks={…} in multiples of src.block_shapes chunks={'x': 700, 'y': 700} 700 = 1.37 × 512 → chunk edges cut through blocks only 4 of 36 blocks sit inside one chunk green = whole; the other 32 decompress twice

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.

What spatial_shuffle changes: the bounding box of one partition Two maps of the same raster extent hold the same 24 zone features as dots in identical positions. On the left, the six dots belonging to one partition are scattered across the extent and their dashed bounding box spans almost the whole raster, so the window read covers the entire mosaic. On the right, after a Hilbert spatial shuffle, a partition of six dots sits together and its bounding box is a small tile, so the window read stays under about 256 megabytes. One partition’s bounding box decides how much raster its task reads dgpd.from_geopandas(zones, npartitions=64) row order: one partition’s features scatter part.total_bounds ≈ the whole raster the task reads the entire mosaic – windowing defeated spatial_shuffle by='hilbert' coloured dots = one partition of 64 dz.spatial_shuffle(by='hilbert', npartitions=64) Hilbert order: that partition is one compact tile part.total_bounds = one small tile six features again, but one window under ~256 MB

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.

Peak memory per task as npartitions rises A bar chart with peak worker memory in megabytes on the vertical axis, from 0 to 1536, and three bars for 16, 64 and 256 vector partitions. Each bar stacks a fixed 256 megabyte GDAL block cache under a raster window that halves in area as partitions multiply: 1024 megabytes at 16 partitions, 256 at 64 and 64 at 256, giving totals of 1280, 512 and 320 megabytes. A dashed line at 512 megabytes marks the target the 64-partition bar meets, and a dashed line near the top marks a single full-mosaic read, tens of gigabytes above the chart. Two knobs, one budget: peak memory held by a single task a task holds one raster window plus one vector partition – never the mosaic 0 256 512 768 1024 1280 1536 peak memory per task (MB) 1280 MB 512 MB 320 MB target: ~256 MB window + 256 MB GDAL cache one src.read() of the whole mosaic: tens of GB, far above this chart npartitions = 16 npartitions = 64 npartitions = 256 window 1024 MB window 256 MB window 64 MB four times the partitions → a quarter of the area per bounding box raster window one per task GDAL block cache GDAL_CACHEMAX=256 + one vector partition (a few MB, shrinks too) window bytes scale as 1 / npartitions

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

Parallel Spatial Workflows with Dask and Joblib Reducing Memory Bottlenecks in Geospatial Workflows Using GeoParquet for Large Spatial Datasets