Repository navigation
Support reusing an existing Grid when opening datasets - #1827
Conversation
Sevans711
added
run-benchmark
There was a problem hiding this comment.
🟡 Changes recommended
Structured node-only and edge-only members can crash, and overlapping chunk aliases can produce incorrect chunk sizes.
2 open issues
What changed in this PR
Adds support for reusing an existing Grid across dataset opens, addressing #1818 without rebuilding topology or cached properties.
Changes:
- Supports Grid inputs in both dataset-opening APIs.
- Preserves shared-grid state and maps member dimensions and chunks.
- Adds regression tests, API documentation, and a synthetic benchmark.
| File | Description |
|---|---|
| uxarray/core/api.py | Implements Grid reuse and member mapping. |
| test/core/test_grid_reuse.py | Tests reuse, chunking, mapping, and input validation. |
| benchmarks/grid_reuse.py | Compares rebuilt and reused Grid inputs. |
Review effort: Balanced
| source_dims = grid._source_dims_dict.copy() | ||
| if grid.source_grid_spec == "Structured": | ||
| lon_dim, lat_dim = source_dims["n_face"] | ||
| if "n_face" in ds.dims and lon_dim not in ds.dims and lat_dim not in ds.dims: |
| chunks = chunks.copy() | ||
| for source_dim, ugrid_dim in grid_filename_or_obj._source_dims_dict.items(): | ||
| if isinstance(ugrid_dim, str) and ugrid_dim in chunks: | ||
| chunks[source_dim] = chunks[ugrid_dim] |
|
Hi @rajeeja, as the person who requested the github copilot review, could you clarify a bit further how important is it for its review comments to be addressed? I don't think it has been particularly common for recent uxarray PRs to get copilot reviews, and I think it could be especially helpful to clarify here because this PR was drafted by a first-time contributor. E.g., are you just seeing it as a helpful starting point for your own review, or are you hoping its suggestions might be treated with the same level of importance/consideration as comments from a human reviewer? Sidenote: reviewing this PR is high on my to-do list, I'm hoping to get to it sometime very soon (probably tomorrow or within the next few days)! |
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
|
Hey @Sevans711, I use Copilot as a starting point for my own review, not as a replacement for human review. Some comments are nitpicks, while others are genuinely useful once verified. I confirmed these two issues:
Please fix these two cases. |
|
@rajeeja Thanks for confirming both cases. Fixed in 9b4fc45:
The two new regression methods exercise both opening APIs, in-memory/Dask members and synthetic NetCDF files, both alias insertion orders, and preservation of the caller's chunk dictionary. They produce 20 failing subcases against the previous head; the revised file passes all 17 tests and 58 subtests. Five additional existing synthetic API/structured-grid tests and all applicable changed-file formatting/lint hooks also pass. An independent synthetic oracle passes 91/91 cases after reproducing 58 failures on the old implementation. Normal Numba JIT was enabled (Linux, Python 3.12). This revision is limited to these two cases. I have not rerun the complete repository suite or the full ASV performance comparison locally; the earlier memory/performance measurements are not claimed resolved. The existing backend-close-callback limitation and PR #1828 dimension-order scope are unchanged. |
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
Thank you for putting this together! Overall, this looks really great, and being able to re-use an existing grid during open_dataset is a very welcome addition to uxarray. I read through all the changes, confirmed locally that tests pass, and tried out this new feature with a couple of small examples. I left a few minor inline comments but after these small points get fixed I think I would be ready to approve!
|
@Sevans711 Thank you for the careful review and for trying the examples. Addressed the four points in 7acbe33:
Validation on this revision: 17 Grid-reuse tests and 58 subtests passed, independently rerun with the same result; all six benchmark setup/time/teardown smoke cases passed; Ruff lint/format and applicable whitespace, EOF, docstring-first and debug-statement hooks passed. Three adjacent synthetic API tests also passed. Three additional fixture-dependent API checks could not complete in this environment (one missing mesh fixture and two read-only data-cache failures). I have not rerun the complete repository suite or full ASV timing comparison, and the new remote CI is not yet confirmed. AI Usage: OpenAI Codex assisted with this revision and independent verification. The changes are ready for another look; I have left the review threads unresolved as requested by the contribution guide. |
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
@sunxiuguo Thank you for fixing all of my previous points (and also for leaving the threads unresolved; I always find it helpful to go through my previous points one-by-one while re-reviewing)!
I'd be happy to approve, but it looks like some of the CI tests failed. Since it was only the windows tests, I am guessing it is caused by #1844, which has been fixed on main (via a workaround for now, which is why that issue isn't closed yet). Would you be able to update your fork & branch for this PR to be up to date with main?
I suspect that then the tests would pass. Once tests are passing, everything else here looks good and I would be happy to approve!
|
@Sevans711 Thanks — I merged current upstream Local validation of that candidate: the 17 Grid-reuse tests (58 subtests) and the longitude-related tests passed. The combined focused run finished with 26 passed and 1 failed: the thread-pool test's child process exited with The new CI run and ASV run are awaiting maintainer approval; neither has executed jobs yet. I have not treated them as passing. AI Usage: OpenAI Codex assisted with the update and verification. |
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
All my previous comments have been resolved, all tests passing, and benchmarking results still look good (the handful of apparent degradations are just a quirk of the ASV suite right now, as can be seen in #1710, not a result of the changes in this PR). Happy to approve, and thank you again @sunxiuguo for all your efforts here!
erogluorhan
left a comment
erogluorhan
left a comment
There was a problem hiding this comment.
Thanks for working on this! I have some inline comments below
|
Putting here a report-back from a discussion I just had with @erogluorhan, to keep you informed, @sunxiuguo: The biggest point we discussed was that Because of that, all the new structured-inputs-related tests proposed in this PR are most likely no longer necessary and should be removed (@erogluorhan flagged some of them with inline comments above). If this also makes it easy to condense the tests a bit (e.g., <300 lines instead of >500 lines), it would be helpful to uxarray maintainers if these could be moved to test_api.py instead of being added into their own file! |
|
Looks functionally clear on my end. Passing an opened |
erogluorhan
left a comment
erogluorhan
left a comment
There was a problem hiding this comment.
Thanks for addressing the comments from the previous cycle. Some more review comments below.
rajeeja
left a comment
rajeeja
left a comment
There was a problem hiding this comment.
Thank you @sunxiuguo for this useful addition that lets datasets share an existing Grid and its cached properties without rebuilding it, with explicit chunking checks to protect that shared state.
Closes #1818
Overview
Allow
open_dataset(existing_grid, data)andopen_mfdataset(existing_grid, paths)to reuse the same Grid, its topology and cached properties without reconstructing it or sending it to a file backend.The current implementation preserves the Grid's existing chunking and source-dimension mapping. Data chunks are mapped without mutating the caller's dictionary, including overlapping source/canonical dimension names. Canonical dimensions are preserved when node and face counts coincide. Grid-construction options
use_dual=Trueand nonemptygrid_kwargsraise actionable errors for an existing Grid.The intended scope is unstructured grid/data inputs. Following the maintainer clarification, structured-input-related tests and usage guidance were removed, and the retained reuse tests were consolidated into
test/core/test_api.py. Ordinary cases use the existing quad-hexagon grid; a small equal-count topology covers dimension/alias regressions. Future warnings/errors for structured lat/lon inputs are tracked separately in #1854.The November 3 revision implements the agreed chunking contract: passing an existing Grid, non-None
chunks, andchunk_grid=True(including its default) raisesValueErrorwith guidance to callGrid.chunk()explicitly. Both opening APIs reject this combination before opening the data.chunks=Nonepreserves the shared Grid, its current chunking, and caches;chunk_grid=Falsepermits data-only chunking. No shared Grid is silently copied or rechunked. The public API docstrings describe the contract, and the tests now select data-only chunking explicitly rather than hiding that option in a fixture.The branch merges main through
79b7aa6e999474fb09f1aec8e52303a19a9aaebd, retaining both histories. The sole merge conflict was inuxarray/core/api.py: the Grid-reuse path is preserved together with upstream's removal of deprecatedlatlonforwarding and its updated helper signature. Upstream changes to the utilities and API tests are retained. Alignment validation and its documentation remain scoped to #1853.Expected Usage
In-memory xarray Dataset members and multi-file data are covered as well.
Validation
November 3 local revision
ValueErrorwas raised. The final reuse group passed 90 tests, including default/explicitchunk_grid=True,chunks=Noneunder either flag, prechunked Grids, data-only chunks, and caller/Grid/cache preservation.python -m pytest test/core test/io test/grid/grid/test_io.py test/grid/grid/test_initialization.py -qpassed 505 tests, with 33 warnings, in 132.90 seconds. This includes the full API tests and the retained upstreamdrop_variablesregression. It is not a full-repository test run.git diff --checkpassed. All six existing synthetic opening benchmark combinations passed a smoke run; no new ASV performance comparison is claimed.723f6d8a273b42aa8fe26ec755ed0cc7048d14cchas tree51d55bb5a3f1197e8d2299d0522867d70652761e, identical to local tested/reviewed candidatebd1f080e424bd8640e7dd5189e19434c4a58b701. The GitHub connector created the published commit metadata, so the commit SHA differs while the source tree is unchanged. Merge commit3d33843e745f880595826e621cac7f1e8d368d5bretains both original headd2c2ad9fand main79b7aa6eas parents. A normal, non-force update was guarded by the expected original head. Remote commit/tree and all four changed files were read back; file bytes, Git blobs, and SHA-256 hashes match the tested source.action_requiredwith zero jobs; they await maintainer approval and have not run successfully. Read the Docs build 35617442 is pending. No new-head hosted full-suite success, Windows success, rendered-documentation-build success, or ASV no-regression result is claimed.Local review revisions — run on 2026-10-27
62f4d411passed 60 targeted tests, 19 deselected: 57 unstructured Grid-reuse cases and 3 unchanged synthetic API tests. The retained reuse group is 295 lines plus 4 import lines, replacing the separate 535-line module.d2c2ad9f, also passed 60 targeted tests, 19 deselected (11.95 seconds). Ruff lint/format, docstring-first/debug-statement hooks, whitespace and patch-applicability checks passed; executable AST was unchanged from62f4d411.Environment for these local checks: Linux, Python 3.12.14; xarray 2026.9.0, Dask 2026.8.0, NumPy 2.5.4.
Hosted results — checked 2026-10-28
At head
d2c2ad9fef5400967cf3d93f9cd64d0ac3aaf1d1, CI run 40573219453 and ASV run 40573219452 both completed successfully on attempt 2.Successful ASV execution is not a no-regression result. The posted comparison, labeled
db33c50bversus7fd57840, reports increases including 1.36× import peak memory, approximately 1.26× face-bounds cold-start RSS, and 1.18×time_two_prod. These observations are not claimed resolved or attributed to a particular change.Earlier implementation evidence and limitations
The original PR description recorded reproduction of the reported netCDF TypeError for both filename and Dataset members, 266 tests plus 28 subtests passing, and an independent 56 tests plus 28 subtests with six benchmark smoke cases. Those are historical results, not reruns on the current head. The synthetic benchmark covers all six file/Dataset/Grid input combinations; its tiny-fixture timings are not production-performance claims.
An earlier local full-suite attempt ended with exit 137 and did not complete. A later combined integration check had 26 passes and one failure when the thread-pool test's child process exited with
-9; its cause was not established. The subsequent hosted successes above do not retroactively turn those local attempts into passes. Existing UxDataset close-callback behavior is not refactored by this PR.PR Checklist
General
Testing & Benchmarking
Documentation and Examples
AI Disclosure
AI Usage: OpenAI Codex assisted with investigation, implementation, synthetic tests, benchmarking, and independent code review.