Support rechunking to seasonal frequency with SeasonalResampler #10519
+227
−36
Conversation
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
added 2 commits
July 10, 2025 00:06
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 15, 2025
dcherian
reviewed
Jul 18, 2025
Comment on lines
+1035
to
+1043
| def _for_chunking(self) -> Self: | ||
| """ | ||
| Return a version of this resampler suitable for chunking. | ||
|
|
||
| For SeasonResampler, this returns a version with drop_incomplete=False | ||
| to prevent data from being silently dropped during chunking operations. | ||
| """ | ||
| return type(self)(seasons=self.seasons, drop_incomplete=False) | ||
|
|
There was a problem hiding this comment.
Suggested change
| def _for_chunking(self) -> Self: | |
| """ | |
| Return a version of this resampler suitable for chunking. | |
| For SeasonResampler, this returns a version with drop_incomplete=False | |
| to prevent data from being silently dropped during chunking operations. | |
| """ | |
| return type(self)(seasons=self.seasons, drop_incomplete=False) |
dcherian
reviewed
Jul 18, 2025
|
|
||
| # Create a temporary resampler that ignores drop_incomplete for chunking | ||
| # This prevents data from being silently dropped during chunking | ||
| resampler_for_chunking = self._for_chunking() |
There was a problem hiding this comment.
Suggested change
| resampler_for_chunking = self._for_chunking() | |
| resampler_for_chunking = type(self)(seasons=self.seasons, drop_incomplete=False) |
dcherian
reviewed
Jul 18, 2025
Comment on lines
+1140
to
+1158
| data = create_test_data() | ||
| for chunks in [1, 2, 3, 4, 5]: | ||
| rechunked = data.chunk({"dim1": chunks}) | ||
| assert rechunked.chunks["dim1"] == (chunks,) * (8 // chunks) + ( | ||
| (8 % chunks,) if 8 % chunks else () | ||
| ) | ||
|
|
||
| rechunked = data.chunk({"dim2": chunks}) | ||
| assert rechunked.chunks["dim2"] == (chunks,) * (9 // chunks) + ( | ||
| (9 % chunks,) if 9 % chunks else () | ||
| ) | ||
|
|
||
| rechunked = data.chunk({"dim1": chunks, "dim2": chunks}) | ||
| assert rechunked.chunks["dim1"] == (chunks,) * (8 // chunks) + ( | ||
| (8 % chunks,) if 8 % chunks else () | ||
| ) | ||
| assert rechunked.chunks["dim2"] == (chunks,) * (9 // chunks) + ( | ||
| (9 % chunks,) if 9 % chunks else () | ||
| ) |
There was a problem hiding this comment.
Suggested change
| data = create_test_data() | |
| for chunks in [1, 2, 3, 4, 5]: | |
| rechunked = data.chunk({"dim1": chunks}) | |
| assert rechunked.chunks["dim1"] == (chunks,) * (8 // chunks) + ( | |
| (8 % chunks,) if 8 % chunks else () | |
| ) | |
| rechunked = data.chunk({"dim2": chunks}) | |
| assert rechunked.chunks["dim2"] == (chunks,) * (9 // chunks) + ( | |
| (9 % chunks,) if 9 % chunks else () | |
| ) | |
| rechunked = data.chunk({"dim1": chunks, "dim2": chunks}) | |
| assert rechunked.chunks["dim1"] == (chunks,) * (8 // chunks) + ( | |
| (8 % chunks,) if 8 % chunks else () | |
| ) | |
| assert rechunked.chunks["dim2"] == (chunks,) * (9 // chunks) + ( | |
| (9 % chunks,) if 9 % chunks else () | |
| ) |
dcherian
reviewed
Jul 18, 2025
| ) | ||
|
|
||
| # Test standard seasons | ||
| rechunked = ds.chunk(x=2, time=SeasonResampler(["DJF", "MAM", "JJA", "SON"])) |
There was a problem hiding this comment.
we'll need to error on a missing season like this:
rechunked = ds.chunk(x=2, time=SeasonResampler(["DJF", "MAM", "SON"]))
dcherian
reviewed
Jul 18, 2025
| chunks = chunks.dropna(name).astype(int) | ||
| chunks_tuple: tuple[int, ...] = tuple(chunks.data.tolist()) | ||
| return chunks_tuple | ||
| return resampler.compute_chunks(name, variable) |
There was a problem hiding this comment.
Suggested change
| return resampler.compute_chunks(name, variable) | |
| newchunks = resampler.compute_chunks(name, variable) | |
| if sum(newchunks) != variable.shape[0]: | |
| raise ValueError(f"Logic bug in rechunking using {resampler!r}. New chunks tuple does not match size of data. Please open an issue.") | |
| return newchunks |
Let's protect ourselves a bit from logic bugs in the resampler
dcherian
reviewed
Jul 18, 2025
| # Test standard seasons | ||
| rechunked = ds.chunk(x=2, time=SeasonResampler(["DJF", "MAM", "JJA", "SON"])) | ||
| # Should have multiple chunks along time dimension | ||
| assert len(rechunked.chunksizes["time"]) > 1 |
There was a problem hiding this comment.
let's assert an actual value here.
dcherian
reviewed
Jul 18, 2025
Comment on lines
+1169
to
+1173
| N = 365 * 2 # 2 years | ||
| if use_cftime: | ||
| time = xr.date_range("2001-01-01", periods=N, freq="D", use_cftime=True) | ||
| else: | ||
| time = xr.date_range("2001-01-01", periods=N, freq="D") |
There was a problem hiding this comment.
Suggested change
| N = 365 * 2 # 2 years | |
| if use_cftime: | |
| time = xr.date_range("2001-01-01", periods=N, freq="D", use_cftime=True) | |
| else: | |
| time = xr.date_range("2001-01-01", periods=N, freq="D") | |
| N = 366 + 365 # 2 years | |
| time = xr.date_range("2000-01-01", periods=N, freq="D", use_cftime=use_cftime) |
By starting in 2000, we can check leap year logic
There was a problem hiding this comment.
Can we parameterize this over calendars too? 360_day, noleap and standard should be good enough.
dcherian
reviewed
Jul 18, 2025
| {"x": 2, "time": SeasonResampler(["DJFM", "AM", "JJA", "SON"])} | ||
| ) | ||
| # Should have multiple chunks along time dimension | ||
| assert len(rechunked.chunksizes["time"]) > 1 |
There was a problem hiding this comment.
here too let's assert actual chunks tuple
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.
whats-new.rstapi.rstusers could not use
SeasonResamplerfor chunking operations in xarray, despite it being a natural fit for seasonal data analysis. When attemptingds.chunk(time=SeasonResampler(["DJF", "MAMJ", "JAS", "ON"])), users encountered obscure errors because the chunking logic was hardcoded to only work withTimeResamplerobjects. This limitation prevented efficient seasonal analysis workflows and forced users to use workarounds or manual chunking strategies.Now Added a generalized chunking approach by adding a
resolve_chunksmethod to theResamplerbase class and updating the chunking logic to work with allResamplerobjects, not justTimeResampler. We also added a_for_chunkingmethod toSeasonResamplerthat ensuresdrop_incomplete=Falseduring chunking operations to prevent silent data loss. The solution maintains full backward compatibility with existingTimeResamplerfunctionality while enabling seamless seasonal chunking