preserve chunked data when creating DataArray from DataArray (#5984)

* variable: add case for data is DataArray instance in `as_compatible_data`
test: add test for preserving chunks DataArray from DataArray creation

* Update xarray/tests/test_dataarray.py

Co-authored-by: Illviljan <[email protected]>

* tests/test_dataarray: check equal chunks in `test_constructor_from_self_described_chunked`

* add whats-new entry

* Update whats-new.rst

Co-authored-by: Illviljan <[email protected]>
Co-authored-by: Mathias Hauser <[email protected]>

Author: 3 people

doc/whats-new.rst

@@ -46,6 +46,8 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Preserve chunks when creating a :py:class:`DataArray` from another :py:class:`DataArray`
+  (:pull:`5984`). By `Fabian Hofmann <https://github.com/FabianHofmann>`_.
 - Properly support :py:meth:`DataArray.ffill`, :py:meth:`DataArray.bfill`, :py:meth:`Dataset.ffill` and :py:meth:`Dataset.bfill` along chunked dimensions (:issue:`6112`).
   By `Joseph Nowak <https://github.com/josephnowak>`_.
 

xarray/core/variable.py

@@ -198,11 +198,13 @@ def as_compatible_data(data, fastpath=False):
 
     Finally, wrap it up with an adapter if necessary.
     """
+    from .dataarray import DataArray
+
     if fastpath and getattr(data, "ndim", 0) > 0:
         # can't use fastpath (yet) for scalars
         return _maybe_wrap_data(data)
 
-    if isinstance(data, Variable):
+    if isinstance(data, (Variable, DataArray)):
         return data.data
 
     if isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES):

xarray/tests/test_dataarray.py

@@ -30,6 +30,7 @@
     ReturnItem,
     assert_allclose,
     assert_array_equal,
+    assert_chunks_equal,
     assert_equal,
     assert_identical,
     has_dask,
@@ -410,6 +411,19 @@ def test_constructor_from_self_described(self):
         actual = DataArray(IndexVariable("foo", ["a", "b"]))
         assert_identical(expected, actual)
 
+    @requires_dask
+    def test_constructor_from_self_described_chunked(self):
+        expected = DataArray(
+            [[-0.1, 21], [0, 2]],
+            coords={"x": ["a", "b"], "y": [-1, -2]},
+            dims=["x", "y"],
+            name="foobar",
+            attrs={"bar": 2},
+        ).chunk()
+        actual = DataArray(expected)
+        assert_identical(expected, actual)
+        assert_chunks_equal(expected, actual)
+
     def test_constructor_from_0d(self):
         expected = Dataset({None: ([], 0)})[None]
         actual = DataArray(0)