CoW: Use weakref callbacks to track dead references

Co-authored-by: José Lucas Silva Mayer <josemayer@usp.br>

Author: wangwillian0

pandas/_libs/internals.pyx

@@ -890,17 +890,29 @@ cdef class BlockValuesRefs:
     """
     cdef:
         public list referenced_blocks
+        public int dead_counter
+        object __weakref__
+        object _weakref_cb
 
     def __cinit__(self, blk: Block | None = None) -> None:
+        def _weakref_cb(item: weakref.ref, selfref: weakref.ref = weakref.ref(self)) -> None:
+            self = selfref()
+            if self is not None:
+                self.dead_counter += 1
+                if self.dead_counter > 256 and self.dead_counter > len(self.referenced_blocks) // 2:
+                    self._clear_dead_references()
+        self._weakref_cb = _weakref_cb
         if blk is not None:
-            self.referenced_blocks = [weakref.ref(blk)]
+            self.referenced_blocks = [weakref.ref(blk, self._weakref_cb)]
         else:
             self.referenced_blocks = []
 
     def _clear_dead_references(self) -> None:
+        old_len = len(self.referenced_blocks)
         self.referenced_blocks = [
             ref for ref in self.referenced_blocks if ref() is not None
         ]
+        self.dead_counter = self.dead_counter - (old_len - len(self.referenced_blocks))
 
     def add_reference(self, blk: Block) -> None:
         """Adds a new reference to our reference collection.
@@ -910,8 +922,7 @@ cdef class BlockValuesRefs:
         blk : Block
             The block that the new references should point to.
         """
-        self._clear_dead_references()
-        self.referenced_blocks.append(weakref.ref(blk))
+        self.referenced_blocks.append(weakref.ref(blk, self._weakref_cb))
 
     def add_index_reference(self, index: object) -> None:
         """Adds a new reference to our reference collection when creating an index.
@@ -921,8 +932,7 @@ cdef class BlockValuesRefs:
         index : Index
             The index that the new reference should point to.
         """
-        self._clear_dead_references()
-        self.referenced_blocks.append(weakref.ref(index))
+        self.referenced_blocks.append(weakref.ref(index, self._weakref_cb))
 
     def has_reference(self) -> bool:
         """Checks if block has foreign references.
@@ -934,6 +944,5 @@ cdef class BlockValuesRefs:
         -------
         bool
         """
-        self._clear_dead_references()
         # Checking for more references than block pointing to itself
-        return len(self.referenced_blocks) > 1
+        return len(self.referenced_blocks) - self.dead_counter > 1