Fuse hypsometric_integral dask path to a single graph evaluation#1212
Merged
Conversation
_hi_dask_numpy did two blocking dask.compute() calls (_unique_finite_zones at one step, _hi_reduce at the next), so the caller paid for two full input scans before the lazy map_blocks output was even returned. _hi_reduce also np.stacked the per-block partials into an (n_blocks, n_zones, 4) array on the scheduler; at 240k blocks * 1000 zones that is ~7.7 GB resident in a single scheduler task. Have each block discover its own local unique zones and return a dict mapping zone id -> (min, max, sum, count). _hi_reduce stream-merges the partial dicts into a global hi_lookup so scheduler peak memory scales with the number of distinct zones, not n_blocks * n_zones. The up-front _unique_finite_zones pass is gone and the whole dask path collapses to a single graph evaluation.
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
_hi_block_statsnow discovers its own local unique zones and returns a dict ofzone_id -> (min, max, sum, count), so there is no up-front_unique_finite_zonespass._hi_reducestream-merges the per-block dicts into a single hi_lookup. Scheduler peak memory now scales with the number of distinct zones instead ofn_blocks * n_zones.dask.compute()calls to one.Motivation
On the previous implementation the dask+numpy path did two full graph evaluations before the user ever called
.compute()on the result:_unique_finite_zones(zones_data)to discover zones globally._hi_reduceon all per-block partials to build the hi_lookup dict.On top of that,
_hi_reducebuiltnp.stack(partials_list)producing an(n_blocks, n_zones, 4)float64 array in a single scheduler task — at 240,000 blocks * 1000 zones that is ~7.7 GB held on the scheduler.The streaming dict-merge keeps scheduler memory proportional to the number of distinct zones (tens of KB for typical inputs) and halves the wall time for the dask backend.
Benchmark
Synthetic input: 512×512 float64, chunks=128 (16 blocks), 20 integer zones.
Test plan
pytest xrspatial/tests/test_hypsometric_integral.py— 29 tests passpytest xrspatial/tests/test_zonal.py— full zonal suite pass