Preserve MultiIndex column fidelity through ColumnAccessor round trips

python/cudf/cudf/core/column_accessor.py

@@ -35,6 +35,36 @@ def _is_bool(val: Any) -> bool:
     return isinstance(val, (bool, np.bool_))
 
 
+def _is_nan_scalar(val: Any) -> bool:
+    return isinstance(val, (float, np.floating)) and val != val
+
+
+def _label_contains_nan(label: Any) -> bool:
+    if isinstance(label, tuple):
+        return any(_is_nan_scalar(lv) for lv in label)
+    return _is_nan_scalar(label)
+
+
+def _canonicalize_nan_label(label: Any) -> Any:
+    """Map float NaN elements of a label to the np.nan singleton.
+
+    No NaN object ever compares equal to any NaN (including itself), so
+    equality alone can never match a NaN-containing label. dict lookups and
+    tuple comparison, however, short-circuit on object *identity* before
+    trying ``==`` — and that shortcut is exactly what this canonicalization
+    targets: with every NaN mapped to the one ``np.nan`` object, two
+    canonicalized labels holding NaN in the same position match by identity.
+    This restores pandas' all-NaNs-are-equal label semantics for labels
+    round-tripped through a pandas Index, which materializes fresh NaN
+    objects on iteration.
+    """
+    if isinstance(label, tuple):
+        return tuple(np.nan if _is_nan_scalar(lv) else lv for lv in label)
+    if _is_nan_scalar(label):
+        return np.nan
+    return label
+
+
 class _NestedGetItemDict(dict):
     """A dictionary whose __getitem__ method accesses nested dicts.
 
@@ -102,6 +132,10 @@ class ColumnAccessor(MutableMapping):
     verify : bool, optional
         For non ColumnAccessor inputs, whether to verify
         column length and data.values() are all Columns
+    pandas_index : pd.Index, optional
+        The source pandas index the keys were taken from, if any.
+        A matching pd.MultiIndex primes the ``to_pandas_index`` cache
+        (see ``_prime_to_pandas_index``).
     """
 
     _data: dict[Hashable, ColumnBase]
@@ -117,6 +151,7 @@ def __init__(
         label_dtype: DtypeObj | None = None,
         verify: bool = True,
         level_dtypes: tuple[DtypeObj, ...] | None = None,
+        pandas_index: pd.Index | None = None,
     ) -> None:
         if isinstance(data, ColumnAccessor):
             self._data = data._data
@@ -125,6 +160,11 @@ def __init__(
             self.rangeindex: bool = data.rangeindex
             self.label_dtype: DtypeObj | None = data.label_dtype
             self._level_dtypes = data._level_dtypes
+            if "to_pandas_index" in data.__dict__:
+                # carry over the primed/cached pandas index: it holds
+                # fidelity (e.g. explicit unsorted level order) that a
+                # rebuild from tuples would lose
+                self.to_pandas_index = data.__dict__["to_pandas_index"]
         elif isinstance(data, MutableMapping):
             # This code path is performance-critical for copies and should be
             # modified with care.
@@ -157,12 +197,50 @@ def __init__(
             raise ValueError(
                 f"data must be a ColumnAccessor or MutableMapping, not {type(data).__name__}"
             )
+        if pandas_index is not None:
+            self._prime_to_pandas_index(pandas_index)
+
+    def _prime_to_pandas_index(self, index: pd.Index) -> None:
+        """Prime the cached ``to_pandas_index`` with the exact source index.
+
+        Rebuilding a pandas MultiIndex from the stored tuple labels re-sorts
+        its levels, losing an explicit unsorted level layout (the level
+        order affects pandas operations that work on level codes, e.g.
+        legacy ``stack(sort=True)``). Keeping the source MultiIndex itself
+        preserves that fidelity. Only a hierarchical columns axis whose
+        length matches the data is primed; anything else is ignored.
+        """
+        if (
+            self.multiindex
+            and isinstance(index, pd.MultiIndex)
+            and len(self._data) == len(index)
+        ):
+            self.to_pandas_index = index
 
     def __iter__(self) -> Iterator:
         return iter(self._data)
 
     def __getitem__(self, key: Hashable) -> ColumnBase:
-        return self._data[key]
+        try:
+            return self._data[key]
+        except KeyError:
+            if _label_contains_nan(key):
+                # NaN labels lose object identity when round-tripped through
+                # a pandas Index; retry with NaNs canonicalized so all NaNs
+                # match by identity, as pandas label semantics require.
+                canon = _canonicalize_nan_label(key)
+                for existing in self._data:
+                    c = _canonicalize_nan_label(existing)
+                    try:
+                        match = c is canon or bool(c == canon)
+                    except TypeError:
+                        # e.g. a pd.NA label: its comparisons return pd.NA,
+                        # whose truthiness raises. Ambiguity is not a match;
+                        # keep scanning for a genuine NaN label.
+                        match = False
+                    if match:
+                        return self._data[existing]
+            raise
 
     def __setitem__(self, key: Hashable, value: ColumnBase) -> None:
         self.set_by_label(key, value)
@@ -322,6 +400,35 @@ def to_pandas_index(self) -> pd.Index:
                     self.names,
                     names=self.level_names,
                 )
+                if (
+                    self._level_dtypes is not None
+                    and len(self._level_dtypes) == result.nlevels
+                ):
+                    # ``from_tuples`` re-infers every level dtype from the
+                    # materialized labels, degrading e.g. categorical levels
+                    # to str, object levels to str once mixed-type labels are
+                    # selected away, and int64 levels with missing entries to
+                    # float64. Restore each preserved level dtype when the
+                    # cast is lossless (round-trips to the inferred values).
+                    new_levels = []
+                    changed = False
+                    for lvl, level_dtype in zip(
+                        result.levels, self._level_dtypes, strict=True
+                    ):
+                        if lvl.dtype != level_dtype:
+                            try:
+                                cast_lvl = lvl.astype(level_dtype)
+                            except (TypeError, ValueError):
+                                pass
+                            else:
+                                # missing-aware equality: ``==`` would treat
+                                # NaN entries as unequal to themselves
+                                if cast_lvl.astype(lvl.dtype).equals(lvl):
+                                    lvl = cast_lvl
+                                    changed = True
+                        new_levels.append(lvl)
+                    if changed:
+                        result = result.set_levels(new_levels)
         else:
             # Determine if we can return a RangeIndex
             if self.rangeindex:

python/cudf/cudf/core/dataframe.py

@@ -615,7 +615,12 @@ def _pd_index_level_dtypes(idx) -> tuple | None:
     dtype cannot be inferred from zero entries).
     """
     if isinstance(idx, pd.MultiIndex):
-        return tuple(idx.get_level_values(i).dtype for i in range(idx.nlevels))
+        # use the levels, not get_level_values: materializing a level that
+        # has missing entries (-1 codes) upcasts e.g. int64 to float64
+        return tuple(level.dtype for level in idx.levels)
+    if isinstance(idx, cudf.MultiIndex):
+        # the per-row columns share their dtype with the levels
+        return tuple(dtype for _, dtype in idx._dtypes)
     return None
 
 
@@ -837,7 +842,7 @@ def _array_to_column_accessor(
         columns_labels = columns
     else:
         columns_labels = pd.RangeIndex(data.shape[1])
-    return ColumnAccessor(
+    ca = ColumnAccessor(
         {
             column_label: as_column(data[:, i], nan_as_null=nan_as_null)
             for column_label, i in zip(
@@ -850,7 +855,9 @@ def _array_to_column_accessor(
         label_dtype=columns_labels.dtype,
         level_names=tuple(columns_labels.names),
         level_dtypes=_pd_index_level_dtypes(columns_labels),
+        pandas_index=columns_labels,
     )
+    return ca
 
 
 @_performance_tracking
@@ -1280,6 +1287,7 @@ def __init__(
                 level_names=tuple(columns.names),
                 label_dtype=columns.dtype,
                 level_dtypes=_pd_index_level_dtypes(columns),
+                pandas_index=columns,
             )
         elif isinstance(data, Mapping):
             # Note: We excluded ColumnAccessor already above
@@ -1352,6 +1360,12 @@ def __init__(
         if dtype:
             self._data = self.astype(dtype)._data
 
+        final_pd_columns = (
+            second_columns if second_columns is not None else columns
+        )
+        if isinstance(final_pd_columns, pd.Index):
+            self._data._prime_to_pandas_index(final_pd_columns)
+
     @classmethod
     def _from_data(  # type: ignore[override]
         cls,
@@ -2554,6 +2568,11 @@ def _fill_same_ca_attributes(
                 )
             elif self._data._level_names == other._data._level_names:
                 ca_attributes["level_names"] = self._data.level_names
+                if self._data.multiindex == other._data.multiindex:
+                    # equal labels can still fail the ``equals`` check above
+                    # on level-dtype differences (e.g. Int8 vs int64); the
+                    # result keeps hierarchical columns like pandas
+                    ca_attributes["multiindex"] = self._data.multiindex
         elif isinstance(other, (dict, Mapping)):
             # Need to fail early on host mapping types because we ultimately
             # convert everything to a dict.
@@ -3237,6 +3256,7 @@ def columns(self, columns):
         rangeindex = False
         label_dtype = None
         level_names = None
+        level_dtypes = None
         if isinstance(columns, (pd.MultiIndex, cudf.MultiIndex)):
             multiindex = True
             if isinstance(columns, cudf.MultiIndex):
@@ -3246,6 +3266,7 @@ def columns(self, columns):
             if pd_columns.nunique(dropna=False) != len(pd_columns):
                 raise ValueError("Duplicate column names are not allowed")
             level_names = list(pd_columns.names)
+            level_dtypes = _pd_index_level_dtypes(pd_columns)
         elif isinstance(columns, (Index, ColumnBase, Series)):
             level_names = (getattr(columns, "name", None),)
             rangeindex = isinstance(columns, cudf.RangeIndex)
@@ -3282,7 +3303,9 @@ def columns(self, columns):
             level_names=level_names,
             label_dtype=label_dtype,
             rangeindex=rangeindex,
+            level_dtypes=level_dtypes,
             verify=False,
+            pandas_index=pd_columns,
         )
 
     def _set_columns_like(self, other: ColumnAccessor) -> None:

python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py

@@ -1391,9 +1391,6 @@ def pytest_unconfigure(config):
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_cast": "TODO: Add a reason for failure",
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_multiindex": "TODO: Add a reason for failure",
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[None]": "TODO: Add a reason for failure",
-    "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[nan0]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='nan') are different",
-    "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key[nan1]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='nan') are different",
-    "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_nan_key_and_columns": "TODO: Add a reason for failure",
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_index": "TODO: Add a reason for failure",
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_index_and_columns": "TODO: Add a reason for failure",
     "tests/frame/test_constructors.py::TestDataFrameConstructors::test_constructor_dict_with_none": "AssertionError: assert nan is None",
@@ -1720,7 +1717,6 @@ def pytest_unconfigure(config):
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[False]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_in_multiindex_columns[True]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_nan_level[False]": "TODO: Add a reason for failure",
-    "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_order_with_unsorted_levels_multi_row_2[False]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[False]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_multiple[True]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_preserve_names[False]": "TODO: Add a reason for failure",
@@ -1729,10 +1725,6 @@ def pytest_unconfigure(config):
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_stack_unstack_wrong_level_name[True-unstack]": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_preserve_types": "TODO: Add a reason for failure",
     "tests/frame/test_stack_unstack.py::TestStackUnstackMultiLevel::test_unstack_with_missing_int_cast_to_float": "TODO: Add a reason for failure",
-    "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=first]": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='('value', nan)') are different",
-    "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=last]": "AssertionError: Attributes of DataFrame.iloc[:, 3] (column name='('value', nan)') are different",
-    "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=second]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='('value', nan)') are different",
-    "tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=third]": "AssertionError: Attributes of DataFrame.iloc[:, 2] (column name='('value', nan)') are different",
     "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_asof": "TODO: Add a reason for failure",
     "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_equals_subclass": "TODO: Add a reason for failure",
     "tests/frame/test_subclass.py::TestDataFrameSubclassing::test_frame_subclassing_and_slicing": "TODO: Add a reason for failure",
@@ -1763,7 +1755,6 @@ def pytest_unconfigure(config):
     "tests/groupby/aggregate/test_aggregate.py::test_agg_str_with_kwarg_axis_1_raises[nunique]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_directory[size]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_empty_key_empty_return": "TODO: Add a reason for failure",
-    "tests/groupby/aggregate/test_aggregate.py::test_order_aggregate_multiple_funcs": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-False]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-True]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-sum-large_int-False]": "TODO: Add a reason for failure",
@@ -2874,7 +2865,6 @@ def pytest_unconfigure(config):
     "tests/reshape/concat/test_categorical.py::TestCategoricalConcat::test_categorical_index_upcast": "TODO: Add a reason for failure",
     "tests/reshape/concat/test_categorical.py::TestCategoricalConcat::test_concat_categorical_datetime": "TODO: Add a reason for failure",
     "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_copy": "TODO: Add a reason for failure",
-    "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_keys_specific_levels": "TODO: Add a reason for failure",
     "tests/reshape/concat/test_concat.py::TestConcatenate::test_concat_order": "TODO: Add a reason for failure",
     "tests/reshape/concat/test_concat.py::test_concat_empty_and_non_empty_frame_regression": "TODO: Add a reason for failure",
     "tests/reshape/concat/test_concat.py::test_concat_ignore_empty_object_float[None-datetime64[ns]]": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='foo') are different",
@@ -3051,8 +3041,6 @@ def pytest_unconfigure(config):
     "tests/reshape/test_melt.py::TestWideToLong::test_raise_of_column_name_value": "TODO: Add a reason for failure",
     "tests/reshape/test_pivot.py::TestPivot::test_pivot_index_is_none": "AssertionError: DataFrame.index are different",
     "tests/reshape/test_pivot.py::TestPivotTable::test_categorical_pivot_index_ordering[False]": "TODO: Add a reason for failure",
-    "tests/reshape/test_pivot.py::TestPivotTable::test_daily": "TODO: Add a reason for failure",
-    "tests/reshape/test_pivot.py::TestPivotTable::test_monthly": "TODO: Add a reason for failure",
     "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_complex_aggfunc": "TODO: Add a reason for failure",
     "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_datetime_tz": "ValueError: Length of names must match number of levels in MultiIndex.",
     "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_index_with_nan[False]": "AssertionError: DataFrame.index are different",
@@ -3874,7 +3862,6 @@ def pytest_unconfigure(config):
     "tests/window/test_timeseries_window.py::TestRollingTS::test_rolling_on_decreasing_index[us]": "TODO: Add a reason for failure",
     "tests/window/test_win_type.py::test_cmov_window_corner[None]": "TODO: Add a reason for failure",
     "tests/window/test_win_type.py::test_win_type_not_implemented": "TODO: Add a reason for failure",
-    "tests/indexing/multiindex/test_loc.py::test_loc_getitem_duplicates_multiindex_empty_indexer[columns_indexer1]": "AssertionError: DataFrame.columns level [0] are different",
 }
 
 # Keep keys in alphabeical order