From a7df365f54ad11283425f077b020af3d902f13d6 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Sun, 6 Sep 2026 15:14:33 -0700 Subject: [PATCH] Fix 0-row result when re-executing a cached prepared statement scanning a python object Re-executing the same statement that scans a pandas/polars/pyarrow object (e.g. two LOAD FROM df queries on one connection) returned 0 rows on every execution after the first. Root cause: with the physical-plan cache enabled, the cached operator tree template is cloned via TableFunctionCall::copy(), which shares the same TableFuncSharedState instance between the template and its copies. The first execution advanced PyArrowTableScanSharedState::currentChunk to the end of the chunk list, and the shared state's resetState() was the base-class no-op, so prepareForReuse() never rewound the cursor and subsequent executions scanned zero chunks. Fix: override resetState() in PyArrowTableScanSharedState to reset currentChunk to 0. Includes a regression test that executes the same LOAD FROM df query three times and asserts each execution returns the full row count. --- src_cpp/include/pyarrow/pyarrow_scan.h | 5 +++++ test/test_scan_pandas_pyarrow.py | 12 ++++++++++++ 2 files changed, 17 insertions(+) diff --git a/src_cpp/include/pyarrow/pyarrow_scan.h b/src_cpp/include/pyarrow/pyarrow_scan.h index 7e0ef4a..e4dbec9 100644 --- a/src_cpp/include/pyarrow/pyarrow_scan.h +++ b/src_cpp/include/pyarrow/pyarrow_scan.h @@ -26,6 +26,11 @@ struct PyArrowTableScanSharedState final : public function::TableFuncSharedState : TableFuncSharedState{numRows}, chunks{std::move(chunks)}, currentChunk{0} {} ArrowArrayWrapper* getNextChunk(); + + // TableFunctionCall::copy() (used by the physical-plan cache on prepared + // statement re-execution) shares the same sharedState instance, so the + // chunk cursor must be rewound here or a re-executed scan yields 0 rows. + void resetState() override { currentChunk = 0; } }; struct PyArrowTableScanFunctionData final : public function::TableFuncBindData { diff --git a/test/test_scan_pandas_pyarrow.py b/test/test_scan_pandas_pyarrow.py index 539f356..d18568e 100644 --- a/test/test_scan_pandas_pyarrow.py +++ b/test/test_scan_pandas_pyarrow.py @@ -141,6 +141,18 @@ def test_pyarrow_primitive(conn_db_empty: ConnDB) -> None: pyarrow_test_helper(establish_connection, sf, thread) +def test_pyarrow_scan_repeated_execution(conn_db_in_mem: ConnDB) -> None: + # Regression test: re-executing the same (implicitly cached) prepared + # statement that scans a python object returned 0 rows on the second run. + # The cached physical plan shares the scan's shared state, whose chunk + # cursor was never rewound between executions. + conn, _ = conn_db_in_mem + df = pd.DataFrame({"a": pd.Series([1, 2, 3, None], dtype="int32[pyarrow]")}) + for _ in range(3): + result = conn.execute("LOAD FROM df RETURN count(*)") + assert result.get_next()[0] == 4 + + def test_pyarrow_time(conn_db_readonly: ConnDB) -> None: conn, _ = conn_db_readonly col1 = pa.array([1000123, 2000123, 3000123], type=pa.duration("s"))