Support rechunking to seasonal frequency with SeasonalResampler #10519

dhruvak001 · 2025-07-09T18:35:33Z

Closes Support rechunking to seasonal frequency with SeasonalResampler #10425
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

users could not use SeasonResampler for chunking operations in xarray, despite it being a natural fit for seasonal data analysis. When attempting ds.chunk(time=SeasonResampler(["DJF", "MAMJ", "JAS", "ON"])), users encountered obscure errors because the chunking logic was hardcoded to only work with TimeResampler objects. 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_chunks method to the Resampler base class and updating the chunking logic to work with all Resampler objects, not just TimeResampler. We also added a _for_chunking method to SeasonResampler that ensures drop_incomplete=False during chunking operations to prevent silent data loss. The solution maintains full backward compatibility with existing TimeResampler functionality while enabling seamless seasonal chunking

dcherian · 2025-07-15T19:53:33Z

xarray/core/types.py

@@ -201,7 +201,7 @@ def copy(
 # FYI in some cases we don't allow `None`, which this doesn't take account of.
 # FYI the `str` is for a size string, e.g. "16MB", supported by dask.
 T_ChunkDim: TypeAlias = str | int | Literal["auto"] | tuple[int, ...] | None  # noqa: PYI051
-T_ChunkDimFreq: TypeAlias = Union["TimeResampler", T_ChunkDim]
+T_ChunkDimFreq: TypeAlias = Union["TimeResampler", "SeasonResampler", T_ChunkDim]


Suggested change

T_ChunkDimFreq: TypeAlias = Union["TimeResampler", "SeasonResampler", T_ChunkDim]

T_ChunkDimFreq: TypeAlias = Union["Resampler", T_ChunkDim]

Good time to generalize here

dcherian · 2025-07-15T19:56:50Z

xarray/groupers.py

@@ -52,6 +52,7 @@
    "EncodedGroups",
    "Grouper",
    "Resampler",
+    "SeasonResampler",


Suggested change

"SeasonResampler",

"SeasonGrouper",

"SeasonResampler",

This was an oversight.

dcherian · 2025-07-15T19:59:34Z

xarray/groupers.py

+        if variable.ndim != 1:
+            raise ValueError(
+                f"chunks={self!r} only supported for 1D variables. "
+                f"Received variable {name!r} with {variable.ndim} dimensions instead."
+            )


I think this check should go in _resolve_frequency in dataset.py

dcherian · 2025-07-15T20:01:07Z

xarray/groupers.py

+        # 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() if hasattr(self, "_for_chunking") else self
+        )


Suggested change

# 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() if hasattr(self, "_for_chunking") else self

)

Instead let's define this compute_chunks method independently on TimeResampler and SeasonResampler. We would like to add FloatResampler (#4008) in the future, so we should keep it general. For example, the "contains datetime" check should be specific to the resamplers

dcherian · 2025-07-15T20:07:48Z

xarray/groupers.py

@@ -169,7 +170,60 @@ class Resampler(Grouper):
    Currently only used for TimeResampler, but could be used for SpaceResampler in the future.
    """

-    pass
+    def resolve_chunks(self, name: Hashable, variable: Variable) -> tuple[int, ...]:


Suggested change

def resolve_chunks(self, name: Hashable, variable: Variable) -> tuple[int, ...]:

def compute_chunks(self, name: Hashable, variable: Variable) -> tuple[int, ...]:

dcherian · 2025-07-15T20:14:26Z

doc/whats-new.rst

@@ -65,7 +69,7 @@ Bug fixes
  creates extra variables that don't match the provided coordinate names, instead
  of silently ignoring them. The error message suggests using the factory method
  pattern with :py:meth:`xarray.Coordinates.from_xindex` and
-  :py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`).
+  :py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`, , :pull:`10503`).


Suggested change

:py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`, , :pull:`10503`).

:py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`, :pull:`10503`).

dcherian · 2025-07-15T20:14:43Z

doc/whats-new.rst

+- Allow skipping the creation of default indexes when opening datasets (:pull:`8051`).
+  By `Benoit Bovy <https://github.com/benbovy>`_ and `Justus Magin <https://github.com/keewis>`_.


Suggested change

- Allow skipping the creation of default indexes when opening datasets (:pull:`8051`).

By `Benoit Bovy <https://github.com/benbovy>`_ and `Justus Magin <https://github.com/keewis>`_.

bad merge

dcherian · 2025-07-15T20:17:38Z

xarray/groupers.py

+        chunks = (
+            DataArray(
+                np.ones(variable.shape, dtype=int),
+                dims=(name,),
+                coords={name: variable},
+            )
+            .resample({name: resampler_for_chunking})
+            .sum()
+        )
+        # When bins (binning) or time periods are missing (resampling)
+        # we can end up with NaNs. Drop them.
+        if chunks.dtype.kind == "f":
+            chunks = chunks.dropna(name).astype(int)
+        chunks_tuple: tuple[int, ...] = tuple(chunks.data.tolist())


If you're interested, a better approach here would be to pass this DataArray to Resampler.factorize() and use the group_indices of the returned EncodedGroups object to figure out the chunks tuple.

This is another reason I'd prefer just sending the variable to a method on the Resampler. That way the resamplers are free to optimize if they can.

dcherian · 2025-07-15T20:19:26Z

xarray/tests/test_dataset.py

+        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 ()
+            )


What is the point of this test?

dcherian · 2025-07-15T20:20:33Z

xarray/tests/test_dataset.py

+
+        # Test standard seasons
+        rechunked = ds.chunk(x=2, time=SeasonResampler(["DJF", "MAM", "JJA", "SON"]))
+        expected = tuple(


why is expected needed?

dcherian · 2025-07-15T20:21:29Z

xarray/tests/test_dataset.py

+        import dask.array
+
+        N = 365 * 2  # 2 years
+        time = xr.date_range("2001-01-01", periods=N, freq="D")


let's parameterize this over use_cftime=[True, False]

Support chunking

0f10d91

github-actions bot added the topic-documentation label Jul 9, 2025

dhruvak001 changed the title ~~Support chunking~~ Support rechunking to seasonal frequency with SeasonalResampler Jul 9, 2025

DHRUVA KUMAR KAUSHAL added 2 commits July 10, 2025 00:06

whats new

3d4b62e

error resolving

527c48c

github-actions bot added the topic-typing label Jul 9, 2025

dcherian removed the topic-documentation label Jul 11, 2025

Merge branch 'main' into issue#10425

138e747

github-actions bot added the topic-documentation label Jul 12, 2025

dcherian reviewed Jul 15, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Support rechunking to seasonal frequency with SeasonalResampler #10519

Support rechunking to seasonal frequency with SeasonalResampler #10519

Uh oh!

dhruvak001 commented Jul 9, 2025

Uh oh!

dcherian Jul 15, 2025 •

edited

Loading

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025 •

edited

Loading

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025 •

edited

Loading

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025

Uh oh!

dcherian Jul 15, 2025

Uh oh!

Uh oh!

	T_ChunkDimFreq: TypeAlias = Union["TimeResampler", "SeasonResampler", T_ChunkDim]
	T_ChunkDimFreq: TypeAlias = Union["Resampler", T_ChunkDim]

	def resolve_chunks(self, name: Hashable, variable: Variable) -> tuple[int, ...]:
	def compute_chunks(self, name: Hashable, variable: Variable) -> tuple[int, ...]:

	:py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`, , :pull:`10503`).
	:py:meth:`Dataset.assign_coords` for advanced use cases (:issue:`10499`, :pull:`10503`).

		- Allow skipping the creation of default indexes when opening datasets (:pull:`8051`).
		By `Benoit Bovy <https://github.com/benbovy>`_ and `Justus Magin <https://github.com/keewis>`_.

Uh oh!

Support rechunking to seasonal frequency with SeasonalResampler #10519

Are you sure you want to change the base?

Support rechunking to seasonal frequency with SeasonalResampler #10519

Uh oh!

Conversation

dhruvak001 commented Jul 9, 2025

Uh oh!

dcherian Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dcherian Jul 15, 2025 •

edited

Loading

dcherian Jul 15, 2025 •

edited

Loading

dcherian Jul 15, 2025 •

edited

Loading