Refactor concat / combine / merge into xarray/structure (#10134)

* Refactor concat / combine / merge into `xarray/structure`

Author: max-sixty

pyproject.toml

@@ -169,7 +169,7 @@ check_untyped_defs = true
 module = [
   "xarray.core.accessor_dt",
   "xarray.core.accessor_str",
-  "xarray.core.alignment",
+  "xarray.structure.alignment",
   "xarray.computation.*",
   "xarray.indexes.*",
   "xarray.tests.*",

xarray/__init__.py

@@ -28,10 +28,7 @@
     where,
 )
 from xarray.conventions import SerializationWarning, decode_cf
-from xarray.core.alignment import align, broadcast
-from xarray.core.combine import combine_by_coords, combine_nested
 from xarray.core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from xarray.core.concat import concat
 from xarray.core.coordinates import Coordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
@@ -44,7 +41,6 @@
 )
 from xarray.core.indexes import Index
 from xarray.core.indexing import IndexSelResult
-from xarray.core.merge import Context, MergeError, merge
 from xarray.core.options import get_options, set_options
 from xarray.core.parallel import map_blocks
 from xarray.core.treenode import (
@@ -55,6 +51,10 @@
 )
 from xarray.core.variable import IndexVariable, Variable, as_variable
 from xarray.namedarray.core import NamedArray
+from xarray.structure.alignment import align, broadcast
+from xarray.structure.combine import combine_by_coords, combine_nested
+from xarray.structure.concat import concat
+from xarray.structure.merge import Context, MergeError, merge
 from xarray.util.print_versions import show_versions
 
 try:

xarray/backends/api.py

@@ -36,11 +36,6 @@
 from xarray.coders import CFDatetimeCoder, CFTimedeltaCoder
 from xarray.core import indexing
 from xarray.core.chunk import _get_chunk, _maybe_chunk
-from xarray.core.combine import (
-    _infer_concat_order_from_positions,
-    _nested_combine,
-    combine_by_coords,
-)
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
 from xarray.core.datatree import DataTree
@@ -50,6 +45,11 @@
 from xarray.core.utils import is_remote_uri
 from xarray.namedarray.daskmanager import DaskManager
 from xarray.namedarray.parallelcompat import guess_chunkmanager
+from xarray.structure.combine import (
+    _infer_concat_order_from_positions,
+    _nested_combine,
+    combine_by_coords,
+)
 
 if TYPE_CHECKING:
     try:

xarray/computation/apply_ufunc.py

@@ -27,10 +27,8 @@
 _V = TypeVar("_V", bound=Union["Dataset", "DataArray"])
 
 from xarray.core import duck_array_ops, utils
-from xarray.core.alignment import deep_align
 from xarray.core.formatting import limit_lines
 from xarray.core.indexes import Index, filter_indexes_from_coords
-from xarray.core.merge import merge_attrs, merge_coordinates_without_align
 from xarray.core.options import _get_keep_attrs
 from xarray.core.utils import (
     is_dict_like,
@@ -39,6 +37,8 @@
 from xarray.core.variable import Variable
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
+from xarray.structure.alignment import deep_align
+from xarray.structure.merge import merge_attrs, merge_coordinates_without_align
 
 if TYPE_CHECKING:
     from xarray.core.coordinates import Coordinates

xarray/computation/computation.py

@@ -20,7 +20,6 @@
 from xarray.compat.array_api_compat import to_like_array
 from xarray.computation.apply_ufunc import apply_ufunc
 from xarray.core import dtypes, duck_array_ops, utils
-from xarray.core.alignment import align
 from xarray.core.common import zeros_like
 from xarray.core.duck_array_ops import datetime_to_numeric
 from xarray.core.options import OPTIONS, _get_keep_attrs
@@ -32,6 +31,7 @@
 from xarray.core.variable import Variable
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
+from xarray.structure.alignment import align
 from xarray.util.deprecation_helpers import deprecate_dims
 
 if TYPE_CHECKING:

xarray/computation/weighted.py

@@ -9,9 +9,9 @@
 from xarray.computation.apply_ufunc import apply_ufunc
 from xarray.computation.computation import dot
 from xarray.core import duck_array_ops, utils
-from xarray.core.alignment import align, broadcast
 from xarray.core.types import Dims, T_DataArray, T_Xarray
 from xarray.namedarray.utils import is_duck_dask_array
+from xarray.structure.alignment import align, broadcast
 
 # Weighted quantile methods are a subset of the numpy supported quantile methods.
 QUANTILE_METHODS = Literal[

xarray/core/common.py

@@ -1216,9 +1216,9 @@ def where(self, cond: Any, other: Any = dtypes.NA, drop: bool = False) -> Self:
         numpy.where : corresponding numpy function
         where : equivalent function
         """
-        from xarray.core.alignment import align
         from xarray.core.dataarray import DataArray
         from xarray.core.dataset import Dataset
+        from xarray.structure.alignment import align
 
         if callable(cond):
             cond = cond(self)

xarray/core/coordinates.py

@@ -13,7 +13,6 @@
 import pandas as pd
 
 from xarray.core import formatting
-from xarray.core.alignment import Aligner
 from xarray.core.indexes import (
     Index,
     Indexes,
@@ -22,7 +21,6 @@
     assert_no_index_corrupted,
     create_default_index_implicit,
 )
-from xarray.core.merge import merge_coordinates_without_align, merge_coords
 from xarray.core.types import DataVars, Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
@@ -31,6 +29,8 @@
     emit_user_level_warning,
 )
 from xarray.core.variable import Variable, as_variable, calculate_dimensions
+from xarray.structure.alignment import Aligner
+from xarray.structure.merge import merge_coordinates_without_align, merge_coords
 
 if TYPE_CHECKING:
     from xarray.core.common import DataWithCoords

xarray/core/dataarray.py

@@ -32,15 +32,10 @@
 from xarray.computation import computation, ops
 from xarray.computation.arithmetic import DataArrayArithmetic
 from xarray.computation.computation import unify_chunks
-from xarray.core import alignment, dtypes, indexing, utils
+from xarray.core import dtypes, indexing, utils
 from xarray.core._aggregations import DataArrayAggregations
 from xarray.core.accessor_dt import CombinedDatetimelikeAccessor
 from xarray.core.accessor_str import StringAccessor
-from xarray.core.alignment import (
-    _broadcast_helper,
-    _get_broadcast_dims_map_common_coords,
-    align,
-)
 from xarray.core.common import AbstractArray, DataWithCoords, get_chunksizes
 from xarray.core.coordinates import (
     Coordinates,
@@ -59,7 +54,6 @@
     isel_indexes,
 )
 from xarray.core.indexing import is_fancy_indexer, map_index_queries
-from xarray.core.merge import PANDAS_TYPES, MergeError
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import (
     Bins,
@@ -88,6 +82,13 @@
 )
 from xarray.plot.accessor import DataArrayPlotAccessor
 from xarray.plot.utils import _get_units_from_attrs
+from xarray.structure import alignment
+from xarray.structure.alignment import (
+    _broadcast_helper,
+    _get_broadcast_dims_map_common_coords,
+    align,
+)
+from xarray.structure.merge import PANDAS_TYPES, MergeError
 from xarray.util.deprecation_helpers import _deprecate_positional_args, deprecate_dims
 
 if TYPE_CHECKING:

xarray/core/dataset.py

@@ -45,20 +45,14 @@
 from xarray.computation import ops
 from xarray.computation.arithmetic import DatasetArithmetic
 from xarray.computation.computation import _ensure_numeric, unify_chunks
+from xarray.core import dtypes as xrdtypes
 from xarray.core import (
-    alignment,
     duck_array_ops,
     formatting,
     formatting_html,
     utils,
 )
-from xarray.core import dtypes as xrdtypes
 from xarray.core._aggregations import DatasetAggregations
-from xarray.core.alignment import (
-    _broadcast_helper,
-    _get_broadcast_dims_map_common_coords,
-    align,
-)
 from xarray.core.common import (
     DataWithCoords,
     _contains_datetime_like_objects,
@@ -83,12 +77,6 @@
     roll_indexes,
 )
 from xarray.core.indexing import is_fancy_indexer, map_index_queries
-from xarray.core.merge import (
-    dataset_merge_method,
-    dataset_update_method,
-    merge_coordinates_without_align,
-    merge_data_and_coords,
-)
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import (
     Bins,
@@ -131,6 +119,18 @@
 from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
 from xarray.namedarray.pycompat import array_type, is_chunked_array, to_numpy
 from xarray.plot.accessor import DatasetPlotAccessor
+from xarray.structure import alignment
+from xarray.structure.alignment import (
+    _broadcast_helper,
+    _get_broadcast_dims_map_common_coords,
+    align,
+)
+from xarray.structure.merge import (
+    dataset_merge_method,
+    dataset_update_method,
+    merge_coordinates_without_align,
+    merge_data_and_coords,
+)
 from xarray.util.deprecation_helpers import _deprecate_positional_args, deprecate_dims
 
 if TYPE_CHECKING:
@@ -144,7 +144,6 @@
     from xarray.computation.weighted import DatasetWeighted
     from xarray.core.dataarray import DataArray
     from xarray.core.groupby import DatasetGroupBy
-    from xarray.core.merge import CoercibleMapping, CoercibleValue
     from xarray.core.resample import DatasetResample
     from xarray.core.types import (
         CFCalendar,
@@ -176,6 +175,7 @@
     )
     from xarray.groupers import Grouper, Resampler
     from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
+    from xarray.structure.merge import CoercibleMapping, CoercibleValue
 
 
 # list of attributes of pd.DatetimeIndex that are ndarrays of time info
@@ -1402,8 +1402,8 @@ def _setitem_check(self, key, value):
         When assigning values to a subset of a Dataset, do consistency check beforehand
         to avoid leaving the dataset in a partially updated state when an error occurs.
         """
-        from xarray.core.alignment import align
         from xarray.core.dataarray import DataArray
+        from xarray.structure.alignment import align
 
         if isinstance(value, Dataset):
             missing_vars = [
@@ -5297,7 +5297,7 @@ def to_stacked_array(
         Dimensions without coordinates: x
 
         """
-        from xarray.core.concat import concat
+        from xarray.structure.concat import concat
 
         stacking_dims = tuple(dim for dim in self.dims if dim not in sample_dims)
 
@@ -9784,8 +9784,8 @@ def curvefit(
         from scipy.optimize import curve_fit
 
         from xarray.computation.computation import apply_ufunc
-        from xarray.core.alignment import broadcast
         from xarray.core.dataarray import _THIS_ARRAY, DataArray
+        from xarray.structure.alignment import broadcast
 
         if p0 is None:
             p0 = {}

xarray/core/datatree.py

@@ -27,7 +27,6 @@
 from xarray.core import utils
 from xarray.core._aggregations import DataTreeAggregations
 from xarray.core._typed_ops import DataTreeOpsMixin
-from xarray.core.alignment import align
 from xarray.core.common import TreeAttrAccessMixin, get_chunksizes
 from xarray.core.coordinates import Coordinates, DataTreeCoordinates
 from xarray.core.dataarray import DataArray
@@ -45,7 +44,6 @@
     datatree_repr as datatree_repr_html,
 )
 from xarray.core.indexes import Index, Indexes
-from xarray.core.merge import dataset_update_method
 from xarray.core.options import OPTIONS as XR_OPTS
 from xarray.core.treenode import NamedNode, NodePath, zip_subtrees
 from xarray.core.types import Self
@@ -62,6 +60,8 @@
 from xarray.core.variable import Variable
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
+from xarray.structure.alignment import align
+from xarray.structure.merge import dataset_update_method
 
 try:
     from xarray.core.variable import calculate_dimensions
@@ -74,7 +74,6 @@
     import pandas as pd
 
     from xarray.core.datatree_io import T_DataTreeNetcdfEngine, T_DataTreeNetcdfTypes
-    from xarray.core.merge import CoercibleMapping, CoercibleValue
     from xarray.core.types import (
         Dims,
         DtCompatible,
@@ -86,6 +85,7 @@
         ZarrWriteModes,
     )
     from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
+    from xarray.structure.merge import CoercibleMapping, CoercibleValue
 
 # """
 # DEVELOPERS' NOTE

xarray/core/groupby.py

@@ -22,17 +22,14 @@
     DataArrayGroupByAggregations,
     DatasetGroupByAggregations,
 )
-from xarray.core.alignment import align, broadcast
 from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
-from xarray.core.concat import concat
 from xarray.core.coordinates import Coordinates, _coordinates_from_variable
 from xarray.core.duck_array_ops import where
 from xarray.core.formatting import format_array_flat
 from xarray.core.indexes import (
     PandasMultiIndex,
     filter_indexes_from_coords,
 )
-from xarray.core.merge import merge_coords
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import (
     Dims,
@@ -54,6 +51,9 @@
 )
 from xarray.core.variable import IndexVariable, Variable
 from xarray.namedarray.pycompat import is_chunked_array
+from xarray.structure.alignment import align, broadcast
+from xarray.structure.concat import concat
+from xarray.structure.merge import merge_coords
 
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike

xarray/core/parallel.py

@@ -8,14 +8,14 @@
 
 import numpy as np
 
-from xarray.core.alignment import align
 from xarray.core.coordinates import Coordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
 from xarray.core.indexes import Index
-from xarray.core.merge import merge
 from xarray.core.utils import is_dask_collection
 from xarray.core.variable import Variable
+from xarray.structure.alignment import align
+from xarray.structure.merge import merge
 
 if TYPE_CHECKING:
     from xarray.core.types import T_Xarray