REF: Refactor reference handling to the BlockManager level for iterrows (#51431)

Author: phofl

pandas/core/frame.py

@@ -1396,10 +1396,7 @@ def iterrows(self) -> Iterable[tuple[Hashable, Series]]:
         for k, v in zip(self.index, self.values):
             s = klass(v, index=columns, name=k).__finalize__(self)
             if using_cow and self._mgr.is_single_block:
-                s._mgr.blocks[0].refs = self._mgr.blocks[0].refs  # type: ignore[union-attr]  # noqa
-                s._mgr.blocks[0].refs.add_reference(  # type: ignore[union-attr]
-                    s._mgr.blocks[0]  # type: ignore[arg-type, union-attr]
-                )
+                s._mgr.add_references(self._mgr)  # type: ignore[arg-type]
             yield k, s
 
     def itertuples(

pandas/core/internals/array_manager.py

@@ -169,6 +169,12 @@ def set_axis(self, axis: AxisInt, new_labels: Index) -> None:
     def get_dtypes(self) -> np.ndarray:
         return np.array([arr.dtype for arr in self.arrays], dtype="object")
 
+    def add_references(self, mgr: BaseArrayManager) -> None:
+        """
+        Only implemented on the BlockManager level
+        """
+        return
+
     def __getstate__(self):
         return self.arrays, self._axes
 

pandas/core/internals/managers.py

@@ -247,6 +247,17 @@ def _has_no_reference_block(self, blkno: int) -> bool:
         """
         return not self.blocks[blkno].refs.has_reference()
 
+    def add_references(self, mgr: BaseBlockManager) -> None:
+        """
+        Adds the references from one manager to another. We assume that both
+        managers have the same block structure.
+        """
+        for i, blk in enumerate(self.blocks):
+            blk.refs = mgr.blocks[i].refs
+            # Argument 1 to "add_reference" of "BlockValuesRefs" has incompatible type
+            # "Block"; expected "SharedBlock"
+            blk.refs.add_reference(blk)  # type: ignore[arg-type]
+
     def get_dtypes(self):
         dtypes = np.array([blk.dtype for blk in self.blocks])
         return dtypes.take(self.blknos)