Dask-native SCRIP corner dedup when opened with chunks= - #1775
Merged
Merged
Conversation
A SCRIP file stores one (lon, lat) pair per face-corner, so the corner table is n_face * n_corners rows even though adjacent faces share almost every vertex. The reader deduped that table with Polars, which needs the whole thing resident: 53.6 GiB of float64 for a 300M-element np4 grid, before the join copy. That is why large np4 SCRIP grids fail to open. When the grid is opened with chunks=, dedup with a dask shuffle instead, which spills partition-by-partition, and build the inverse index lazily with map_blocks against the (small) unique-node table. The eager path is unchanged.
Without a distributed client dask picks an in-memory shuffle. On the 226M-face ne2048np4 grid that peaked at 17.9 GiB against 12.1 GiB for the disk shuffle, for the same answer. Pin it, but only when the caller has not chosen a method and no client is running, since p2p is the better choice on a real cluster. Also drop the intermediate DataFrame once its two columns are out. The unique-node table is 3.4 GiB at this size, so holding it past its last use is not free. Together these take the full open from ~31.0 to ~29.3 GiB peak (3 runs each).
The block lookup relied on a sort-after-join to keep each block aligned with its inputs. Polars' own docs say not to rely on an observed join order without requesting one, so request it with maintain_order="left" and check the row count survived. Misalignment here would not raise -- it would build every face from the wrong corners. Lift the closure to module level so that contract is testable, and widen the tests: parametrize over chunks= spellings including "auto" (what callers actually pass), compare corners in winding order rather than sorted, cover the radians path, and assert the connectivity is still lazy after open so a future .compute() in the reader fails here rather than only on a multi-GB file.
This was referenced Sep 19, 2026
Closed
The eager reader deduplicated corners with a Polars hash-join. Hashing is the faster algorithm and that is why it was written that way; what it costs is space, because the join holds the corner table, a hash table keyed on every row of it, and the joined result at the same time. np.lexsort trades the better asymptotic time for a bounded footprint -- an argsort and two gathers over index arrays. On a 4M-face, 48M-corner SCRIP grid the whole open goes from 13.9 GiB to 3.6 GiB at roughly twice the wall time. That is the right trade for a format whose corner table is the largest thing in the file and the reason large meshes fail to open at all. The dask path keeps its hash-based shuffle. Measured on a 3.6-billion-row corner table, sorting there peaked at 191 GiB against the shuffle's 118 -- lexsort must materialize what the shuffle deliberately streams, so the two paths want opposite algorithms for the same reason. No benchmark regression: the only SCRIP grid in ASV is outCSne8 at 384 faces, where the sort is faster than the join (4.2 ms against 6.6 ms) because Polars' fixed setup dominates at that size, and open_grid happens in setup() which ASV does not time. The crossover is near 10k faces. Geometry is unchanged. Node numbering differs -- sorted rather than first-appearance -- so the tests compare the unique node set and each face's corners in winding order, both identical on every SCRIP mesh in the suite.
Sevans711
added
run-benchmark
Collaborator
|
A few quick questions, before I take a closer look:
|
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Contributor
Author
|
Fair point—I thought about this while working on the dedup paths. The eager path should now be deterministic because it uses a lexicographic sort, so it likely fixes #1738 for non-chunked opens. The Dask path may produce a different node order because of the shuffle. I don’t think eager and chunked loading necessarily need identical node numbering as long as both produce the same mesh with correctly remapped connectivity. That said, consistent ordering would be nice. I’ll measure the memory and runtime overhead of sorting the Dask-generated unique-node table; if it’s reasonably small, I can make both paths produce the same ordering here. Otherwise, I’d prefer to keep this PR focused on bounded-memory loading and handle canonical ordering separately in #1738. |
Sevans711
requested changes
Sep 23, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Mostly looks good, just some small inline comments!
The inconsistent ordering of nodes/edges between the chunks=... and non-dask and small/large numpy cases doesn't need to block this PR, but can be reassessed/addressed while looking into #1738.
Contributor
Author
|
Thanks, fixed those in 085b255. Let's let ASV finish. |
Sevans711
approved these changes
Sep 23, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Previous comments have been resolved, and there is no longer any performance degradation in the ASV benchmarks suite. Approving!
- Replace the per-block polars join with a numba binary search over the unique table, sorted once; the join rehashed all unique nodes per block. - Pass shuffle_method to drop_duplicates; dask swapped a config "disk" default for "tasks", so the pin never applied. - Free the sorted copies before building ids on the eager path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
NaN != NaN, so the lexsort dedup made every NaN corner (a decoded _FillValue) its own node, while the dask path merged them. Compare NaN as equal, only when a NaN is present, so both paths build the same mesh. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Give an in-memory corner array the other's chunks before the dask dedup; map_blocks passed it whole to every block, pairing corners with the wrong latitudes. - Return empty arrays from the eager dedup when there are no corners instead of indexing a zero-length array. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Collaborator
|
@rajeeja I added two crash fixes, a faster lookup algorithm, and better NaN handling. I'd be curious to see what the new performance numbers are on the big meshes. Also, the binary search algorithm included here could potentially be optimized with a galloping search. I was seeing ~1.7x over the existing lexsort/binary search on oQU120, but it depends heavily on how relatively-sorted the grid is. On a huge grid, it's probably a better option. I was thinking that could be a further PR if we wanted to try it. The sorting/searching may not take long enough to bother though. |
The disk pin was justified by a 17.9 vs 12.1 GiB measurement that compared dask's default against an explicit method, not disk against tasks: the pin was set through dask.config, which drop_duplicates overrides, so its "disk" arm never ran a disk shuffle. Passing the method as an argument made the pin effective and so made it measurable, and disk then loses on both axes -- 315.5 s/86.8 GiB against tasks' 249.1 s/85.4 GiB on the 300M-face CONUS RRM grid for byte-identical connectivity, and 2.71 s/7.01 GiB against 2.25 s/6.54 GiB at 12M faces. Dask's own default matches tasks. Drop the pin and forward only a method the caller configured. This retires _distributed_client_active, which existed to decide whether to pin, and corrects the test that asserted the pinned default. Also record the locality the binary-search lookup depends on, measured as the normalized mean jump between consecutive lookups (1/3 being random): 0.0001 on the CONUS RRM grid, 0.0003 on a 74M-corner ESMF mesh, 0.0007-0.0028 on regular lat-lon, 0.02-0.08 on MPAS SFC and FESOM.
Contributor
Author
|
@cmdupuis3 This is great work, thank you. The lookup change is what makes this PR actually work at the scale it was written for, and both crash fixes are real bugs. I ran all three commits against the CONUS RRM grid on Chrysalis plus a locality survey, and your work turned up something that led me to change one of my own decisions. I reproduced each defect against its own parent before reviewing the fix, and all three hold up:
Full-scale numbers on the 63 GB CONUS RRM grid — 299,999,162 faces, 3,599,989,944 corners, 500,001,289 unique nodes, on 253 GB nodes:
Your commit message caught a real problem:
On your locality question — whether most grids store faces and nodes close together or whether some end up basically random. Good question to ask, since the binary search depends on the answer, so I measured it. Mean jump between consecutive lookups over the sorted unique-node table, normalized by table length, where 1/3 is the random baseline:
On galloping:
Where this leaves us:
|
cmdupuis3
approved these changes
Sep 25, 2026
Collaborator
|
@rajeeja Can you briefly summarize your understanding of what was changed after I approved? I'm trying to understand your message but there are a lot of bullet points and it is difficult for me to summarize and/or determine what is most important. |
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
3 participants
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.
Closes #1774
Overview
A SCRIP file stores one (lon, lat) pair per face-corner, so the corner table is
n_face × n_cornersrows and shared vertices are restated once per touching face._to_ugriddeduplicated that with Polars, which needs the whole table resident — 53.6 GiB for a 300M-face np4 grid — and never dispatched onchunks=, so large files could not be opened at all.Each path now gets the dedup that suits it:
chunks=np.lexsortchunks=Measured
The 63 GB file now opens: 299,999,162 faces, 500,001,289 unique nodes.
Sorting was also tried on the dask path and is worse there — 191 GiB against the shuffle's 118, because
lexsortmust materialize what the shuffle streams. The two paths want opposite algorithms for the same reason.Benchmarks
No ASV regression. The only SCRIP grid in ASV is
outCSne8at 384 faces, where the sort is faster (4.2 ms vs 6.6 ms) because Polars' fixed setup dominates at that size, andopen_gridruns in untimedsetup(). Crossover is near 10k faces; above it the trade is roughly 1:1, time for memory.Correctness
Node numbering differs — sorted rather than first-appearance — so tests compare the unique node set and each face's corners in winding order, since winding determines area sign. Identical on every SCRIP mesh in the suite, and chunk-size independent.
Two findings from measurement: the dask shuffle defaults to an in-memory method without a
distributedclient (pinning it to disk saved 32%), and Polars does not guarantee join order, somaintain_order="left"is requested rather than assumed.