Use ._data in Variable._replace (#10969)

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

Author: dcherian

xarray/coding/times.py

@@ -1379,9 +1379,9 @@ def __init__(
         self.time_unit = time_unit
 
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
-        if np.issubdtype(
-            variable.data.dtype, np.datetime64
-        ) or contains_cftime_datetimes(variable):
+        if np.issubdtype(variable.dtype, np.datetime64) or contains_cftime_datetimes(
+            variable
+        ):
             dims, data, attrs, encoding = unpack_for_encoding(variable)
 
             units = encoding.pop("units", None)
@@ -1499,7 +1499,7 @@ def __init__(
         self._emit_decode_timedelta_future_warning = False
 
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
-        if np.issubdtype(variable.data.dtype, np.timedelta64):
+        if np.issubdtype(variable.dtype, np.timedelta64):
             dims, data, attrs, encoding = unpack_for_encoding(variable)
             dtype = encoding.get("dtype", None)
             units = encoding.pop("units", None)

xarray/core/variable.py

@@ -972,7 +972,7 @@ def _replace(
         if dims is _default:
             dims = copy.copy(self._dims)
         if data is _default:
-            data = copy.copy(self.data)
+            data = copy.copy(self._data)
         if attrs is _default:
             attrs = copy.copy(self._attrs)
         if encoding is _default:
@@ -1228,7 +1228,7 @@ def _pad_options_dim_to_index(
         if fill_with_shape:
             return [
                 pad_option.get(d, (n, n))
-                for d, n in zip(self.dims, self.data.shape, strict=True)
+                for d, n in zip(self.dims, self.shape, strict=True)
             ]
         return [pad_option.get(d, (0, 0)) for d in self.dims]
 
@@ -1304,7 +1304,7 @@ def pad(
 
         # workaround for bug in Dask's default value of stat_length https://github.com/dask/dask/issues/5303
         if stat_length is None and mode in ["maximum", "mean", "median", "minimum"]:
-            stat_length = [(n, n) for n in self.data.shape]  # type: ignore[assignment]
+            stat_length = [(n, n) for n in self.shape]  # type: ignore[assignment]
 
         pad_width_by_index = self._pad_options_dim_to_index(pad_width)
 
@@ -1469,14 +1469,15 @@ def set_dims(self, dim, shape=None):
         if self.dims == expanded_dims:
             # don't use broadcast_to unless necessary so the result remains
             # writeable if possible
-            expanded_data = self.data
+            expanded_data = self._data
         elif shape is None or all(
             s == 1 for s, e in zip(shape, dim, strict=True) if e not in self_dims
         ):
             # "Trivial" broadcasting, i.e. simply inserting a new dimension
             # This is typically easier for duck arrays to implement
             # than the full "broadcast_to" semantics
             indexer = (None,) * (len(expanded_dims) - self.ndim) + (...,)
+            # TODO: switch this to ._data once we teach ExplicitlyIndexed arrays to handle indexers with None.
             expanded_data = self.data[indexer]
         else:  # elif shape is not None:
             dims_map = dict(zip(dim, shape, strict=True))

xarray/tests/__init__.py

@@ -31,6 +31,7 @@
     DuckArrayWrapper,
     FirstElementAccessibleArray,
     InaccessibleArray,
+    IndexableArray,
     UnexpectedDataAccess,
 )
 

xarray/tests/arrays.py

@@ -41,6 +41,16 @@ def __getitem__(self, key):
         return self.array[tuple_idxr]
 
 
+class IndexableArray(InaccessibleArray):
+    """An InaccessibleArray subclass that supports indexing."""
+
+    def __getitem__(self, key):
+        return type(self)(self.array[key])
+
+    def transpose(self, axes):
+        return type(self)(self.array.transpose(axes))
+
+
 class DuckArrayWrapper(utils.NDArrayMixin):
     """Array-like that prevents casting to array.
     Modeled after cupy."""

xarray/tests/test_variable.py

@@ -32,6 +32,7 @@
 from xarray.core.variable import as_compatible_data, as_variable
 from xarray.namedarray.pycompat import array_type
 from xarray.tests import (
+    IndexableArray,
     assert_allclose,
     assert_array_equal,
     assert_equal,
@@ -3255,3 +3256,15 @@ def test_timedelta_conversion(values, unit) -> None:
     dims = ["time"] if isinstance(values, np.ndarray | pd.Index) else []
     var = Variable(dims, values)
     assert var.dtype == np.dtype(f"timedelta64[{unit}]")
+
+
+def test_explicitly_indexed_array_preserved() -> None:
+    """Test that methods using ._data preserve ExplicitlyIndexed arrays.
+
+    Regression test for methods that should use ._data instead of .data
+    to avoid loading lazy arrays into memory.
+    """
+    arr = IndexableArray(np.array([1, 2, 3]))
+    var = Variable(["x"], arr)
+    result = var.drop_encoding()
+    assert isinstance(result._data, indexing.ExplicitlyIndexed)