Repository navigation
perf(utils): 2-3x faster _StartCountStride (21% on MFDataset reads) - #1524
Merged
Merged
Conversation
Skip numpy conversion for plain ints, slices and Ellipsis, and classify each index array once instead of in three passes. Build start, count and stride from Python lists and convert them once, and broadcast index arrays into the chunk grid instead of using np.apply_along_axis. Output, error types and messages are unchanged. The MFDataset reads in test_multifile_selected_decoding go from 165.75 ms to 131.35 ms (21%). Co-Authored-By: Aiden <aiden@weco.ai>
Collaborator
|
Looks good, thank you. |
Collaborator
|
Can you add a Changelog entry please? |
Refs Unidata#1524 Co-Authored-By: Aiden <aiden@weco.ai>
Contributor
Author
|
Thank you for the review. I added a Changelog entry in 7eeb2c8. It is under a new |
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.
TL;DR
_StartCountStrideinsrc/netCDF4/utils.pyfaster. EveryVariableandMFDatasetread and write calls it.test_multifile_selected_decodingreads: 165.75 ms -> 131.35 ms, 21% faster.pytest test/gives the same result before and after (151 passed, 4 environment failures).What & Why
Before, the function called
np.asarrayandnp.iterableon each index in three separate passes.It did this also for plain ints and slices.
It wrote each value into small numpy arrays one element at a time.
It used
np.apply_along_axisto expand index arrays into the chunk grid.Now, plain ints, slices and
Ellipsisskip the numpy conversion.Each index array is checked once, and the result is reused in the later passes.
The per-dimension values go into Python lists, and the function converts them to arrays once.
Index arrays are reshaped and broadcast into the chunk grid instead of
apply_along_axis.The datashape padding work only runs when data is written.
The output stays identical.
The same checks run in the same order, so the error types and messages do not change.
The returned
start,count,strideandindicesarrays have the same values, dtypes and shapes.Results
Time per call of
_StartCountStride(shape(100, 50, 40)), and one full read for comparison:v[3, :, 2:10]v[[1, 5, 9], :, 0]v[:, mask, :]v[[], :, 1]v[[1,4,7], [2,8], :]v[:]test_multifile_selected_decodingThe gain matters most for many small reads, for example
MFDatasetreads.That case calls the function once for the aggregate and once per file.
For large reads the index work is a very small part of the time, so nothing changes.
Why this is safe and easy to maintain
The change stays inside one private function and keeps its structure and comments.
Behaviour and output do not change, including error messages and the order of the checks.
test_empty_indexing.py,test_utils.py,test_slicing.pyand the rest oftest/cover these paths.Validation
cd test && python -m pytest -q .ncdumpbinary, no zstd plugin).test_multifile_selected_decodingreads.The check also covered fancy, ellipsis and invalid indices: values, dtypes, shapes, masks and exception messages.
AI assistance
The candidate changes were written by Codex (gpt-6.1-sol) and tracked with Weco Observe.
Each candidate was timed and checked for identical output with an automated harness.
This candidate was picked for readability, and the test suite was run on it.
Weco share link: https://dashboard.weco.ai/share/94aIQh5OIBPswhVZFl0zqgU1yAhTbTiV