Parallel Spatial Workflows with Dask and Joblib

TL;DR: Use joblib.Parallel(n_jobs=8, backend="loky") for independent units — tiles, folds, regions — and dask_geopandas.spatial_shuffle(by="hilbert") before dask_geopandas.sjoin when the data will not fit in memory. Send workers coordinate arrays or GeoParquet paths, never live geometry, and expand every tile by the variogram range before selecting conditioning data.

Why This Matters #

Spatial work parallelises well because most of it is embarrassingly parallel and nobody notices. A kriging surface over a large area is not one linear system; it is hundreds of small ones, because kriging weights decay to nothing beyond the variogram range and a global solve is impossible anyway. Twelve thousand samples would need a 12,000 × 12,000 covariance matrix — 1.15 GB before the right-hand side — and the right-hand side for a 640,000-cell output grid is measured in terabytes. Tiling is not an optimisation here. It is what makes the problem exist at all, and once you have tiles you have independent units, and once you have independent units you have parallelism for free.

The trap is that the cheapest parallel tools were designed for arrays, and spatial data is not arrays. A GeoDataFrame carries GEOS objects that must be reconstructed in every worker, a spatial index that is deliberately dropped on pickling, and a CRS that GDAL and PROJ do not want shared across threads. Get those three details wrong and an eight-core run is slower than the loop it replaced. This page sits inside Memory-Efficient Processing, and the two concerns are the same concern: parallelism multiplies your peak memory by the worker count, so the techniques in reducing memory bottlenecks in geospatial workflows are a prerequisite rather than an alternative.

Environment and Version Pinning #

pip install \
  "geopandas==1.1.1" "shapely==2.1.1" "pyarrow==19.0.1" \
  "numpy==2.2.6" "scipy==1.15.2" \
  "joblib==1.5.1" "threadpoolctl==3.6.0" \
  "dask==2025.5.1" "distributed==2025.5.1" "dask-geopandas==0.5.0" \
  "pykrige==1.7.2" "gstools==1.7.0"

import time
import numpy as np
import geopandas as gpd
import gstools as gs
from joblib import Parallel, delayed, parallel_config
from pykrige.ok import OrdinaryKriging

Step-by-Step Implementation #

1. Build a sample field with a known variogram #

The halo width later depends on the variogram range, so the range has to be a measured quantity rather than a guess. Simulating the field with gstools means the model that generated it and the model fitted to it can be compared directly.

rng = np.random.default_rng(2026)
SIDE = 20_000.0          # 20 km square, EPSG:32630, metres
N_SAMPLES = 12_000

xs = rng.uniform(0.0, SIDE, N_SAMPLES)
ys = rng.uniform(0.0, SIDE, N_SAMPLES)

# Exponential covariance, practical range = 3 * len_scale = 1200 m.
# gstools does not simulate the nugget, so add it as measurement noise.
model = gs.Exponential(dim=2, var=1270.0, len_scale=400.0)
srf = gs.SRF(model, mean=420.0, seed=20260807)
pb = srf((xs, ys)) + rng.normal(0.0, np.sqrt(180.0), N_SAMPLES)

samples = gpd.GeoDataFrame(
    {"pb": pb}, geometry=gpd.points_from_xy(xs, ys), crs="EPSG:32630"
)
samples.to_parquet("data/pb_samples.parquet")

bin_centre, gamma = gs.vario_estimate(
    (xs, ys), pb, bin_edges=np.arange(0.0, 3000.0, 100.0),
    sampling_size=4000, sampling_seed=7,
)
fitted = gs.Exponential(dim=2)
fitted.fit_variogram(bin_centre, gamma, nugget=True)
print(fitted)
print(f"practical range = {3 * fitted.len_scale:.1f} m")

Exponential(dim=2, var=1263.4, len_scale=397.2, nugget=182.6)
practical range = 1191.6 m

2. Define the tiling and the worker function #

The worker takes plain float arrays and a bounds tuple. Nothing in its signature is a geometry, a GeoDataFrame or a CRS object, which is the single decision that makes the rest of the run cheap.

TILE = 2500.0            # tile side, metres
HALO = 1191.6            # >= the practical range
CELL = 25.0              # output raster resolution
VARIO = {
    "variogram_model": "exponential",
    # PyKrige's exponential takes [partial sill, practical range, nugget]
    "variogram_parameters": [1263.4, 1191.6, 182.6],
}

tiles = [
    (i * TILE, j * TILE, (i + 1) * TILE, (j + 1) * TILE)
    for j in range(int(SIDE // TILE))
    for i in range(int(SIDE // TILE))
]
print(f"{len(tiles)} tiles, output grid "
      f"{int(SIDE // CELL)}x{int(SIDE // CELL)} cells")


def krige_tile(bounds, x, y, z, halo=HALO, cell=CELL, vario=VARIO):
    """Krige one tile, conditioning on every sample within `halo` of it."""
    x0, y0, x1, y1 = bounds
    sel = ((x >= x0 - halo) & (x <= x1 + halo) &
           (y >= y0 - halo) & (y <= y1 + halo))
    n = int(sel.sum())
    if n < 12:
        raise ValueError(f"tile {bounds} has only {n} conditioning points")

    gx = np.arange(x0 + cell / 2, x1, cell)
    gy = np.arange(y0 + cell / 2, y1, cell)

    ok = OrdinaryKriging(
        x[sel], y[sel], z[sel], exact_values=False, **vario
    )
    zhat, var = ok.execute("grid", gx, gy)
    return bounds, np.asarray(zhat, dtype="float32"), n

64 tiles, output grid 800x800 cells

3. Run the tiles with joblib and mosaic the result #

x_arr = samples.geometry.x.to_numpy()
y_arr = samples.geometry.y.to_numpy()
z_arr = samples["pb"].to_numpy()

with parallel_config(backend="loky", inner_max_num_threads=1):
    for n_jobs in (1, 2, 4, 8):
        t0 = time.perf_counter()
        results = Parallel(n_jobs=n_jobs)(
            delayed(krige_tile)(b, x_arr, y_arr, z_arr) for b in tiles
        )
        print(f"n_jobs={n_jobs}: {time.perf_counter() - t0:6.1f} s")

counts = np.array([n for _, _, n in results])
print(f"conditioning points per tile: min {counts.min()}, "
      f"median {int(np.median(counts))}, max {counts.max()}")

nrow = ncol = int(SIDE // CELL)
mosaic = np.full((nrow, ncol), np.nan, dtype="float32")
for (x0, y0, _, _), zhat, _ in results:
    r0, c0 = int(y0 // CELL), int(x0 // CELL)
    mosaic[r0:r0 + zhat.shape[0], c0:c0 + zhat.shape[1]] = zhat
print(f"mosaic {mosaic.shape}, unfilled cells: {int(np.isnan(mosaic).sum())}")

n_jobs=1:   48.6 s
n_jobs=2:   25.9 s
n_jobs=4:   13.7 s
n_jobs=8:    7.4 s
conditioning points per tile: min 401, median 707, max 781
mosaic (800, 800), unfilled cells: 0

4. Measure the seams, with and without the halo #

A tile-edge artefact is easy to miss on a colour ramp and trivial to measure: compare each pair of cells straddling a tile boundary against the ordinary cell-to-cell step inside a tile.

seam_cols = np.arange(1, int(SIDE // TILE)) * int(TILE // CELL)
seam_jump = np.abs(mosaic[:, seam_cols] - mosaic[:, seam_cols - 1])
interior = np.abs(np.diff(mosaic[:, 10:90], axis=1))

print(f"max seam jump    : {seam_jump.max():7.2f} mg/kg")
print(f"median seam jump : {np.median(seam_jump):7.2f} mg/kg")
print(f"median interior  : {np.median(interior):7.2f} mg/kg")

With HALO = 1191.6:

max seam jump    :    0.42 mg/kg
median seam jump :    0.05 mg/kg
median interior  :    0.06 mg/kg

Re-running the whole thing with HALO = 0.0 — each tile kriged only from the samples it contains — gives:

max seam jump    :   41.68 mg/kg
median seam jump :    6.31 mg/kg
median interior  :    0.06 mg/kg

The halo-free mosaic finishes in 2.1 seconds on the same eight workers instead of 7.4, and it is not a surface. A median seam jump of 6.31 mg/kg against a median interior step of 0.06 mg/kg is a grid of visible lines a hundred times the natural gradient, and the worst cell is off by 41.68 mg/kg — comparable to the field’s own standard deviation of 38.1.

5. Partition and join out of core with dask-geopandas #

Kriging tiles are independent, so joblib was enough. A join is not: every left partition may match rows in any right partition, so the framework has to reason about which pairs can possibly intersect. That is what dask-geopandas adds over plain Dask, and it only works once both frames carry spatial_partitions.

import dask_geopandas as dgpd
from dask.distributed import Client

client = Client(n_workers=8, threads_per_worker=1,
                processes=True, memory_limit="4GB")

buildings = dgpd.read_parquet("data/buildings.parquet")   # 4.2 M polygons
areas = dgpd.read_parquet("data/output_areas.parquet")    # 34 211 polygons
print(buildings.npartitions, areas.npartitions, buildings.spatial_partitions)

buildings = buildings.spatial_shuffle(by="hilbert", npartitions=64)
areas = areas.spatial_shuffle(by="hilbert", npartitions=16)
print(type(buildings.spatial_partitions).__name__,
      len(buildings.spatial_partitions))

64 16 None
GeoSeries 64

With bounding geometry attached to every partition, the candidate pairs collapse. You can count them yourself before paying for the join:

left_boxes = gpd.GeoDataFrame(geometry=buildings.spatial_partitions)
right_boxes = gpd.GeoDataFrame(geometry=areas.spatial_partitions)
pairs = gpd.sjoin(left_boxes, right_boxes, predicate="intersects")
print(f"partition pairs to evaluate: {len(pairs)} of "
      f"{buildings.npartitions * areas.npartitions}")

joined = dgpd.sjoin(buildings, areas, how="inner", predicate="intersects")
per_area = joined.groupby("oa_code").size().compute()
print(per_area.describe()[["count", "mean", "min", "50%", "max"]])

partition pairs to evaluate: 118 of 1024
count    34211.000000
mean       122.706000
min          1.000000
50%        104.000000
max       1943.000000

Read the two numbers together: 118 partition pairs instead of 1,024 is 8.7 times less work scheduled, before any geometry is touched. Skipping the shuffle does not fail — it silently runs the full cross product, because row-order partitions each span the whole study area and every bounding box therefore intersects every other. The same reasoning governs the single-machine case in optimizing GeoPandas spatial joins for large datasets; Dask just applies it one level up, to partitions rather than rows.

Interpreting the Output #

The joblib ladder is the honest measure of whether parallelism helped. Going from one worker to eight took 48.6 seconds down to 7.4, a speedup of 6.6 and a parallel efficiency of 82%. On independent, evenly sized units that is roughly what good looks like: the missing 18% is process start-up, the per-task pickling of three 96 kB arrays, and the tail effect of 64 tasks not dividing evenly into eight workers. An efficiency above 90% usually means the tasks are long; below 50% means something is serial, and on spatial work that something is almost always either payload serialisation or a thread pool fighting itself.

The seam metrics are the correctness check, and they are more important than the timings. Compare the median seam jump against the median interior step, not against zero: 0.05 versus 0.06 mg/kg says the seam is indistinguishable from ordinary cell-to-cell variation, which is the definition of a continuous mosaic. Warning sign: a maximum seam jump more than about three times the median interior step, which means at least one tile had too little conditioning data on one side.

For the Dask side, the number to watch is the partition-pair count. If it is close to left.npartitions * right.npartitions, spatial partitioning is not doing anything and you have a distributed cross product. Values in the range of two to four times max(npartitions) indicate compact, well-separated partitions. The Dask dashboard’s task stream is the second check: long uniform bars are healthy, a red band of communication means partitions are being shuffled repeatedly.

Critical Best Practices #

Never let a live GeoDataFrame cross the process boundary #

Anything the worker function closes over is pickled once per task, in the parent, one task at a time. Shapely 2 serialises a whole geometry column through a vectorised WKB round trip, which is much better than Shapely 1.8’s per-object pickling, but the receiving worker still has to reconstruct every GEOS object. On a 250,000-polygon building layer that is 0.41 seconds to write and 0.96 seconds to parse, against 0.74 seconds of actual join work. Pass a GeoParquet path, a WKB byte array via gdf.geometry.to_wkb(), or bare coordinate arrays as the kriging example does — and note that joblib memory-maps NumPy arrays above max_nbytes (1 MB by default), so large float arrays become free while geometry never does. Storing the partitions as GeoParquet in the first place, as covered in using GeoParquet for large spatial datasets, is what makes the path-passing pattern possible.

Size the halo from the variogram, not from intuition #

The correct halo is the distance beyond which kriging weights are negligible, which for a fitted model is the practical range — three times len_scale for exponential, the range parameter itself for spherical. Anything shorter leaves prediction cells near the tile edge with no data on one side, and ordinary kriging fills the gap by re-estimating a local mean from a lopsided sample. Anything much longer is wasted: the conditioning set grows with the square of the halo width, and PyKrige’s cost grows with the square of the conditioning count again.

Prefer processes to threads for anything touching GDAL or PROJ #

GDAL dataset handles are not safe to share across threads, PROJ transformation objects must not be reused across threads, and Shapely holds the GIL for part of every predicate call. The loky backend spawns fresh interpreters, which is slower to start and immune to all three problems; backend="threading" is only appropriate for kernels that spend their time inside NumPy or BLAS. In dask.distributed, that means processes=True, threads_per_worker=1 for any workflow that opens rasters or reprojects, which is also the configuration assumed in chunked raster processing with Dask-GeoPandas.

Stop joblib and BLAS from oversubscribing the CPU #

PyKrige solves dense systems through NumPy, and NumPy’s BLAS defaults to one thread per core. Eight joblib workers each spawning eight BLAS threads is 64 threads on eight cores, and the resulting cache thrashing routinely halves throughput. parallel_config(backend="loky", inner_max_num_threads=1) sets the relevant environment variables inside each worker; threadpoolctl.threadpool_limits(limits=1, user_api="blas") does the same job for code you cannot wrap. Symptomatically this looks like a speedup that stalls around 2–3× no matter how many workers you add.

Accept that the spatial index is rebuilt in every worker #

GeoPandas deliberately excludes _sindex from the pickled state, and there is no way to send a built index to a worker — the underlying GEOS STRtree is not serialisable. Every worker that calls sjoin, clip or any predicate pays to build it. The consequence is that fine-grained tasks are wasteful: a task that spends 1.22 seconds indexing to do 0.74 seconds of work has an index-to-work ratio above one. Make tasks larger, hand each worker a whole partition rather than a row group, and set batch_size in joblib so short tasks are dispatched in groups.

Troubleshooting #

Next Steps #

Parallel workers multiply peak memory by the worker count, so pair this with reducing memory bottlenecks in geospatial workflows before raising n_jobs, and store the partitions the workers read as described in using GeoParquet for large spatial datasets.

Frequently Asked Questions #

When should I reach for Dask instead of joblib?

Reach for Dask when the working set does not fit in one process, or when the units of work depend on each other. A join, an overlay or a dissolve needs data from more than one partition, so it needs a task graph and a shuffle; that is Dask territory. Kriging a thousand tiles, fitting one variogram per region, or running cross-validation folds are independent units, and joblib does those with a single Parallel call and no scheduler, no cluster and no serialisation of an execution graph.

Why is passing a GeoDataFrame to joblib workers so expensive?

Anything captured by the function you pass to Parallel is pickled once per task, in the parent process, one task at a time. Shapely 2 serialises a geometry column through a vectorised WKB round trip, which is far better than Shapely 1.8 but still means the worker must rebuild every GEOS object on arrival. On 250,000 building polygons that is roughly 0.41 seconds to write and 0.96 seconds to parse, against 0.74 seconds of actual join work. Send coordinate arrays or GeoParquet paths instead.

How wide should the halo around each tile be?

At least the practical range of the variogram, because beyond that distance the kriging weights are effectively zero and adding more conditioning data changes nothing. For the exponential model fitted here the practical range is three times the length scale, 1,192 metres, so a 1,192 metre halo drops the worst seam discontinuity from 41.68 to 0.42 mg/kg. A wider halo is safe but not free: the conditioning set grows with the square of the halo width.

Can I use threads instead of processes for spatial work?

Rarely. GDAL objects are not safe to share between threads, PROJ transformation objects must not be reused across threads, and Shapely predicate calls hold the GIL for part of their work. Threads do help for pure NumPy or BLAS-heavy kernels that release the GIL, but any task that opens a raster, reprojects, or touches a spatial index should run in its own process. Use the loky backend in joblib, or processes=True with threads_per_worker=1 in dask.distributed.

Related

• Chunked Raster Processing with Dask-GeoPandas — the raster half of the same problem, with window alignment instead of halos
• Reducing Memory Bottlenecks in Geospatial Workflows — what to fix before multiplying peak memory by eight
• Using GeoParquet for Large Spatial Datasets — the storage format that makes path-passing cheap
• Optimizing GeoPandas Spatial Joins for Large Datasets — the single-process join this parallelises

← Back to Memory-Efficient Processing