Add dask backends for rasterize()#997
Merged
Merged
Conversation
Tile-based parallel rasterization via dask.delayed. New `chunks` parameter splits the output raster into tiles, rasterizes each independently, and assembles with da.block. Polygons are rasterized per-tile (scanline fill is tile-exact). Lines and points are pre-rasterized at full-raster scale then distributed to tiles to avoid Bresenham boundary artifacts. Also adds dask-geopandas input support (.compute() eagerly) and cuspatial optional import stub.
Removes three functions (_bresenham_collect, _precompute_line_pixels, _pixels_for_tile) that materialized all Bresenham pixels in the scheduler process before any dask task ran. Line segments and points are now kept in their compact representation (5 arrays for segments, 3 for points) and filtered per-tile via bbox overlap. Bresenham runs inside each tile worker using the existing @ngjit _burn_lines_cpu (numpy) or GPU burn_lines kernel (cupy). Segments are offset to tile-local coords; since Bresenham is translation- invariant, the pixel path is exact. Fixes: pure-Python _bresenham_collect (now removed), O(T*P) per-tile pixel scanning (now O(T*S) segment filtering where S << P), and O(total_line_pixels) upfront memory allocation (now O(segments)).
Shapely geometry pickle goes through __reduce__ with WKB encoding per object. When a polygon intersects many tiles, that serialization cost multiplies. Serialize to WKB once up front (~20x cheaper to pickle as raw bytes), pass the bytes into each delayed task, and deserialize inside the worker via shapely.from_wkb. Applied to both _run_dask_numpy and _run_dask_cupy orchestrators.
Expand TestDaskCupy from 8 to 30 tests, matching TestDaskNumpy coverage: line/point parity across chunk sizes, chunk boundary tests for lines and points, single-chunk, output type/shape/coords, all 6 merge modes, holes, all_touched, dtype, int chunks shorthand, and GeoDataFrame input.
The merge parameter now accepts a callable in addition to the built-in
string modes. For CPU backends, pass a @ngjit function. For GPU backends,
pass a @cuda.jit(device=True) function. Signature:
merge_fn(old_val, new_val, is_first) -> float64
Internally, the integer-dispatch _merge_pixel is replaced by
_apply_merge which calls the merge function directly. Built-in modes
are now @ngjit functions (_merge_last, _merge_first, etc.) looked up
from _MERGE_FUNCTIONS. GPU kernels are compiled per merge function
via a factory/cache pattern in _ensure_gpu_kernels.
Adds 11 new tests (TestCustomMerge, TestCustomMergeDask,
TestCustomMergeGPU) and a rasterize user guide notebook stub.
The merge function signature is now (pixel, props, is_first) -> float64
where props is a 1D float64 array carrying one or more feature
properties. Built-in modes (last, first, max, min, sum, count) use
props[0]. Custom functions can access any column.
New 'columns' parameter on rasterize() extracts multiple GeoDataFrame
columns into the props array. For (geometry, value) pairs, props is
a length-1 array containing the value.
Internally, the pipeline now uses geometry-index + props-table instead
of carrying scalar values. Polygon edges use edge_geom_id (already
existed) to index into a poly_props table. Line segments and points
carry geom_idx arrays that index into line_props / point_props tables.
The edge_value array is removed entirely.
Example:
@ngjit
def weighted_density(pixel, props, is_first):
density = props[0] / props[1]
if is_first:
return density
return pixel + density
rasterize(gdf, columns=['population', 'area'],
merge=weighted_density, ...)
…i-column, and dask examples
…ations - CSR construction: difference-array counting O(n_edges + height) vs O(n_edges * avg_height) - Shell sort replaces insertion sort in scanline fill (CPU + GPU): O(n^4/3) vs O(n^2) - Pre-compute segment bboxes once before dask tile loop - Pre-transfer shared GPU props tables once in dask+cupy path - Boundary segment extraction: direct coord extraction, no LineString wrappers - Edge loop fallback: pre-allocated arrays instead of per-edge list appends - Vectorized edge extraction: early del of intermediates to cut peak memory - Shared _geometry_bboxes() call in list-of-pairs input parsing
… dedup _parse_input - Remove per-element np.int32() wrapping in _extract_points_loop and _extract_lines_loop - Use object array + boolean indexing for WKB tile filtering (replaces nonzero + list comprehension) - Hoist geometry.tolist() and total_bounds above column/columns branch in _parse_input - Add early return in _build_row_csr_numba for empty edge input - Use np.arange instead of list(range()) for poly_ids in tile workers
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
chunksparameter onrasterize()for dask-backed parallel output. Splits the output raster into tiles, rasterizes each tile viadask.delayed, assembles withda.block.cuspatialoptional import stub for future GPU geometry parsing.daskandgpuextras insetup.cfg.Test plan
TestDaskNumpytests: polygon/line/point/mixed parity vs numpy across chunk sizes(5,5),(3,7),(100,100). Chunk boundary, empty tile, all 6 merge modes, holes,all_touched, dtype, coords, GeoDataFrame, dask-geopandas.TestDaskCupytests: same structure with GPU backend.dask-geopandas).