Support duplicate dimensions in .chunk (#9099)

* Allow duplicate dimensions in chunking

* Address review comments

* fix whats-new

* add comment

* Update xarray/tests/test_dask.py

---------

Co-authored-by: Deepak Cherian <deepak@cherian.net>
Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Author: mraspaud

doc/whats-new.rst

@@ -22,7 +22,8 @@ v2024.06.1 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
-
+- Allow chunking for arrays with duplicated dimension names (:issue:`8759`, :pull:`9099`).
+  By `Martin Raspaud <https://github.com/mraspaud>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
@@ -73,7 +74,6 @@ Bug fixes
   support arbitrary kwargs such as ``order`` for polynomial interpolation (:issue:`8762`).
   By `Nicolas Karasiak <https://github.com/nkarasiak>`_.
 
-
 Documentation
 ~~~~~~~~~~~~~
 - Add link to CF Conventions on packed data and sentence on type determination in the I/O user guide (:issue:`9041`, :pull:`9045`).

xarray/namedarray/core.py

@@ -812,7 +812,12 @@ def chunk(
             chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
 
         if is_dict_like(chunks):
-            chunks = {self.get_axis_num(dim): chunk for dim, chunk in chunks.items()}
+            # This method of iteration allows for duplicated dimension names, GH8579
+            chunks = {
+                dim_number: chunks[dim]
+                for dim_number, dim in enumerate(self.dims)
+                if dim in chunks
+            }
 
         chunkmanager = guess_chunkmanager(chunked_array_type)
 

xarray/tests/test_dask.py

@@ -638,6 +638,13 @@ def counting_get(*args, **kwargs):
 
         assert count[0] == 1
 
+    def test_duplicate_dims(self):
+        data = np.random.normal(size=(4, 4))
+        arr = DataArray(data, dims=("x", "x"))
+        chunked_array = arr.chunk({"x": 2})
+        assert chunked_array.chunks == ((2, 2), (2, 2))
+        assert chunked_array.chunksizes == {"x": (2, 2)}
+
     def test_stack(self):
         data = da.random.normal(size=(2, 3, 4), chunks=(1, 3, 4))
         arr = DataArray(data, dims=("w", "x", "y"))