Fix DataFrame.stack/unstack pandas incompatibilities

python/cudf/cudf/core/column/numerical.py

@@ -911,20 +911,23 @@ def as_numerical_column(self, dtype: DtypeObj) -> NumericalColumn:
                 self.dtype
             ):
                 # Short-circuit the cast if the dtypes are equivalent
-                # but not the same type object.
+                # but not the same type object. Do NOT mutate self._dtype:
+                # the column object may be shared with the caller's frame.
                 if (
                     is_pandas_nullable_extension_dtype(dtype)
                     and isinstance(self.dtype, np.dtype)
                     and self.dtype.kind == "f"
                 ):
-                    # If the dtype is a pandas nullable extension type, we need to
-                    # float column doesn't have any NaNs.
+                    # NaNs must become nulls before viewing as a masked dtype.
                     res = self.nans_to_nulls()
-                    res._dtype = dtype
-                    return res
-                else:
-                    self._dtype = dtype
-                    return self
+                    return cast(
+                        "NumericalColumn",
+                        ColumnBase.create(res.plc_column, dtype),
+                    )
+                return cast(
+                    "NumericalColumn",
+                    ColumnBase.create(self.plc_column, dtype),
+                )
             if self.dtype.kind == "f" and dtype.kind in "iu":
                 if not is_pandas_nullable_extension_dtype(dtype) and (
                     self.nan_count > 0

python/cudf/cudf/core/column_accessor.py

@@ -35,6 +35,32 @@ def _is_bool(val: Any) -> bool:
     return isinstance(val, (bool, np.bool_))
 
 
+def _is_nan_scalar(val: Any) -> bool:
+    return isinstance(val, float) 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.
+
+    Distinct float('nan') objects hash and compare unequal, so NaN-containing
+    labels round-tripped through a pandas Index (which materializes fresh NaN
+    objects on iteration) never match dict keys. Canonicalizing to the np.nan
+    singleton restores pandas' all-NaNs-are-equal label semantics (dict/tuple
+    comparison uses the per-element identity shortcut).
+    """
+    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.
 
@@ -125,6 +151,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.
@@ -162,7 +193,19 @@ 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
+                # compare equal, matching pandas label semantics.
+                canon = _canonicalize_nan_label(key)
+                for existing in self._data:
+                    c = _canonicalize_nan_label(existing)
+                    if c is canon or c == canon:
+                        return self._data[existing]
+            raise
 
     def __setitem__(self, key: Hashable, value: ColumnBase) -> None:
         self.set_by_label(key, value)
@@ -322,6 +365,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: