allow chunking to drop non-dimensional coords (#532)

d-chambers · web-flow · commit 4242dc7ba687 · 2025-09-10T08:11:06.000+01:00
* allow chunking to drop non-dimensional coords

* rabbit review
diff --git a/dascore/utils/coordmanager.py b/dascore/utils/coordmanager.py
@@ -18,9 +18,10 @@ def merge_coord_managers(
     coord_managers: Sequence[dc.CoordManager],
     dim: str,
     snap_tolerance: float | None = None,
+    drop_conflicting: bool = False,
 ) -> dc.CoordManager:
     """
-    Merger coordinate managers along a specified dimension.
+    Merge coordinate managers along a specified dimension.
 
     Parameters
     ----------
@@ -34,6 +35,9 @@ def merge_coord_managers(
         start/end to be joined together. If they don't meet this requirement
         an [CoordMergeError](`dascore.exceptions.CoordMergeError`) is raised.
         If None, no checks are performed.
+    drop_conflicting
+        If True, drop conflicting (non-dimensional) coordinates, otherwise
+        raise an exception if they occur.
     """
 
     def _get_dims(managers):
@@ -68,9 +72,15 @@ def _get_non_merge_coords(managers, non_merger_names):
                 dims = managers[0].dim_map[coord_name]
                 out[coord_name] = (dims, first)
                 continue
+            # Simply skip conflicting
+            elif drop_conflicting:
+                # These are non dimensional coords
+                if not any(coord_name in x.dims for x in managers):
+                    continue
             msg = (
                 f"Non merging coordinates {coord_name} are not equal. "
-                "Coordinate managers cannot be merged."
+                "Coordinate managers cannot be merged. Try using "
+                "spool.chunk with conflict='drop'."
             )
             raise CoordMergeError(msg)
         return out
diff --git a/dascore/utils/patch.py b/dascore/utils/patch.py
@@ -392,9 +392,11 @@ def _maybe_step(df, dim):
             return get_middle_value(col)
         return None
 
-    def _get_new_coord(df, merge_dim, coords):
+    def _get_new_coord(df, merge_dim, coords, drop_conflicting=False):
         """Get new coordinates, also validate anticipated sampling."""
-        new_coord = merge_coord_managers(coords, dim=merge_dim)
+        new_coord = merge_coord_managers(
+            coords, dim=merge_dim, drop_conflicting=drop_conflicting
+        )
         expected_step = _maybe_step(df, merge_dim)
         if not pd.isnull(expected_step):
             new_coord = new_coord.snap(merge_dim)[0]
@@ -416,7 +418,10 @@ def _get_new_coord(df, merge_dim, coords):
     coords = [x.coords for x in patches]
     attrs = [x.attrs for x in patches]
     new_data = np.concatenate(datas, axis=axis)
-    new_coord = _get_new_coord(df, merge_dim, coords)
+    # Determine if conflicting non-dimensional coords should be dropped.
+    conf = merge_kwargs.get("conflicts", None)
+    drop_conf_coords = True if conf in {"drop", "keep_first"} else False
+    new_coord = _get_new_coord(df, merge_dim, coords, drop_conf_coords)
     coord = new_coord.coord_map[merge_dim] if merge_dim in dims else None
     new_attrs = combine_patch_attrs(attrs, merge_dim, coord=coord, **merge_kwargs)
     patch = dc.Patch(data=new_data, coords=new_coord, attrs=new_attrs, dims=dims)
diff --git a/tests/test_core/test_patch_chunk.py b/tests/test_core/test_patch_chunk.py
@@ -270,6 +270,21 @@ def adjacent_spool_different_attrs(self, adjacent_spool_no_overlap):
         # since
         return dc.spool(out)
 
+    @pytest.fixture(scope="class")
+    def patches_conflicting_private_coord(self, random_patch):
+        """Create two patches that have conflicting private coords."""
+        dist_ax = random_patch.dims.index("distance")
+        rand = np.random.RandomState(42)
+        c1 = rand.random(random_patch.shape[dist_ax])
+        c2 = rand.random(c1.shape)
+
+        time = random_patch.get_coord("time")
+        p1 = random_patch.update_coords(_bad_coord=("distance", c1))
+        p2 = random_patch.update_coords(
+            _bad_coord=("distance", c2), time=time + time.coord_range()
+        )
+        return p1, p2
+
     def test_merge_unequal_other(self, distance_adjacent):
         """When distance values are not equal time shouldn't be merge-able."""
         with pytest.raises(CoordMergeError):
@@ -462,3 +477,24 @@ def test_chunk_patches_with_non_coord(self, random_patch):
         chunked = spool.chunk(time=None)
         # Since the time dims are NaN, this can't work.
         assert not len(chunked)
+
+    def test_merge_with_conflicting_private_coords(
+        self,
+        patches_conflicting_private_coord,
+    ):
+        """
+        Private coords that conflict should be dropped and not block merge
+        when conflict="drop".
+
+        Otherwise they should raise.
+        """
+        p1, p2 = patches_conflicting_private_coord
+        merged_spool = dc.spool([p1, p2]).chunk(time=None, conflict="drop")
+        merge_patch = merged_spool[0]
+        assert len(merged_spool) == 1
+        # Since the private coords conflicted, they should have been dropped.
+        coord_names = list(merge_patch.coords.coord_map)
+        assert not any([x.startswith("_") for x in coord_names])
+        # Without conflict drop this should raise.
+        with pytest.raises(CoordMergeError, match="conflict"):
+            dc.spool([p1, p2]).chunk(time=None)[0]
diff --git a/tests/test_utils/test_coordmanager_utils.py b/tests/test_utils/test_coordmanager_utils.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 from dascore.core.coords import CoordArray, CoordMonotonicArray, CoordRange
@@ -23,6 +24,21 @@ def _get_offset_coord_manager(self, cm, from_max=True, **kwargs):
         new, _ = cm.update_from_attrs({attr_name: start + value})
         return new
 
+    @pytest.fixture(scope="class")
+    def conflicting_non_dim_coords(self, cm_basic):
+        """Get two coord managers with conflicting non-dimensional coordinates."""
+        dist_ax = cm_basic.dims.index("distance")
+        rand = np.random.RandomState(42)
+        c1 = rand.random(cm_basic.shape[dist_ax])
+        c2 = rand.random(c1.shape)
+
+        time = cm_basic.get_coord("time")
+        cm1 = cm_basic.update_coords(_bad_coord=("distance", c1))
+        cm2 = cm1.update_coords(
+            _bad_coord=("distance", c2), time=time + time.coord_range()
+        )
+        return cm1, cm2
+
     def test_merge_simple(self, cm_basic):
         """Ensure we can merge simple, contiguous, coordinates together."""
         cm1 = cm_basic
@@ -136,3 +152,16 @@ def test_slightly_different_dt(self, cm_dt_small_diff):
         cm = cm_dt_small_diff
         coord = cm.coord_map["time"]
         assert coord.sorted
+
+    def test_conflicting_non_dimensional_coords(self, conflicting_non_dim_coords):
+        """
+        Ensure conflicting non-dimensional coords can be merged if drop_conflict=True,
+        Otherwise raise.
+        """
+        c1, c2 = conflicting_non_dim_coords
+
+        out = merge_coord_managers([c1, c2], dim="time", drop_conflicting=True)
+        assert not any([x.startswith("_") for x in out.coord_map])
+
+        with pytest.raises(CoordMergeError, match="cannot be merged"):
+            merge_coord_managers([c1, c2], dim="time", drop_conflicting=False)
