Skip to content

Fixes #879: memory-guarded KDTree with tiled fallback for proximity#892

Merged
brendancol merged 1 commit intomasterfrom
fix/879-proximity-kdtree-oom
Feb 25, 2026
Merged

Fixes #879: memory-guarded KDTree with tiled fallback for proximity#892
brendancol merged 1 commit intomasterfrom
fix/879-proximity-kdtree-oom

Conversation

@brendancol
Copy link
Contributor

Summary

  • Replaces unbounded target-coordinate accumulation in _process_dask_kdtree() with a streaming count pass that caches coords within a 25% memory budget
  • Adds memory decision: if global tree estimate < 50% available RAM → fast lazy da.map_blocks path (functionally identical to before); otherwise → tiled expanding-ring KDTree fallback (eager, memory-safe)
  • Tiled path validates correctness via boundary-distance lower bound, expanding the search ring until provably correct or covering the full raster
  • Emits ResourceWarning when tiled fallback activates

Test plan

  • All 32 existing proximity tests pass (no regression)
  • 7 new tests covering the tiled fallback:
    • Dense raster matches numpy baseline (Euclidean + Manhattan)
    • Single target in corner (max ring expansion)
    • All-targets raster (all zeros)
    • No targets (all NaN)
    • ResourceWarning emitted on tiled fallback
    • Global path correctness after Phase 0 restructure

The dask KDTree path in proximity() accumulated all target coordinates
in memory before building a global cKDTree, causing OOM on dense rasters.

Replace the unbounded accumulation with a two-phase approach:
- Phase 0: stream chunks counting targets, cache coords within 25% budget
- Memory decision: if global tree estimate < 50% available RAM, build a
  single cKDTree (fast, lazy via da.map_blocks); otherwise fall back to
  a tiled expanding-ring KDTree (eager but memory-safe)

The tiled path builds local trees per output chunk, expanding the search
ring until a boundary-distance validation proves correctness. A
ResourceWarning is emitted when the tiled fallback activates.
@brendancol brendancol self-assigned this Feb 25, 2026
@brendancol brendancol merged commit 5b9c830 into master Feb 25, 2026
10 checks passed
@brendancol brendancol deleted the fix/879-proximity-kdtree-oom branch February 26, 2026 14:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant