ENH: Improve ref-tracking for group keys

pandas/core/groupby/grouper.py

@@ -947,9 +947,19 @@ def is_in_obj(gpr) -> bool:
             # For the CoW case, we need an equality check as the identity check
             # no longer works (each Series from column access is a new object)
             try:
-                return gpr.equals(obj[gpr.name])
+                obj_gpr_column = obj[gpr.name]
             except (AttributeError, KeyError, IndexError, InvalidIndexError):
                 return False
+            if isinstance(gpr, Series) and isinstance(obj_gpr_column, Series):
+                # Item "SingleArrayManager" of "Union[SingleArrayManager,
+                # SingleBlockManager]" has no attribute "references_same_values"
+                # Argument 1 to "references_same_values" of "BaseBlockManager" has
+                # incompatible type "Union[SingleArrayManager, SingleBlockManager]";
+                # expected "BaseBlockManager
+                return gpr._mgr.references_same_values(  # type: ignore[union-attr]
+                    obj_gpr_column._mgr, 0  # type: ignore[arg-type]
+                )
+            return False
         try:
             return gpr is obj[gpr.name]
         except (KeyError, IndexError, InvalidIndexError):

pandas/core/internals/managers.py

@@ -11,6 +11,7 @@
     cast,
 )
 import warnings
+import weakref
 
 import numpy as np
 
@@ -247,6 +248,14 @@ def _has_no_reference_block(self, blkno: int) -> bool:
         """
         return not self.blocks[blkno].refs.has_reference()
 
+    def references_same_values(self, mgr: BaseBlockManager, blkno: int) -> bool:
+        """
+        Checks if two blocks from two different block managers reference the
+        same underlying values.
+        """
+        ref = weakref.ref(self.blocks[blkno])
+        return ref in mgr.blocks[blkno].refs.referenced_blocks
+
     def get_dtypes(self):
         dtypes = np.array([blk.dtype for blk in self.blocks])
         return dtypes.take(self.blknos)