fix(xorq): use WeakKeyDictionary for _expr_count_cache to prevent id() reuse (#889)

`_expr_count_cache` was keyed by `id(expr)`, which is the CPython memory
address. When a GC'd expression's address was reused by a new expression,
the cache returned a stale count for the new object.

Concretely: `test_filter_pushes_down` cached count=4 for a filtered expr;
after GC that address was reused by the 2-row search result in the next
test, causing `filtered_rows` to read 4 instead of 2.

Switching to `weakref.WeakKeyDictionary` (keyed by the expression object
directly) means entries are evicted when the expression is GC'd, eliminating
the false hit. Also update the `TestExprCountMemoization` assertions that
referenced `id(stub)` directly.

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: paddymul

buckaroo/xorq_buckaroo.py

@@ -12,6 +12,7 @@
 
 import logging
 import traceback
+import weakref
 from io import BytesIO
 from typing import Any
 
@@ -45,14 +46,13 @@ def _is_pandas(obj: Any) -> bool:
 # with one entry per unique expression object observed; entries become
 # unreachable when the expression is GC'd, so for a long-running
 # session it tracks the live set of expressions naturally.
-_expr_count_cache: dict[int, int] = {}
+_expr_count_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
 
 
 def _expr_count(expr_or_df: Any) -> int:
     if _is_pandas(expr_or_df):
         return len(expr_or_df)
-    key = id(expr_or_df)
-    cached = _expr_count_cache.get(key)
+    cached = _expr_count_cache.get(expr_or_df)
     if cached is not None:
         return cached
     try:
@@ -65,7 +65,7 @@ def _expr_count(expr_or_df: Any) -> int:
         # next call retries the backend.
         logger.exception("_expr_count: backend count failed; not caching")
         return 0
-    _expr_count_cache[key] = result
+    _expr_count_cache[expr_or_df] = result
     return result
 
 

tests/unit/test_xorq_buckaroo_widget.py

@@ -95,7 +95,7 @@ def test_expr_count_memoizes_repeat_calls(self):
         assert len(xorq_buckaroo._expr_count_cache) == 1, (
             f"expected exactly 1 cached entry for 3 calls on the same "
             f"expression; got {len(xorq_buckaroo._expr_count_cache)}. "
-            f"The cache key must be id(expr)."
+            f"The cache key must be the expression object itself."
         )
 
     def test_expr_count_separate_expressions_cache_separately(self):
@@ -159,17 +159,17 @@ def count(self):
         stub = _StubExpr()
 
         first = xorq_buckaroo._expr_count(stub)
-        assert id(stub) not in xorq_buckaroo._expr_count_cache, (
+        assert stub not in xorq_buckaroo._expr_count_cache, (
             "failed _expr_count call must not write to the cache; "
-            f"found cached value {xorq_buckaroo._expr_count_cache.get(id(stub))!r}"
+            f"found cached value {xorq_buckaroo._expr_count_cache.get(stub)!r}"
         )
 
         second = xorq_buckaroo._expr_count(stub)
         assert second == 42, (
             f"after backend recovery, _expr_count must return the real "
             f"count, not a cached failure sentinel; got {second!r}"
         )
-        assert xorq_buckaroo._expr_count_cache.get(id(stub)) == 42, (
+        assert xorq_buckaroo._expr_count_cache.get(stub) == 42, (
             "successful call following a failure must populate the cache "
             "with the real count"
         )