Skip to content

Commit 93fd1d7

Browse files
committed
fix: prepared-statement cache id-recycling corruption for DataFrame/Arrow params
_pybind_value_signature used raw id(value) for pointer-type parameters (DataFrames, Arrow tables, Polars DataFrames). Python recycles object ids after garbage collection (~99.7% of the time a new allocation reuses a freed id). On a shared connection (conn_db_readonly), this caused the cache to return a stale prepared statement compiled for a completely different DataFrame — e.g. string columns from test_pyarrow_string reused for dictionary-encoded int32 columns in test_pyarrow_dict — producing data corruption (dictionary indices read as string offsets). Fix: replace raw id(value) with a weakref-validated generation counter. _pybind_pointer_tracker[id(obj)] -> (gen, weakref.ref(obj)) On each call: If id is found AND the weakref is still alive → reuse the cached generation (same object → cache hit, correct). If the weakref is dead (original object was GCd and id recycled) → purge the stale entry and assign a fresh generation. If not found → assign a fresh generation and store the weakref. This preserves correct caching for repeated calls with the same object while being immune to id recycling. Fixes test_pyarrow_dict intermittent AssertionError.
1 parent b03788a commit 93fd1d7

1 file changed

Lines changed: 26 additions & 1 deletion

File tree

src_py/connection.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import threading
88
import uuid
99
import warnings
10+
import weakref
1011
from typing import TYPE_CHECKING, Any
1112
from weakref import WeakSet
1213

@@ -117,6 +118,16 @@ def _capi_param_signature(parameters: dict[str, Any]) -> tuple:
117118
)
118119

119120

121+
# Maps id(object) to (generation, weakref) for pointer-type parameters
122+
# (DataFrames, Arrow tables, etc.). We use weakref to detect when the
123+
# original object has been garbage-collected. If the object is gone and
124+
# its id is recycled, the tracker assigns a new generation so the
125+
# prepared-statement cache does not return a stale entry compiled for
126+
# the old (now freed) object.
127+
_pybind_pointer_gen = 0
128+
_pybind_pointer_tracker: dict[int, tuple[int, weakref.ref]] = {}
129+
130+
120131
def _pybind_int_signature(value: int) -> tuple[str]:
121132
if -(2**7) <= value <= 2**7 - 1:
122133
return ("int8",)
@@ -175,7 +186,21 @@ def _pybind_value_signature(value: Any) -> tuple:
175186
return ("string",)
176187
module_name = type(value).__module__
177188
if module_name.startswith(("pandas", "polars", "pyarrow")):
178-
return ("pointer", module_name, type(value).__name__, id(value))
189+
global _pybind_pointer_gen, _pybind_pointer_tracker
190+
obj_id = id(value)
191+
entry = _pybind_pointer_tracker.get(obj_id)
192+
if entry is not None:
193+
gen, weak = entry
194+
if weak() is not None:
195+
return ("pointer", module_name, type(value).__name__, gen)
196+
# Original object was GC'd and a new object now occupies the
197+
# same memory address. Purge the stale entry and fall
198+
# through to assign a fresh generation.
199+
del _pybind_pointer_tracker[obj_id]
200+
_pybind_pointer_gen += 1
201+
gen = _pybind_pointer_gen
202+
_pybind_pointer_tracker[obj_id] = (gen, weakref.ref(value))
203+
return ("pointer", module_name, type(value).__name__, gen)
179204
if isinstance(value, dict):
180205
items = list(value.items())
181206
if (

0 commit comments

Comments
 (0)