From d1896d24c0dab09761f9e83a99ed7311331bb4f2 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 13:58:52 -0700 Subject: [PATCH 1/2] fix: prepared-statement cache id-recycling corruption for DataFrame/Arrow params MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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. --- src_py/connection.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/src_py/connection.py b/src_py/connection.py index 806f47a..fa90299 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -7,6 +7,7 @@ import threading import uuid import warnings +import weakref from typing import TYPE_CHECKING, Any from weakref import WeakSet @@ -117,6 +118,16 @@ def _capi_param_signature(parameters: dict[str, Any]) -> tuple: ) +# Maps id(object) to (generation, weakref) for pointer-type parameters +# (DataFrames, Arrow tables, etc.). We use weakref to detect when the +# original object has been garbage-collected. If the object is gone and +# its id is recycled, the tracker assigns a new generation so the +# prepared-statement cache does not return a stale entry compiled for +# the old (now freed) object. +_pybind_pointer_gen = 0 +_pybind_pointer_tracker: dict[int, tuple[int, weakref.ref]] = {} + + def _pybind_int_signature(value: int) -> tuple[str]: if -(2**7) <= value <= 2**7 - 1: return ("int8",) @@ -175,7 +186,21 @@ def _pybind_value_signature(value: Any) -> tuple: return ("string",) module_name = type(value).__module__ if module_name.startswith(("pandas", "polars", "pyarrow")): - return ("pointer", module_name, type(value).__name__, id(value)) + global _pybind_pointer_gen, _pybind_pointer_tracker + obj_id = id(value) + entry = _pybind_pointer_tracker.get(obj_id) + if entry is not None: + gen, weak = entry + if weak() is not None: + return ("pointer", module_name, type(value).__name__, gen) + # Original object was GC'd and a new object now occupies the + # same memory address. Purge the stale entry and fall + # through to assign a fresh generation. + del _pybind_pointer_tracker[obj_id] + _pybind_pointer_gen += 1 + gen = _pybind_pointer_gen + _pybind_pointer_tracker[obj_id] = (gen, weakref.ref(value)) + return ("pointer", module_name, type(value).__name__, gen) if isinstance(value, dict): items = list(value.items()) if ( From 03c589e8445c060e285039e34ea32b0fb300ed58 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 14:07:45 -0700 Subject: [PATCH 2/2] fix: add schema fingerprint to pointer cache key for in-place mutation safety Include a schema fingerprint (shape + column names + dtypes for pandas, field names + types for pyarrow/polars) in the prepared-statement cache key for pointer-type parameters. Without this, mutating a DataFrame in place between execute() calls (e.g. df.drop(columns=[...], inplace=True)) would reuse the cached prepared statement compiled for the original schema, causing data corruption. The cache key tuple expands from 4 to 5 elements: ("pointer", module, type_name, gen) -> ("pointer", module, type_name, gen, schema_fp) The fingerprint is computed fresh on every call; two objects with the same identity and schema share the cache entry, while schema mutations produce a different key and trigger a fresh prepare. --- src_py/connection.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/src_py/connection.py b/src_py/connection.py index fa90299..8ca77b3 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -128,6 +128,31 @@ def _capi_param_signature(parameters: dict[str, Any]) -> tuple: _pybind_pointer_tracker: dict[int, tuple[int, weakref.ref]] = {} +def _get_pointer_schema_fingerprint(value: Any) -> int: + """ + Compute a hash of the schema for a pandas/pyarrow/polars object. + + Two objects with the same shape, column names, and column dtypes + produce the same fingerprint. If the object is mutated in place + (e.g. columns dropped, renamed, or dtypes changed) the fingerprint + changes, causing a cache miss so a fresh prepared statement is + compiled for the new schema. + """ + module_name = type(value).__module__ + if module_name == "pandas.core.frame": + cols = list(value.columns) + return hash( + (value.shape, tuple(cols), tuple(str(value[col].dtype) for col in cols)) + ) + if module_name.startswith("pyarrow"): + schema = value.schema + return hash(tuple((f.name, str(f.type)) for f in schema)) + if module_name.startswith("polars"): + schema = value.schema + return hash(tuple(sorted((name, str(dtype)) for name, dtype in schema.items()))) + return 0 + + def _pybind_int_signature(value: int) -> tuple[str]: if -(2**7) <= value <= 2**7 - 1: return ("int8",) @@ -188,11 +213,12 @@ def _pybind_value_signature(value: Any) -> tuple: if module_name.startswith(("pandas", "polars", "pyarrow")): global _pybind_pointer_gen, _pybind_pointer_tracker obj_id = id(value) + fp = _get_pointer_schema_fingerprint(value) entry = _pybind_pointer_tracker.get(obj_id) if entry is not None: gen, weak = entry if weak() is not None: - return ("pointer", module_name, type(value).__name__, gen) + return ("pointer", module_name, type(value).__name__, gen, fp) # Original object was GC'd and a new object now occupies the # same memory address. Purge the stale entry and fall # through to assign a fresh generation. @@ -200,7 +226,7 @@ def _pybind_value_signature(value: Any) -> tuple: _pybind_pointer_gen += 1 gen = _pybind_pointer_gen _pybind_pointer_tracker[obj_id] = (gen, weakref.ref(value)) - return ("pointer", module_name, type(value).__name__, gen) + return ("pointer", module_name, type(value).__name__, gen, fp) if isinstance(value, dict): items = list(value.items()) if (