Avoid converting chunked data to Numpy arrays (i.e. .values calls) - #1588
Open
cmdupuis3 wants to merge 34 commits into
Open
Avoid converting chunked data to Numpy arrays (i.e. .values calls)#1588cmdupuis3 wants to merge 34 commits into
.values calls)#1588cmdupuis3 wants to merge 34 commits into
Conversation
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Collaborator
Author
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Collaborator
Author
|
There's some spurious asv results from machine variability, but the benchmarks consistenly show peak-mem reductions and some speedups for cross-sections. |
Collaborator
Author
|
pre-commit.ci autofix |
Sevans711
self-requested a review
July 27, 2026 16:18
Sevans711
requested changes
Jul 27, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Hi @cmdupuis3, thank you for proposing these changes! Overall these look like good clean changes which should help improve the scalability of uxarray.
I have some notes/suggestions/requested changes. Primarily, I noticed that there are missing regression tests. This may be especially important for places with significant changes (more significant than just replacing obj.values with obj.data):
UxDataArray.integrate()UxDataArrayCrossSectionAccessor.__call__()uxarray.plot.matplotlib._nearest_neighbor_resample()RemapAccessor.apply_weights()
It may be nice to add a regression test for UxDataset.to_xarray() too, even though the changes here are minor, since that feels like a very core part of the functionality.
The tests could be similar to what you added already in test_topological_agg.py (plus my additional request on that file): ensure that numpy and dask inputs ultimately give the same values (plus assert that the dask inputs lead to dask outputs).
Other parts of the code changes here might benefit from similar tests, but I don't know if that should be necessary. For example, the changes in _geos just replace obj.values with obj.data; does a regression test need to be added for that or no? Curious to hear from @erogluorhan and/or @rajeeja on this question in particular.
Misc. note: I'm not sure if I fully understand the changes to uxarray/cross_sections files. I need to make sure to take a closer look at those during a subsequent review.
erogluorhan
changed the title
.values calls) in the code
erogluorhan
changed the title
.values calls) in the code.values calls)
A stash pop was committed in 9cfd53e with its markers intact, making the module a syntax error. Both sides were additive, so keep both. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sevans711
requested changes
Aug 10, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
This looks almost ready to merge, I just have a few minor comments remaining!
Requires response:
- I noticed the
mpas_ocean.Integrate.time_integrate('480km')benchmark has worsened slightly. Is this expected behavior? Do you think this is real, a rounding error, or something else? - I would suggest to remove all of the
pytest.importorskipcalls; I left an inline comment with more details.
The rest of this comment does not require response. It includes the notes I wrote down, for future reference.
Other Benchmarking:
- The cross sections benchmarks seem to have improved, which is great!
- I'm ignoring peakmem "improvements" (as per #1605)
- Even though other benchmarks didn't show notable improvements, I think that is fine, because they didn't get worse, either. Also, this PR isn't fully closing the original issue, anyway, it is just addressing part of it.
Testing:
I checked all the cases where I previously requested tests for correctness (numpy & dask give same result) and regression tests (check dask inputs cause dask outputs); here is what I found:
- (Looks done to me! Confirmed tests pass here, fail on main.)
UxDataArray.integrate() - (Looks done to me! Confirmed tests pass here, fail on main. Also, the tests do a good job covering a variety of possible failure cases, by including multiple dimensions and some nan / some non-nan values.)
UxDataArrayCrossSectionAccessor.__call__() - (Not done, but now I'm less sure if it's necessary; left an inline comment to discuss.)
uxarray.plot.matplotlib._nearest_neighbor_resample() - (Looks done to me! Confirmed tests pass here, fail on main. Also, the tests do a good job covering nontrivial cases, like handling multiple dimensions.)
RemapAccessor.apply_weights() - (Not done, but optional; no need to block merging for this.)
UxDataset.to_xarray()
Previously I noted I need to look closer at cross sections. After doing so and considering the new cross sections tests, I am feeling convinced that the cross sections changes should be working as intended. Noting that the extra dims included in those tests might also help with #1461.
Collaborator
Author
It isn't real, you can rerun the benchmarks and it'll probably disappear. There's always a chance with the smaller benchmarks that they'll trigger due to machine variability. 1.1x is pretty suspect. |
This was referenced Aug 10, 2026
Contributor
|
While reviewing this I noticed the numpy aggregation paths drop coords and attrs while your new dask paths preserve them via |
Collaborator
Author
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Sevans711
approved these changes
Aug 14, 2026
Sevans711
left a comment
Sevans711
left a comment
Collaborator
There was a problem hiding this comment.
Looks to me like this is ready to merge into main!
Any guesses as to why the mpas_ocean.CrossSections.time_const_lat benchmarks show improvements now, but didn't show improvements in prior iterations of this PR? (This question is non-blocking, since it's an improvement! Just curious if you have insights about this.)
Collaborator
Author
|
@Sevans711 They've shown improvements for a while now, that's not recent. There are some changes to how cross-sections are caluculated on chunked grids that served as the |
Collaborator
Ah, woops! I confused myself by looking at my prior comment mentioning the benchmarks. Thank you for clarifying! |
Member
|
I have copilot access on github and will ask it to do a review. |
Contributor
There was a problem hiding this comment.
Pull request overview
Improves lazy execution by replacing eager .values conversions with Dask/xarray-compatible operations.
Changes:
- Adds lazy remapping, aggregation, integration, and cross-section paths.
- Preserves chunked backing arrays during grid and format conversions.
- Adds Dask parity tests for key operations.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
uxarray/remap/structured.py |
Preserves lazy arrays during reshaping. |
uxarray/remap/apply_weights.py |
Applies sparse weights blockwise. |
uxarray/remap/accessor.py |
Documents lazy weight application. |
uxarray/plot/matplotlib.py |
Loads only sampled faces. |
uxarray/io/_scrip.py |
Preserves lazy SCRIP coordinates. |
uxarray/io/_geos.py |
Preserves GEOS backing arrays. |
uxarray/grid/slice.py |
Adds lazy connectivity remapping. |
uxarray/grid/grid.py |
Avoids unnecessary grid materialization. |
uxarray/cross_sections/sample.py |
Removes eager Numba gathering. |
uxarray/cross_sections/dataarray_accessor.py |
Implements lazy cross-section gathering. |
uxarray/core/dataset.py |
Preserves backing arrays in duals. |
uxarray/core/dataarray.py |
Adds lazy integration and transformations. |
uxarray/core/aggregation.py |
Implements Dask topological aggregation. |
test/test_cross_sections.py |
Tests lazy cross-sections. |
test/precomputed_weights_test.py |
Tests lazy sparse remapping. |
test/grid/integrate/test_basic.py |
Tests lazy integration. |
test/core/test_topological_agg.py |
Tests Dask aggregations. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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
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.
Partly addresses #1583
Overview
This PR is to resolve suboptimal usage of
.valuesthroughout the repo, primarily by deferring to lazy xarray and dask operations. See the table of usage sites in Issue #1583.This specific PR is for the first two tables of usage sites; dead code and bugfixes aren't really included here. Aside from simple cases, some cases were solved by branching on whether the data type was already chunked, so scalability for some parts will depend on if you're using numpy or xarray/dask arrays at those points.
PR Checklist
General
Testing
Documentation