From 9e497188a39b6822900c0399e279161c4a39e297 Mon Sep 17 00:00:00 2001 From: ericyuanhui <285521263@qq.com> Date: Thu, 9 Jul 2026 19:24:07 +0800 Subject: [PATCH 1/8] fix python prepare params implicit cause RSS raising Signed-off-by: ericyuanhui <285521263@qq.com> --- examples/ice_disk_to_arrow_memory.py | 12 +- src_py/connection.py | 18 ++- test/test_pybind_implicit_prepare_cache.py | 149 +++++++++++++++++++++ 3 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 test/test_pybind_implicit_prepare_cache.py diff --git a/examples/ice_disk_to_arrow_memory.py b/examples/ice_disk_to_arrow_memory.py index 52df48a..e64fa51 100644 --- a/examples/ice_disk_to_arrow_memory.py +++ b/examples/ice_disk_to_arrow_memory.py @@ -77,11 +77,17 @@ def main() -> None: db = lb.Database(":memory:") conn = db.connect() if args.layout == "flat": - register_flat(conn, args.data_dir, args.node_table, args.rel_table, src_table, dst_table) + register_flat( + conn, args.data_dir, args.node_table, args.rel_table, src_table, dst_table + ) else: - register_csr(conn, args.data_dir, args.node_table, args.rel_table, src_table, dst_table) + register_csr( + conn, args.data_dir, args.node_table, args.rel_table, src_table, dst_table + ) - result = conn.execute(f"MATCH (a:{src_table})-[r:{args.rel_table}]->(b:{dst_table}) RETURN COUNT(*)") + result = conn.execute( + f"MATCH (a:{src_table})-[r:{args.rel_table}]->(b:{dst_table}) RETURN COUNT(*)" + ) print(result.get_next()[0]) diff --git a/src_py/connection.py b/src_py/connection.py index fa41b89..98e6097 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -52,6 +52,7 @@ def __init__(self, database: Database, num_threads: int = 0): self._query_timeout_ms = 0 self._query_results: WeakSet[QueryResult] = WeakSet() self._capi_scan_tables: set[str] = set() + self._pybind_implicit_prepared_cache: dict[str, Any] = {} self.database._register_connection(self) self.init_connection() @@ -116,6 +117,7 @@ def close(self) -> None: for query_result in list(self._query_results): query_result.close() self._query_results.clear() + self._pybind_implicit_prepared_cache.clear() if self._connection is not None and not self.database.is_closed: self._connection.close() @@ -460,9 +462,23 @@ def _execute_with_pybind( return py_connection.query(query) query, parameters = self._normalize_parameters_for_pybind(query, parameters) - prepared = py_connection.prepare(query, parameters) + prepared = self._get_or_prepare_pybind_statement( + py_connection, query, parameters + ) return py_connection.execute(prepared, parameters) + def _get_or_prepare_pybind_statement( + self, + py_connection: Any, + query: str, + parameters: dict[str, Any], + ) -> Any: + prepared = self._pybind_implicit_prepared_cache.get(query) + if prepared is None: + prepared = py_connection.prepare(query, parameters) + self._pybind_implicit_prepared_cache[query] = prepared + return prepared + def _maybe_raise_scan_unsupported_object(self, query: str) -> None: match = re.search( r"\bLOAD\s+FROM\s+([A-Za-z_][A-Za-z0-9_]*)\b", query, re.IGNORECASE diff --git a/test/test_pybind_implicit_prepare_cache.py b/test/test_pybind_implicit_prepare_cache.py new file mode 100644 index 0000000..7df3719 --- /dev/null +++ b/test/test_pybind_implicit_prepare_cache.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import ladybug as lb +import ladybug.connection as lb_connection +import pytest + + +class _FakeResult: + def isSuccess(self) -> bool: + return True + + def hasNextQueryResult(self) -> bool: + return False + + +class _FakePreparedStatement: + def __init__(self, query: str, parameters: dict[str, object]): + self.query = query + self.parameters = dict(parameters) + + +class _FakePybindConnection: + def __init__(self) -> None: + self.prepare_calls: list[tuple[str, dict[str, object]]] = [] + self.execute_calls: list[tuple[_FakePreparedStatement, dict[str, object]]] = [] + self.query_calls: list[str] = [] + self.closed = False + + def prepare( + self, query: str, parameters: dict[str, object] + ) -> _FakePreparedStatement: + self.prepare_calls.append((query, dict(parameters))) + return _FakePreparedStatement(query, parameters) + + def execute( + self, prepared: _FakePreparedStatement, parameters: dict[str, object] + ) -> _FakeResult: + self.execute_calls.append((prepared, dict(parameters))) + return _FakeResult() + + def query(self, query: str) -> _FakeResult: + self.query_calls.append(query) + return _FakeResult() + + def close(self) -> None: + self.closed = True + + +class _FakeBackendConnection: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class _FakeDatabase: + def __init__(self) -> None: + self._use_pybind_backend = True + self._database = object() + self.is_closed = False + self.registered_connections: list[lb.Connection] = [] + + def _register_connection(self, connection: lb.Connection) -> None: + self.registered_connections.append(connection) + + def _unregister_connection(self, connection: lb.Connection) -> None: + self.registered_connections.remove(connection) + + def init_database(self) -> None: + return None + + +@pytest.fixture +def fake_pybind_connection(monkeypatch: pytest.MonkeyPatch) -> _FakePybindConnection: + fake_pybind = _FakePybindConnection() + + monkeypatch.setattr(lb_connection, "get_pybind_module", lambda: SimpleNamespace()) + monkeypatch.setattr( + lb.Connection, + "init_connection", + lambda self: setattr(self, "_connection", _FakeBackendConnection()), + ) + monkeypatch.setattr( + lb.Connection, "_get_pybind_connection", lambda self: fake_pybind + ) + return fake_pybind + + +def test_pybind_implicit_prepare_reuses_same_query( + fake_pybind_connection: _FakePybindConnection, +) -> None: + conn = lb.Connection(_FakeDatabase()) + + conn.execute("RETURN $value", {"value": 1}) + conn.execute("RETURN $value", {"value": 2}) + + assert fake_pybind_connection.query_calls == [] + assert len(fake_pybind_connection.prepare_calls) == 1 + assert [call[0].query for call in fake_pybind_connection.execute_calls] == [ + "RETURN $value", + "RETURN $value", + ] + assert [call[1] for call in fake_pybind_connection.execute_calls] == [ + {"value": 1}, + {"value": 2}, + ] + + +def test_pybind_implicit_prepare_does_not_share_different_queries( + fake_pybind_connection: _FakePybindConnection, +) -> None: + conn = lb.Connection(_FakeDatabase()) + + conn.execute("RETURN $value", {"value": 1}) + conn.execute("RETURN $other", {"other": 1}) + + assert [call[0] for call in fake_pybind_connection.prepare_calls] == [ + "RETURN $value", + "RETURN $other", + ] + + +def test_pybind_no_parameter_query_skips_prepare_cache( + fake_pybind_connection: _FakePybindConnection, +) -> None: + conn = lb.Connection(_FakeDatabase()) + + conn.execute("RETURN 1") + + assert fake_pybind_connection.prepare_calls == [] + assert fake_pybind_connection.query_calls == ["RETURN 1"] + + +def test_pybind_close_clears_implicit_prepare_cache( + fake_pybind_connection: _FakePybindConnection, +) -> None: + conn = lb.Connection(_FakeDatabase()) + + conn.execute("RETURN $value", {"value": 1}) + + assert set(conn._pybind_implicit_prepared_cache) == {"RETURN $value"} + + conn.close() + + assert conn._pybind_implicit_prepared_cache == {} + assert fake_pybind_connection.closed is True From 7d514fa5a8ef49a8fed88fcc2fe73f5f4568408f Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 09:26:08 -0700 Subject: [PATCH 2/8] Fix test_pybind_implicit_prepare_cache for C-API CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs fixed in the test: 1. Add close() method to _FakeResult — QueryResult.close() calls self._query_result.close() internally; the missing method caused AttributeError during teardown on C-API CI. 2. Fix _get_pybind_connection monkeypatch to store fake_pybind in self._py_connection — the bare lambda returned fake_pybind but never stored it, so Connection.close() never called fake_pybind_connection.close(), causing 'assert fake_pybind_connection.closed is True' to fail. The test runs on both pybind and C-API CI jobs. These changes ensure it passes on C-API without affecting the pybind behavior. --- test/test_pybind_implicit_prepare_cache.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/test/test_pybind_implicit_prepare_cache.py b/test/test_pybind_implicit_prepare_cache.py index 7df3719..cb740de 100644 --- a/test/test_pybind_implicit_prepare_cache.py +++ b/test/test_pybind_implicit_prepare_cache.py @@ -14,6 +14,9 @@ def isSuccess(self) -> bool: def hasNextQueryResult(self) -> bool: return False + def close(self) -> None: + return + class _FakePreparedStatement: def __init__(self, query: str, parameters: dict[str, object]): @@ -83,9 +86,13 @@ def fake_pybind_connection(monkeypatch: pytest.MonkeyPatch) -> _FakePybindConnec "init_connection", lambda self: setattr(self, "_connection", _FakeBackendConnection()), ) - monkeypatch.setattr( - lb.Connection, "_get_pybind_connection", lambda self: fake_pybind - ) + + def _get_pybind_connection(self: lb.Connection) -> _FakePybindConnection: + if self._py_connection is None: + self._py_connection = fake_pybind + return self._py_connection + + monkeypatch.setattr(lb.Connection, "_get_pybind_connection", _get_pybind_connection) return fake_pybind From 34500cede83730bbae3edba9eb47c40718cf4824 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 10:45:22 -0700 Subject: [PATCH 3/8] Fix implicit prepared-statement cache (pybind + C-API) Two failure modes were present in the implicit prepared-statement cache on Connection: 1. C-API path: every parameterized execute() called _prepare() which triggered a fresh lbug_connection_prepare, never reusing the previously prepared statement. RSS would grow without bound on workloads that re-issued the same query. 2. pybind path: the cached PreparedStatement was reused, but when a caller's parameter *shape* changed (e.g. MAP -> STRUCT, or two STRUCTs with different field names), the upstream C++ plan (cachedPreparedStatementManager) stayed bound to the first shape's plan while the prepared statement's parameterMap was rebound in place. canReuseCachedPlanWith then returned true on the third call and the executor ran the stale plan against the new shape, segfaulting in Value::copyFromRowLayout. Reproduced by test_parameter.py::test_dict_conversion (MAP, STRUCT, STRUCT). The C-API's "fix" of binding to a fresh temp per execute (commit 4afc0ad) traded correctness for cost: it always re-prepares, defeating the cache entirely. Fix: * Cache key is (query, param_signature) instead of just query. The signature is a structural hash of the parameter *type* (not values): MAP vs STRUCT, STRUCT field names, MAP key/value element types. Two callers passing the same query with different shapes get independent prepared statements whose plans stay consistent with their bound values, avoiding the upstream mismatch. * Per-connection RLock (_prepared_cache_lock) serializes prepare + bind + execute so concurrent threads/async tasks cannot tear the cached entry's mutable bound state. The C++ side already takes its own mtx inside executeWithParams; the Python lock only needs to cover the C-API _bound_values map and the cache itself. * close() walks the cache and calls close() on each entry to destroy the C++ lbug_prepared_statement; the C-API wrapper has no __del__ so the entries would otherwise leak their _prepared_statement and _bound_values C++ allocations. * _get_or_prepare_pybind_statement and _execute_with_pybind hold the lock around prepare+bind+execute. * Tests: - test_pybind_implicit_prepare_cache.py: updated for the new tuple cache key and added two tests that pin the new shape-segregation behavior (3 distinct shapes -> 3 prepares; same shape, different values -> 1 prepare). - test_pybind_implicit_prepare_cache_threading.py: NEW. Reproduces the corruption in the pybind backend with 2/8/32 threads on a single shared connection. Currently FAILS on the pybind backend (to be debugged separately) and PASSES on the C-API backend. --- src_py/connection.py | 172 ++++++++++++++++-- test/test_pybind_implicit_prepare_cache.py | 51 +++++- ...pybind_implicit_prepare_cache_threading.py | 67 +++++++ 3 files changed, 275 insertions(+), 15 deletions(-) create mode 100644 test/test_pybind_implicit_prepare_cache_threading.py diff --git a/src_py/connection.py b/src_py/connection.py index 98e6097..92638fa 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -1,8 +1,10 @@ from __future__ import annotations +import contextlib import inspect import json import re +import threading import uuid import warnings from typing import TYPE_CHECKING, Any @@ -27,6 +29,94 @@ from typing_extensions import Self +def _capi_value_signature(value: Any) -> tuple: + """ + Compute a stable, hashable signature for a single parameter value. + + The signature mirrors the *structural* type that ``lbug`` will infer + when binding the parameter: for a Python ``dict`` that has ``"key"`` + and ``"value"`` list fields of equal length, the binder produces a + ``MAP(K, V)`` and the signature records the element type of both + lists. Any other string-keyed ``dict`` produces a ``STRUCT`` and the + signature records each field's name and child signature. This + granularity is required because two STRUCTs with different field + names produce different C++ plans, and routing a second shape + through the first shape's cached plan would misread fields. + + The signature is independent of *value contents* (so callers passing + the same query with different numeric values still share the cache) + but is sensitive to anything that would change the binder's plan: + nesting, field names, and element type. + """ + if value is None: + return ("null",) + if isinstance(value, bool): + return ("bool",) + # ``bool`` is a subclass of ``int``; check it first. + if isinstance(value, int): + return ("int",) + if isinstance(value, float): + return ("float",) + if isinstance(value, (bytes, bytearray, memoryview)): + return ("blob",) + # ``CAPIJsonParameter`` is a frozen dataclass wrapping a JSON string. + # Treat it as a string for caching purposes. + if type(value).__name__ == "CAPIJsonParameter": + return ("string",) + if isinstance(value, str): + return ("string",) + if isinstance(value, dict): + if ( + set(value.keys()) == {"key", "value"} + and isinstance(value["key"], list) + and isinstance(value["value"], list) + and len(value["key"]) == len(value["value"]) + ): + # ``MAP`` - the binder treats the two parallel lists as the + # key and value columns. Record the per-list element type so + # that ``MAP(INT,INT)`` vs ``MAP(STRING,INT)`` get distinct + # plans. + key_sig = ( + ("list", ("null",)) + if not value["key"] + else ("list", _capi_value_signature(value["key"][0])) + ) + val_sig = ( + ("list", ("null",)) + if not value["value"] + else ("list", _capi_value_signature(value["value"][0])) + ) + return ("map", key_sig, val_sig) + if all(isinstance(k, str) for k in value): + # ``STRUCT`` - the binder infers one field per key. Field + # names matter because they appear in the plan's expressions + # and column metadata. + return ( + "struct", + tuple(sorted((k, _capi_value_signature(v)) for k, v in value.items())), + ) + return ("dict",) + if isinstance(value, (list, tuple)): + if not value: + return ("list", ("null",)) + return ("list", _capi_value_signature(value[0])) + return ("unknown", type(value).__name__) + + +def _capi_param_signature(parameters: dict[str, Any]) -> tuple: + """ + Return a hashable signature of parameter *types* (not values). + + The signature is independent of value contents but stable across + calls that pass the same kind of parameters. The keys are sorted so + that two callers that pass the same parameters in different orders + share a prepared statement. + """ + return tuple( + sorted((key, _capi_value_signature(val)) for key, val in parameters.items()) + ) + + class Connection: """Connection to a database.""" @@ -52,7 +142,19 @@ def __init__(self, database: Database, num_threads: int = 0): self._query_timeout_ms = 0 self._query_results: WeakSet[QueryResult] = WeakSet() self._capi_scan_tables: set[str] = set() - self._pybind_implicit_prepared_cache: dict[str, Any] = {} + # Implicit prepared-statement cache, shared by the pybind and + # C-API paths. The key is (query, param_signature) so that two + # callers passing the same query with *different* parameter + # shapes (e.g. MAP vs STRUCT, or two STRUCTs with different + # field names) get independent prepared statements. Each cached + # entry's plan stays consistent with its bound values, avoiding + # the upstream "cached plan vs rebound parameterMap" mismatch. + self._pybind_implicit_prepared_cache: dict[tuple[str, tuple], Any] = {} + # Serializes prepare / bind / execute on a single connection so + # that the cached entry's mutable bound state cannot be torn by + # concurrent callers (multi-threaded users or AsyncConnection's + # thread-pool). + self._prepared_cache_lock = threading.RLock() self.database._register_connection(self) self.init_connection() @@ -117,7 +219,23 @@ def close(self) -> None: for query_result in list(self._query_results): query_result.close() self._query_results.clear() - self._pybind_implicit_prepared_cache.clear() + # Destroy the C++ ``lbug_prepared_statement`` held by every cached + # entry before we drop the references. ``lbug_prepared_statement`` + # does not own a Python ``__del__`` so the entries would otherwise + # leak their ``_prepared_statement`` and ``_bound_values`` C++ + # allocations. Pybind entries are safe to drop because their state + # is held in a ``shared_ptr`` that self-cleans on refcount=0. + with self._prepared_cache_lock: + for prepared in self._pybind_implicit_prepared_cache.values(): + close_fn = getattr(prepared, "close", None) + if callable(close_fn): + with contextlib.suppress(RuntimeError): + # The C++ connection may already be torn down + # (close ordering between connection and cache). + # Best-effort: the entries are about to be dropped + # anyway. + close_fn() + self._pybind_implicit_prepared_cache.clear() if self._connection is not None and not self.database.is_closed: self._connection.close() @@ -462,10 +580,11 @@ def _execute_with_pybind( return py_connection.query(query) query, parameters = self._normalize_parameters_for_pybind(query, parameters) - prepared = self._get_or_prepare_pybind_statement( - py_connection, query, parameters - ) - return py_connection.execute(prepared, parameters) + with self._prepared_cache_lock: + prepared = self._get_or_prepare_pybind_statement( + py_connection, query, parameters + ) + return py_connection.execute(prepared, parameters) def _get_or_prepare_pybind_statement( self, @@ -473,10 +592,26 @@ def _get_or_prepare_pybind_statement( query: str, parameters: dict[str, Any], ) -> Any: - prepared = self._pybind_implicit_prepared_cache.get(query) + # Caller must hold ``self._prepared_cache_lock``. + cache_key = (query, _capi_param_signature(parameters)) + prepared = self._pybind_implicit_prepared_cache.get(cache_key) if prepared is None: prepared = py_connection.prepare(query, parameters) - self._pybind_implicit_prepared_cache[query] = prepared + self._pybind_implicit_prepared_cache[cache_key] = prepared + return prepared + + def _get_or_prepare_capi_statement( + self, + query: str, + parameters: dict[str, Any], + ) -> PreparedStatement: + # Caller must hold ``self._prepared_cache_lock``. + cache_key = (query, _capi_param_signature(parameters)) + cached = self._pybind_implicit_prepared_cache.get(cache_key) + if cached is not None: + return cached + prepared = self._prepare(query, parameters) + self._pybind_implicit_prepared_cache[cache_key] = prepared return prepared def _maybe_raise_scan_unsupported_object(self, query: str) -> None: @@ -581,12 +716,21 @@ def execute( query, parameters = self._normalize_parameters_for_capi( query, parameters ) - prepared_statement = ( - self._prepare(query, parameters) if isinstance(query, str) else query - ) - query_result_internal = self._connection.execute( - prepared_statement._prepared_statement, parameters - ) + # Hold the lock across prepare + bind + execute so that + # concurrent callers (threads / async tasks) cannot tear the + # cached entry's mutable bound state. The C++ side has its + # own ``mtx`` around ``executeWithParams``; the lock here + # only protects the C-API ``_bound_values`` map and the + # cache itself. + with self._prepared_cache_lock: + prepared_statement = ( + self._get_or_prepare_capi_statement(query, parameters) + if isinstance(query, str) + else query + ) + query_result_internal = self._connection.execute( + prepared_statement._prepared_statement, parameters + ) if not query_result_internal.isSuccess(): raise RuntimeError(query_result_internal.getErrorMessage()) for table_name in scan_tables_to_drop: diff --git a/test/test_pybind_implicit_prepare_cache.py b/test/test_pybind_implicit_prepare_cache.py index cb740de..f98aedf 100644 --- a/test/test_pybind_implicit_prepare_cache.py +++ b/test/test_pybind_implicit_prepare_cache.py @@ -148,9 +148,58 @@ def test_pybind_close_clears_implicit_prepare_cache( conn.execute("RETURN $value", {"value": 1}) - assert set(conn._pybind_implicit_prepared_cache) == {"RETURN $value"} + # Cache key is now (query, parameter-type signature) so that a + # single query with different parameter shapes (e.g. MAP vs STRUCT) + # does not share a prepared statement. + assert len(conn._pybind_implicit_prepared_cache) == 1 + cached_key = next(iter(conn._pybind_implicit_prepared_cache)) + assert cached_key[0] == "RETURN $value" conn.close() assert conn._pybind_implicit_prepared_cache == {} assert fake_pybind_connection.closed is True + + +def test_pybind_implicit_prepare_segregates_by_parameter_shape( + fake_pybind_connection: _FakePybindConnection, +) -> None: + """The same query with different parameter shapes must not share a plan.""" + conn = lb.Connection(_FakeDatabase()) + + # Same keys ``key``/``value`` with matching list lengths -> MAP. + conn.execute("RETURN $st", {"st": {"key": [1, 2, 3], "value": [3, 7, 98]}}) + # First key is ``key1`` -> STRUCT. + conn.execute("RETURN $st", {"st": {"key1": [1, 2, 3], "value": [3, 7, 98]}}) + # ``key``/``value`` but mismatched list lengths -> STRUCT (different + # signature from the first call AND from the previous STRUCT because + # the field names differ). + conn.execute("RETURN $st", {"st": {"key": [1, 2], "value": [3, 7, 98, 4]}}) + + # 3 distinct structural shapes -> 3 separate prepared statements. + # Sharing any of them would either misbind the parameter or reroute + # a new shape through a stale plan. + assert len(fake_pybind_connection.prepare_calls) == 3 + assert all(call[0] == "RETURN $st" for call in fake_pybind_connection.prepare_calls) + assert [call[1] for call in fake_pybind_connection.execute_calls] == [ + {"st": {"key": [1, 2, 3], "value": [3, 7, 98]}}, + {"st": {"key1": [1, 2, 3], "value": [3, 7, 98]}}, + {"st": {"key": [1, 2], "value": [3, 7, 98, 4]}}, + ] + + +def test_pybind_implicit_prepare_reuses_same_struct_shape( + fake_pybind_connection: _FakePybindConnection, +) -> None: + """Two calls that produce the same structural signature reuse the cache.""" + conn = lb.Connection(_FakeDatabase()) + + # Same fields, different values -> identical signature. + conn.execute("RETURN $st", {"st": {"key1": [1, 2, 3], "value": [3, 7, 98]}}) + conn.execute("RETURN $st", {"st": {"key1": [4, 5, 6], "value": [7, 8, 9]}}) + + assert len(fake_pybind_connection.prepare_calls) == 1 + assert [call[1] for call in fake_pybind_connection.execute_calls] == [ + {"st": {"key1": [1, 2, 3], "value": [3, 7, 98]}}, + {"st": {"key1": [4, 5, 6], "value": [7, 8, 9]}}, + ] diff --git a/test/test_pybind_implicit_prepare_cache_threading.py b/test/test_pybind_implicit_prepare_cache_threading.py new file mode 100644 index 0000000..6c0fdd6 --- /dev/null +++ b/test/test_pybind_implicit_prepare_cache_threading.py @@ -0,0 +1,67 @@ +""" +Thread-safety regression tests for the implicit prepared-statement cache. + +The cache lives on ``Connection._pybind_implicit_prepared_cache`` and is +shared by the pybind and C-API execute paths. When two threads execute +the same parameterized query on the same connection, the cached +prepared statement is shared; if its mutable bound state is torn by +concurrent bind calls, the executor reads garbage and returns the wrong +value. These tests reproduce that scenario and would catch a +regression where the cache is reintroduced without proper +serialization. +""" + +from __future__ import annotations + +import threading + +import ladybug as lb +import pytest + + +@pytest.mark.parametrize( + ("num_threads", "iters"), + [ + # Light load - already covered by the async tests in spirit, but + # we re-test on a synchronous shared connection to make sure the + # per-connection lock covers the multi-threaded case. + (2, 5), + # Moderate load - the smallest case where the race used to + # surface (1 thread never reproduces it). + (8, 5), + # Heavy load - the original failure mode. + (32, 5), + ], +) +def test_shared_connection_concurrent_same_query(num_threads: int, iters: int) -> None: + """Many threads on a single connection must each see its own bound value.""" + db = lb.Database(":memory:", buffer_pool_size=2**28) + conn = lb.Connection(db) + + errors: list[tuple[int, list]] = [] + errors_lock = threading.Lock() + + def worker(base: int) -> None: + for i in range(iters): + v = base + i + result = conn.execute("RETURN $n", {"n": v}) + try: + got = result.get_next() + finally: + result.close() + if got != [v]: + with errors_lock: + errors.append((v, got)) + + threads = [ + threading.Thread(target=worker, args=(i * 1000,)) for i in range(num_threads) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert errors == [], ( + f"{len(errors)}/{num_threads * iters} executions returned the " + f"wrong value. First few: {errors[:5]}" + ) From 2cde7382e886e30066478c41faf16d4ec621e0e4 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 10:56:58 -0700 Subject: [PATCH 4/8] Fix pybind prepared cache signatures --- src_py/connection.py | 105 ++++++++++++++++++++- test/test_pybind_implicit_prepare_cache.py | 17 ++++ 2 files changed, 117 insertions(+), 5 deletions(-) diff --git a/src_py/connection.py b/src_py/connection.py index 92638fa..806f47a 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -117,6 +117,100 @@ def _capi_param_signature(parameters: dict[str, Any]) -> tuple: ) +def _pybind_int_signature(value: int) -> tuple[str]: + if -(2**7) <= value <= 2**7 - 1: + return ("int8",) + if 0 <= value <= 2**8 - 1: + return ("uint8",) + if -(2**15) <= value <= 2**15 - 1: + return ("int16",) + if 0 <= value <= 2**16 - 1: + return ("uint16",) + if -(2**31) <= value <= 2**31 - 1: + return ("int32",) + if 0 <= value <= 2**32 - 1: + return ("uint32",) + return ("int64",) + + +def _pybind_homogeneous_list_signature( + value: list[Any] | tuple[Any, ...], +) -> tuple | None: + first_non_null = next((item for item in value if item is not None), None) + if first_non_null is None: + return ("list", ("any",)) + if not isinstance(first_non_null, (bool, int, float)): + return None + first_type = type(first_non_null) + if any(item is not None and type(item) is not first_type for item in value): + return None + if isinstance(first_non_null, bool): + return ("list", ("bool",)) + if isinstance(first_non_null, int): + return ("list", ("int64",)) + return ("list", ("double",)) + + +def _pybind_value_signature(value: Any) -> tuple: + """ + Compute the parameter type signature used by the pybind binder. + + This intentionally differs from the C-API signature for Python ints: + pybind narrows scalar ints to INT8/UINT8/INT16/... based on the value + seen during prepare. Reusing a prepared statement across those + widths corrupts later executions because updateParameter() mutates + the value without revalidating the logical type. + """ + if value is None: + return ("any",) + if isinstance(value, bool): + return ("bool",) + if isinstance(value, int): + return _pybind_int_signature(value) + if isinstance(value, float): + return ("double",) + if isinstance(value, (bytes, bytearray, memoryview)): + return ("blob",) + if isinstance(value, str): + return ("string",) + module_name = type(value).__module__ + if module_name.startswith(("pandas", "polars", "pyarrow")): + return ("pointer", module_name, type(value).__name__, id(value)) + if isinstance(value, dict): + items = list(value.items()) + if ( + len(items) == 2 + and items[0][0] == "key" + and items[1][0] == "value" + and isinstance(items[0][1], list) + and isinstance(items[1][1], list) + and len(items[0][1]) == len(items[1][1]) + ): + return ( + "map", + _pybind_value_signature(items[0][1])[1], + _pybind_value_signature(items[1][1])[1], + ) + return ( + "struct", + tuple((str(k), _pybind_value_signature(v)) for k, v in items), + ) + if isinstance(value, (list, tuple)): + homogeneous = _pybind_homogeneous_list_signature(value) + if homogeneous is not None: + return homogeneous + if not value: + return ("list", ("any",)) + return ("list", _pybind_value_signature(value[0])) + return ("unknown", type(value).__name__) + + +def _pybind_param_signature(parameters: dict[str, Any]) -> tuple: + return tuple( + sorted((key, _pybind_value_signature(val)) for key, val in parameters.items()) + ) + + class Connection: """Connection to a database.""" @@ -143,12 +237,13 @@ def __init__(self, database: Database, num_threads: int = 0): self._query_results: WeakSet[QueryResult] = WeakSet() self._capi_scan_tables: set[str] = set() # Implicit prepared-statement cache, shared by the pybind and - # C-API paths. The key is (query, param_signature) so that two + # C-API paths. The key is (query, backend_param_signature) so that two # callers passing the same query with *different* parameter # shapes (e.g. MAP vs STRUCT, or two STRUCTs with different - # field names) get independent prepared statements. Each cached - # entry's plan stays consistent with its bound values, avoiding - # the upstream "cached plan vs rebound parameterMap" mismatch. + # field names) or pybind-native integer widths get independent + # prepared statements. Each cached entry's plan stays consistent + # with its bound values, avoiding the upstream "cached plan vs + # rebound parameterMap" mismatch. self._pybind_implicit_prepared_cache: dict[tuple[str, tuple], Any] = {} # Serializes prepare / bind / execute on a single connection so # that the cached entry's mutable bound state cannot be torn by @@ -593,7 +688,7 @@ def _get_or_prepare_pybind_statement( parameters: dict[str, Any], ) -> Any: # Caller must hold ``self._prepared_cache_lock``. - cache_key = (query, _capi_param_signature(parameters)) + cache_key = (query, _pybind_param_signature(parameters)) prepared = self._pybind_implicit_prepared_cache.get(cache_key) if prepared is None: prepared = py_connection.prepare(query, parameters) diff --git a/test/test_pybind_implicit_prepare_cache.py b/test/test_pybind_implicit_prepare_cache.py index f98aedf..e0466ca 100644 --- a/test/test_pybind_implicit_prepare_cache.py +++ b/test/test_pybind_implicit_prepare_cache.py @@ -130,6 +130,23 @@ def test_pybind_implicit_prepare_does_not_share_different_queries( ] +def test_pybind_implicit_prepare_segregates_by_integer_width( + fake_pybind_connection: _FakePybindConnection, +) -> None: + conn = lb.Connection(_FakeDatabase()) + + conn.execute("RETURN $value", {"value": 1}) + conn.execute("RETURN $value", {"value": 1000}) + conn.execute("RETURN $value", {"value": 2}) + + assert len(fake_pybind_connection.prepare_calls) == 2 + assert [call[1] for call in fake_pybind_connection.execute_calls] == [ + {"value": 1}, + {"value": 1000}, + {"value": 2}, + ] + + def test_pybind_no_parameter_query_skips_prepare_cache( fake_pybind_connection: _FakePybindConnection, ) -> None: From 7ef55e863bf605ad2ad2515384d414684789ff27 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 11:25:17 -0700 Subject: [PATCH 5/8] fix: test failures --- src_py/prepared_statement.py | 21 +++++++++++++++++++ ...pybind_implicit_prepare_cache_threading.py | 5 ++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src_py/prepared_statement.py b/src_py/prepared_statement.py index 25efe9a..de997db 100644 --- a/src_py/prepared_statement.py +++ b/src_py/prepared_statement.py @@ -33,6 +33,27 @@ def __init__( self._prepared_statement = connection._connection.prepare(query, parameters) self._connection = connection + def close(self) -> None: + """ + Release the underlying C-API prepared statement resources. + + The C-API ``PreparedStatement`` (from ``_lbug_capi.py``) holds a + ``lbug_prepared_statement`` C struct that must be destroyed explicitly; + it is NOT garbage-collected when the Python wrapper is dropped. + The pybind variant is managed by ``shared_ptr`` and can be dropped + without an explicit call. + + ``Connection.close()`` iterates the implicit prepared-statement cache + and calls this method on every cached entry so that resources are + freed when the connection is closed. + """ + close_fn = getattr(self._prepared_statement, "close", None) + if callable(close_fn): + close_fn() + + def __del__(self) -> None: + self.close() + def is_success(self) -> bool: """ Check if the prepared statement is successfully prepared. diff --git a/test/test_pybind_implicit_prepare_cache_threading.py b/test/test_pybind_implicit_prepare_cache_threading.py index 6c0fdd6..bd4ef96 100644 --- a/test/test_pybind_implicit_prepare_cache_threading.py +++ b/test/test_pybind_implicit_prepare_cache_threading.py @@ -35,7 +35,10 @@ ) def test_shared_connection_concurrent_same_query(num_threads: int, iters: int) -> None: """Many threads on a single connection must each see its own bound value.""" - db = lb.Database(":memory:", buffer_pool_size=2**28) + # Use an explicit max_db_size (1 GB) so the C-API backend does not default + # to the library's 8 TB mmap region, which fails on CI runners with tight + # virtual-address limits. + db = lb.Database(":memory:", buffer_pool_size=2**28, max_db_size=2**30) conn = lb.Connection(db) errors: list[tuple[int, list]] = [] From 554520447e23a7805c916b1b0682ac9199296215 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 11:56:09 -0700 Subject: [PATCH 6/8] test_fsm: exclude INDEX pages from *_delete page-range checks test_fsm_reclaim_node_table_delete and test_fsm_reclaim_rel_table_delete were failing because get_used_page_ranges() pulled in PK index pages (start_page_idx 31..51 in the vPerson fixture) and asserted they were a subset of the post-checkpoint free pages. match ... delete only reclaims data pages; the table and its persistent hash index stay around for future inserts, so those pages legitimately remain in use. Add get_used_node_data_page_ranges() that filters storage_info by table_type = 'NODE' and switch the two _delete tests to use it. The drop-table tests keep the original get_used_page_ranges() since drop removes the index along with the table. Likely broken by ladybug 9cd338e3f ('Fix hash index storage accounting and show built-in PK indexes'), which gave the hash index a getStorageEntries() implementation -- combined with 8a8c0f13f ('Report ART index storage in info UDFs') that wired appendStorageInfoForIndex into storage_info, this caused PK index pages to start showing up in storage_info output for the first time. The test was originally added in 8a55bc6 / C++ eaf97b755 ('Reclaim pages for fully-deleted node groups') when storage_info returned only NODE/REL data pages. --- test/test_fsm.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/test/test_fsm.py b/test/test_fsm.py index fba92e7..7b9855f 100644 --- a/test/test_fsm.py +++ b/test/test_fsm.py @@ -24,6 +24,24 @@ def get_used_page_ranges(conn, table, column=None): return used_pages +def get_used_node_data_page_ranges(conn, table): + """ + Like get_used_page_ranges, but excludes INDEX pages. + + Useful for `match ... delete` scenarios where the table and its indexes persist + while only the data pages are reclaimed. + """ + used_pages = [] + storage_info = conn.execute( + f'call storage_info("{table}") where table_type = "NODE" return start_page_idx, num_pages' + ) + while storage_info.has_next(): + cur_tuple = storage_info.get_next() + if cur_tuple[1] > 0: + used_pages.append(cur_tuple) + return used_pages + + def get_total_used_pages(conn): return conn.execute("call file_info() return num_pages").get_next()[0] @@ -181,7 +199,7 @@ def test_fsm_reclaim_node_table_recopy(fsm_node_table_setup) -> None: def test_fsm_reclaim_node_table_delete(fsm_node_table_setup) -> None: _, conn = fsm_node_table_setup - used_pages = get_used_page_ranges(conn, "person") + used_pages = get_used_node_data_page_ranges(conn, "person") conn.execute("match (p:person) delete p") prevent_data_file_truncation(conn) conn.execute("checkpoint") @@ -201,7 +219,7 @@ def test_fsm_reclaim_rel_table(fsm_rel_table_setup) -> None: def test_fsm_reclaim_rel_table_delete(fsm_node_table_setup) -> None: _, conn = fsm_node_table_setup - used_pages = get_used_page_ranges(conn, "person") + used_pages = get_used_node_data_page_ranges(conn, "person") conn.execute("match (p:person) delete p") prevent_data_file_truncation(conn) conn.execute("checkpoint") From 530c5ec3f7f020103d39d70e491e69c5923e6dfd Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 09:12:48 -0700 Subject: [PATCH 7/8] fix: memory corruption due to accessing out of bounds output vector index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In numpy_scan.cpp, the OBJECT/STRING scan path uses pos = i + offset (the absolute row index in the source array) as the output vector index for setNull() when handling None or NaN values: ```cpp // BEFORE (buggy): outputVector->setNull(pos, true /* isNull */); // pos = i + offset ``` The output vector is sized for DEFAULT_VECTOR_CAPACITY (= 2048) entries and expects 0-based indices within each batch. But pos is the absolute dataframe row index. When the COPY spans multiple batches (i.e., when offset >= DEFAULT_VECTOR_CAPACITY), pos exceeds the null mask buffer, causing a heap buffer overrun → the segfault / free(): invalid next size / munmap_chunk(): invalid pointer crashes seen on Linux. --- src_cpp/numpy/numpy_scan.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src_cpp/numpy/numpy_scan.cpp b/src_cpp/numpy/numpy_scan.cpp index 0dc885c..c172205 100644 --- a/src_cpp/numpy/numpy_scan.cpp +++ b/src_cpp/numpy/numpy_scan.cpp @@ -172,7 +172,7 @@ void NumpyScan::scan(PandasColumnBindData* bindData, uint64_t count, uint64_t of !py::isinstance(val)) { if (val == Py_None || (py::isinstance(val) && std::isnan(PyFloat_AsDouble(val)))) { - outputVector->setNull(pos, true /* isNull */); + outputVector->setNull(i, true /* isNull */); continue; } if (!py::isinstance(val)) { From bc5efe137e4e4fb2c7fc2e96960ee95293ac6db2 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Thu, 9 Jul 2026 12:19:27 -0700 Subject: [PATCH 8/8] test_pyarrow_primitive: avoid NaN in float32 test data and make tables_equal NaN-aware The test generated float32 values by taking random 32-bit patterns and reinterpreting them as IEEE 754 floats. About 0.4% of those patterns are NaN (all-ones exponent, nonzero mantissa), so the floatcol contained dozens of NaN values at the larger scale factors. pyarrow.Array == Array follows the Arrow spec (NaN != NaN), so tables_equal()'s element-wise comparison would report two tables that differ only by NaN positions as not-equal. Locally, pa.Table.from_pandas converts NaN to None when building the expected patable, masking the issue. In CI (different pyarrow/pandas versions) the NaN survives in both the expected and scanned tables and the comparison fails with 'tables are not equal'. Two-part fix: - generate_primitive() for float32 now retries until it gets a non-NaN bit pattern. Subnormals, +/-Inf, signed zero, etc. are all still exercised, just not NaN. - tables_equal() routes floating columns through numpy.array_equal(equal_nan=True), which is the correct semantic for a round-trip scan test. Non-floating columns keep the cheap pyarrow.Array != Array check. Related: ladybug #647 (data corruption in the DataFrame scan path on Linux, fixed in 530c5ec for numpy_scan.cpp). This is a different bug in the test itself, not in the scan path. --- test/test_scan_pandas_pyarrow.py | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/test/test_scan_pandas_pyarrow.py b/test/test_scan_pandas_pyarrow.py index 7cfbab4..539f356 100644 --- a/test/test_scan_pandas_pyarrow.py +++ b/test/test_scan_pandas_pyarrow.py @@ -24,10 +24,16 @@ def generate_primitive(dtype): if dtype.startswith("uint64"): return random.randrange(0, 18446744073709551615) if dtype.startswith("float32"): - random_bits = random.getrandbits(32) - random_bytes = struct.pack("