Skip to content

Commit

Permalink
CoW: Use weakref callbacks to track dead references
Browse files Browse the repository at this point in the history
Co-authored-by: José Lucas Silva Mayer <[email protected]>
  • Loading branch information
wangwillian0 and josemayer committed Oct 16, 2023
1 parent 364c9cb commit 9cc2a3e
Showing 1 changed file with 16 additions and 7 deletions.
23 changes: 16 additions & 7 deletions pandas/_libs/internals.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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) // 8:
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.
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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

0 comments on commit 9cc2a3e

Please sign in to comment.