Fix assert_equal with check_dim_order=False for mixed dimension orders (#10718)

* Fix assert_equal with check_dim_order=False for Datasets with mixed dimension orders

Fixes #10704

The bug: assert_equal with check_dim_order=False was failing when comparing
Datasets containing variables with different dimension orders. It would even
fail when comparing a Dataset to itself.

The fix: Transpose both objects to a canonical dimension order using the
intersection of their dimensions. The ellipsis (...) handles any dimensions
unique to either object, making the solution general and elegant.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Address review feedback: avoid sorting non-sortable dimension names

- Use list(common_dims) instead of sorted(common_dims) since dimensions
  only need to be hashable, not sortable
- Add test case for datasets with non-sortable dimension names (e.g., int and str)
- Transpose both a and b to the same canonical order for consistency

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add full test coverage for maybe_transpose_dims changes

- Add test for no common dimensions path
- Add test for Variable type specifically (not just DataArray)
- Now all code paths in maybe_transpose_dims are covered

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix DataTree handling in maybe_transpose_dims

The reviewer was correct - the DataTree handling was wrong. We were only
transposing b, but for consistency we need to transpose both a and b
to the same canonical order, just like we do for Dataset.

This fixes the issue and adds a comprehensive test case.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trigger CI re-run for flaky Zarr test

* Simplify DataTree handling in maybe_transpose_dims

Use single map_over_datasets call instead of two separate calls.
map_over_datasets supports tuple returns and unpacks them automatically,
making the explicit lambda extraction unnecessary.

Co-authored-by: Claude <noreply@anthropic.com>

* Fix RTD builds by restricting test-nightly to linux/macOS platforms

The `test-nightly` environment uses pandas nightly wheels from PyPI,
which currently don't have win-64 builds available. This causes
`pixi lock` to fail when solving for all platforms.

RTD builds fail because they have no lock file cache (unlike GitHub
Actions CI which caches pixi.lock). When RTD runs `pixi install -e doc`,
pixi must generate the lock file from scratch, which fails on the
unsolvable test-nightly/win-64 combination.

This restriction can be removed once pandas nightly provides win-64
wheels again.

Co-authored-by: Claude <noreply@anthropic.com>

* Remove redundant numpy imports in test functions

Co-authored-by: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Mathias Hauser <mathause@users.noreply.github.com>

Author: 4 people

xarray/testing/assertions.py

@@ -86,23 +86,33 @@ def assert_isomorphic(a: DataTree, b: DataTree):
 
 
 def maybe_transpose_dims(a, b, check_dim_order: bool):
-    """Helper for assert_equal/allclose/identical"""
+    """Helper for assert_equal/allclose/identical
+
+    Returns (a, b) tuple with dimensions transposed to canonical order if needed.
+    """
 
     __tracebackhide__ = True
 
+    if check_dim_order:
+        return a, b
+
     def _maybe_transpose_dims(a, b):
         if not isinstance(a, Variable | DataArray | Dataset):
-            return b
-        if set(a.dims) == set(b.dims):
-            # Ensure transpose won't fail if a dimension is missing
-            # If this is the case, the difference will be caught by the caller
-            return b.transpose(*a.dims)
-        return b
-
-    if check_dim_order:
-        return b
+            return a, b
+
+        # Find common dimensions and transpose both to canonical order
+        common_dims = set(a.dims) & set(b.dims)
+        if common_dims:
+            # Use order from the intersection, with ellipsis for any unique dims
+            canonical_order = list(common_dims) + [...]
+            # For Datasets, we need to transpose both to the same order
+            # For Variable/DataArray, we could just transpose b, but for consistency
+            # and simplicity we transpose both
+            return a.transpose(*canonical_order), b.transpose(*canonical_order)
+        return a, b
 
     if isinstance(a, DataTree):
+        # map_over_datasets supports tuple returns and unpacks them automatically
         return map_over_datasets(_maybe_transpose_dims, a, b)
 
     return _maybe_transpose_dims(a, b)
@@ -140,7 +150,7 @@ def assert_equal(a, b, check_dim_order: bool = True):
     assert type(a) is type(b) or (
         isinstance(a, Coordinates) and isinstance(b, Coordinates)
     )
-    b = maybe_transpose_dims(a, b, check_dim_order)
+    a, b = maybe_transpose_dims(a, b, check_dim_order)
     if isinstance(a, Variable | DataArray):
         assert a.equals(b), formatting.diff_array_repr(a, b, "equals")
     elif isinstance(a, Dataset):
@@ -228,7 +238,7 @@ def assert_allclose(
     """
     __tracebackhide__ = True
     assert type(a) is type(b)
-    b = maybe_transpose_dims(a, b, check_dim_order)
+    a, b = maybe_transpose_dims(a, b, check_dim_order)
 
     equiv = functools.partial(
         _data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes

xarray/tests/test_assertions.py

@@ -115,6 +115,26 @@ def test_assert_equal_transpose_datatree() -> None:
 
     xr.testing.assert_equal(a, b, check_dim_order=False)
 
+    # Test with mixed dimension orders in datasets (the tricky case)
+    ds_mixed = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([4, 5]), dims=("a", "b")),
+            "bar": xr.DataArray(np.ones([5, 4]), dims=("b", "a")),
+        }
+    )
+    ds_mixed2 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([5, 4]), dims=("b", "a")),
+            "bar": xr.DataArray(np.ones([4, 5]), dims=("a", "b")),
+        }
+    )
+
+    tree1 = xr.DataTree.from_dict({"node": ds_mixed})
+    tree2 = xr.DataTree.from_dict({"node": ds_mixed2})
+
+    # Should work with check_dim_order=False
+    xr.testing.assert_equal(tree1, tree2, check_dim_order=False)
+
 
 @pytest.mark.filterwarnings("error")
 @pytest.mark.parametrize(
@@ -241,6 +261,101 @@ def __array__(
         assert len(w) == 0
 
 
+def test_assert_equal_dataset_check_dim_order():
+    """Test for issue #10704 - check_dim_order=False with Datasets containing mixed dimension orders."""
+    # Dataset with variables having different dimension orders
+    dataset_1 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([4, 5]), dims=("a", "b")),
+            "bar": xr.DataArray(np.ones([5, 4]), dims=("b", "a")),
+        }
+    )
+
+    dataset_2 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([5, 4]), dims=("b", "a")),
+            "bar": xr.DataArray(np.ones([4, 5]), dims=("a", "b")),
+        }
+    )
+
+    # These should be equal when ignoring dimension order
+    xr.testing.assert_equal(dataset_1, dataset_2, check_dim_order=False)
+    xr.testing.assert_allclose(dataset_1, dataset_2, check_dim_order=False)
+
+    # Should also work when comparing dataset to itself
+    xr.testing.assert_equal(dataset_1, dataset_1, check_dim_order=False)
+    xr.testing.assert_allclose(dataset_1, dataset_1, check_dim_order=False)
+
+    # But should fail with check_dim_order=True
+    with pytest.raises(AssertionError):
+        xr.testing.assert_equal(dataset_1, dataset_2, check_dim_order=True)
+    with pytest.raises(AssertionError):
+        xr.testing.assert_allclose(dataset_1, dataset_2, check_dim_order=True)
+
+    # Test with non-sortable dimension names (int and str)
+    dataset_mixed_1 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([4, 5]), dims=(1, "b")),
+            "bar": xr.DataArray(np.ones([5, 4]), dims=("b", 1)),
+        }
+    )
+
+    dataset_mixed_2 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([5, 4]), dims=("b", 1)),
+            "bar": xr.DataArray(np.ones([4, 5]), dims=(1, "b")),
+        }
+    )
+
+    # Should work with mixed types when ignoring dimension order
+    xr.testing.assert_equal(dataset_mixed_1, dataset_mixed_2, check_dim_order=False)
+    xr.testing.assert_equal(dataset_mixed_1, dataset_mixed_1, check_dim_order=False)
+
+
+def test_assert_equal_no_common_dims():
+    """Test assert_equal when objects have no common dimensions."""
+    # DataArrays with completely different dimensions
+    da1 = xr.DataArray(np.zeros([4, 5]), dims=("x", "y"))
+    da2 = xr.DataArray(np.zeros([3, 2]), dims=("a", "b"))
+
+    # Should fail even with check_dim_order=False since dims are different
+    with pytest.raises(AssertionError):
+        xr.testing.assert_equal(da1, da2, check_dim_order=False)
+
+    # Datasets with no common dimensions
+    ds1 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([4]), dims=("x",)),
+            "bar": xr.DataArray(np.ones([5]), dims=("y",)),
+        }
+    )
+    ds2 = xr.Dataset(
+        {
+            "foo": xr.DataArray(np.zeros([3]), dims=("a",)),
+            "bar": xr.DataArray(np.ones([2]), dims=("b",)),
+        }
+    )
+
+    # Should fail since dimensions are completely different
+    with pytest.raises(AssertionError):
+        xr.testing.assert_equal(ds1, ds2, check_dim_order=False)
+
+
+def test_assert_equal_variable_transpose():
+    """Test assert_equal with transposed Variable objects."""
+    # Variables with transposed dimensions
+    var1 = xr.Variable(("x", "y"), np.zeros([4, 5]))
+    var2 = xr.Variable(("y", "x"), np.zeros([5, 4]))
+
+    # Should fail with check_dim_order=True
+    with pytest.raises(AssertionError):
+        xr.testing.assert_equal(var1, var2, check_dim_order=True)
+
+    # Should pass with check_dim_order=False
+    xr.testing.assert_equal(var1, var2, check_dim_order=False)
+    xr.testing.assert_allclose(var1, var2, check_dim_order=False)
+
+
 class CustomIndex(Index):
     """Custom index without equals() implementation for testing."""