Avoid auto creation of indexes in concat

xarray/core/concat.py

@@ -8,6 +8,7 @@
 
 from xarray.core import dtypes, utils
 from xarray.core.alignment import align, reindex_variables
+from xarray.core.coordinates import Coordinates
 from xarray.core.duck_array_ops import lazy_array_equiv
 from xarray.core.indexes import Index, PandasIndex
 from xarray.core.merge import (
@@ -646,14 +647,26 @@ def get_indexes(name):
             # preserves original variable order
             result_vars[name] = result_vars.pop(name)
 
-    result = type(datasets[0])(result_vars, attrs=result_attrs)
-
-    absent_coord_names = coord_names - set(result.variables)
+    absent_coord_names = coord_names - set(result_vars)
     if absent_coord_names:
         raise ValueError(
             f"Variables {absent_coord_names!r} are coordinates in some datasets but not others."
         )
-    result = result.set_coords(coord_names)
+    coord_vars = {
+        name: result_var
+        for name, result_var in result_vars.items()
+        if name in coord_names
+    }
+    coords = Coordinates(coord_vars, indexes=result_indexes)
+
+    # TODO: this is just the complement of the set of coord_vars
+    result_data_vars = {
+        name: result_var
+        for name, result_var in result_vars.items()
+        if name not in coord_names
+    }
+
+    result = type(datasets[0])(result_data_vars, coords=coords, attrs=result_attrs)
     result.encoding = result_encoding
 
     result = result.drop_vars(unlabeled_dims, errors="ignore")
@@ -665,10 +678,6 @@ def get_indexes(name):
         else:
             index_vars = index.create_variables()
         result[dim] = index_vars[dim]
-        result_indexes[dim] = index
-
-    # TODO: add indexes at Dataset creation (when it is supported)
-    result = result._overwrite_indexes(result_indexes)
 
     return result
 

xarray/core/coordinates.py

@@ -22,7 +22,11 @@
     assert_no_index_corrupted,
     create_default_index_implicit,
 )
-from xarray.core.merge import merge_coordinates_without_align, merge_coords
+from xarray.core.merge import (
+    assert_valid_explicit_coords,
+    merge_coordinates_without_align,
+    merge_coords,
+)
 from xarray.core.types import DataVars, Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
@@ -306,6 +310,11 @@ def __init__(
                 else:
                     variables[name] = var
 
+            dims = set(d for var in variables.values() for d in var.dims)
+            assert_valid_explicit_coords(
+                variables, dims=dims, explicit_coords=list(variables)
+            )
+
         if indexes is None:
             indexes = {}
         else:

xarray/core/merge.py

@@ -564,7 +564,7 @@ def merge_coords(
 
 def assert_valid_explicit_coords(
     variables: Mapping[Any, Any],
-    dims: Mapping[Any, int],
+    dims: Iterable[Any],
     explicit_coords: Iterable[Hashable],
 ) -> None:
     """Validate explicit coordinate names/dims.
@@ -721,16 +721,16 @@ def merge_core(
         collected, prioritized, compat=compat, combine_attrs=combine_attrs
     )
 
-    dims = calculate_dimensions(variables)
+    sizes = calculate_dimensions(variables)
 
     coord_names, noncoord_names = determine_coords(coerced)
     if compat == "minimal":
         # coordinates may be dropped in merged results
         coord_names.intersection_update(variables)
     if explicit_coords is not None:
-        assert_valid_explicit_coords(variables, dims, explicit_coords)
+        assert_valid_explicit_coords(variables, sizes.keys(), explicit_coords)
         coord_names.update(explicit_coords)
-    for dim, size in dims.items():
+    for dim, size in sizes.items():
         if dim in variables:
             coord_names.add(dim)
     ambiguous_coords = coord_names.intersection(noncoord_names)
@@ -745,7 +745,7 @@ def merge_core(
         combine_attrs,
     )
 
-    return _MergeResult(variables, coord_names, dims, out_indexes, attrs)
+    return _MergeResult(variables, coord_names, sizes, out_indexes, attrs)
 
 
 def merge(

xarray/tests/__init__.py

@@ -4,7 +4,9 @@
 import platform
 import string
 import warnings
+from collections.abc import Iterable
 from contextlib import contextmanager, nullcontext
+from typing import Any, Callable
 from unittest import mock  # noqa: F401
 
 import numpy as np
@@ -208,6 +210,105 @@ def __getitem__(self, key):
         raise UnexpectedDataAccess("Tried accessing data.")
 
 
+HANDLED_ARRAY_FUNCTIONS: dict[str, Callable] = {}
+
+
+def implements(numpy_function):
+    """Register an __array_function__ implementation for ManifestArray objects."""
+
+    def decorator(func):
+        HANDLED_ARRAY_FUNCTIONS[numpy_function] = func
+        return func
+
+    return decorator
+
+
+@implements(np.concatenate)
+def concatenate(
+    arrays: Iterable[ConcatenatableArray], /, *, axis=0
+) -> ConcatenatableArray:
+    if any(not isinstance(arr, ConcatenatableArray) for arr in arrays):
+        raise TypeError
+
+    result = np.concatenate([arr.array for arr in arrays], axis=axis)
+    return ConcatenatableArray(result)
+
+
+@implements(np.stack)
+def stack(arrays: Iterable[ConcatenatableArray], /, *, axis=0) -> ConcatenatableArray:
+    if any(not isinstance(arr, ConcatenatableArray) for arr in arrays):
+        raise TypeError
+
+    result = np.stack([arr.array for arr in arrays], axis=axis)
+    return ConcatenatableArray(result)
+
+
+@implements(np.result_type)
+def result_type(*arrays_and_dtypes) -> np.dtype:
+    """Called by xarray to ensure all arguments to concat have the same dtype."""
+    first_dtype, *other_dtypes = (np.dtype(obj) for obj in arrays_and_dtypes)
+    for other_dtype in other_dtypes:
+        if other_dtype != first_dtype:
+            raise ValueError("dtypes not all consistent")
+    return first_dtype
+
+
+@implements(np.broadcast_to)
+def broadcast_to(
+    x: ConcatenatableArray, /, shape: tuple[int, ...]
+) -> ConcatenatableArray:
+    """
+    Broadcasts an array to a specified shape, by either manipulating chunk keys or copying chunk manifest entries.
+    """
+    if not isinstance(x, ConcatenatableArray):
+        raise TypeError
+
+    result = np.broadcast_to(x.array, shape=shape)
+    return ConcatenatableArray(result)
+
+
+class ConcatenatableArray(utils.NDArrayMixin):
+    """Disallows loading or coercing to an index but does support concatenation / stacking."""
+
+    # TODO only reason this is different from InaccessibleArray is to avoid it being a subclass of ExplicitlyIndexed
+
+    HANDLED_ARRAY_FUNCTIONS = [concatenate, stack, result_type]
+
+    def __init__(self, array):
+        self.array = array
+
+    def get_duck_array(self):
+        raise UnexpectedDataAccess("Tried accessing data")
+
+    def __array__(self, dtype: np.typing.DTypeLike = None):
+        raise UnexpectedDataAccess("Tried accessing data")
+
+    def __getitem__(self, key):
+        raise UnexpectedDataAccess("Tried accessing data.")
+
+    def __array_function__(self, func, types, args, kwargs) -> Any:
+        if func not in HANDLED_ARRAY_FUNCTIONS:
+            return NotImplemented
+
+        # Note: this allows subclasses that don't override
+        # __array_function__ to handle ManifestArray objects
+        if not all(issubclass(t, ConcatenatableArray) for t in types):
+            return NotImplemented
+
+        return HANDLED_ARRAY_FUNCTIONS[func](*args, **kwargs)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs) -> Any:
+        """We have to define this in order to convince xarray that this class is a duckarray, even though we will never support ufuncs."""
+        return NotImplemented
+
+    def astype(self, dtype: np.dtype, /, *, copy: bool = True) -> ConcatenatableArray:
+        """Needed because xarray will call this even when it's a no-op"""
+        if dtype != self.dtype:
+            raise NotImplementedError()
+        else:
+            return self
+
+
 class FirstElementAccessibleArray(InaccessibleArray):
     def __getitem__(self, key):
         tuple_idxr = key.tuple

xarray/tests/test_concat.py

@@ -12,6 +12,7 @@
 from xarray.core.coordinates import Coordinates
 from xarray.core.indexes import PandasIndex
 from xarray.tests import (
+    ConcatenatableArray,
     InaccessibleArray,
     assert_array_equal,
     assert_equal,
@@ -978,6 +979,34 @@ def test_concat_str_dtype(self, dtype, dim) -> None:
 
         assert np.issubdtype(actual.x2.dtype, dtype)
 
+    def test_concat_avoids_index_auto_creation(self) -> None:
+        # TODO once passing indexes={} directly to Dataset constructor is allowed then no need to create coords first
+        coords = Coordinates(
+            {"x": ConcatenatableArray(np.array([1, 2, 3]))}, indexes={}
+        )
+        datasets = [
+            Dataset(
+                {"a": (["x", "y"], ConcatenatableArray(np.zeros((3, 3))))},
+                coords=coords,
+            )
+            for _ in range(2)
+        ]
+        # should not raise on concat
+        combined = concat(datasets, dim="x")
+        assert combined["a"].shape == (6, 3)
+        assert combined["a"].dims == ("x", "y")
+
+        # nor have auto-created any indexes
+        assert combined.indexes == {}
+
+        # should not raise on stack
+        combined = concat(datasets, dim="z")
+        assert combined["a"].shape == (2, 3, 3)
+        assert combined["a"].dims == ("z", "x", "y")
+
+        # nor have auto-created any indexes
+        assert combined.indexes == {}
+
 
 class TestConcatDataArray:
     def test_concat(self) -> None:
@@ -1051,6 +1080,35 @@ def test_concat_lazy(self) -> None:
         assert combined.shape == (2, 3, 3)
         assert combined.dims == ("z", "x", "y")
 
+    def test_concat_avoids_index_auto_creation(self) -> None:
+        # TODO once passing indexes={} directly to DataArray constructor is allowed then no need to create coords first
+        coords = Coordinates(
+            {"x": ConcatenatableArray(np.array([1, 2, 3]))}, indexes={}
+        )
+        arrays = [
+            DataArray(
+                ConcatenatableArray(np.zeros((3, 3))),
+                dims=["x", "y"],
+                coords=coords,
+            )
+            for _ in range(2)
+        ]
+        # should not raise on concat
+        combined = concat(arrays, dim="x")
+        assert combined.shape == (6, 3)
+        assert combined.dims == ("x", "y")
+
+        # nor have auto-created any indexes
+        assert combined.indexes == {}
+
+        # should not raise on stack
+        combined = concat(arrays, dim="z")
+        assert combined.shape == (2, 3, 3)
+        assert combined.dims == ("z", "x", "y")
+
+        # nor have auto-created any indexes
+        assert combined.indexes == {}
+
     @pytest.mark.parametrize("fill_value", [dtypes.NA, 2, 2.0])
     def test_concat_fill_value(self, fill_value) -> None:
         foo = DataArray([1, 2], coords=[("x", [1, 2])])

xarray/tests/test_coordinates.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -8,7 +9,7 @@
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
 from xarray.core.indexes import PandasIndex, PandasMultiIndex
-from xarray.core.variable import IndexVariable
+from xarray.core.variable import IndexVariable, Variable
 from xarray.tests import assert_identical, source_ndarray
 
 
@@ -69,6 +70,12 @@ def test_init_dim_sizes_conflict(self) -> None:
         with pytest.raises(ValueError):
             Coordinates(coords={"foo": ("x", [1, 2]), "bar": ("x", [1, 2, 3, 4])})
 
+    def test_init_var_shares_name_with_dim(self) -> None:
+        # regression test for GH #8883
+        var = Variable(data=np.arange(6).reshape(2, 3), dims=["x", "y"])
+        with pytest.raises(ValueError):
+            Coordinates(coords={"x": var}, indexes={})
+
     def test_from_pandas_multiindex(self) -> None:
         midx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("one", "two"))
         coords = Coordinates.from_pandas_multiindex(midx, "x")