Skip to content

Commit dd285ce

Browse files
committed
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.
1 parent eed4d5d commit dd285ce

2 files changed

Lines changed: 17 additions & 0 deletions

File tree

‎src_cpp/include/pyarrow/pyarrow_scan.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ struct PyArrowTableScanSharedState final : public function::TableFuncSharedState
2626
: TableFuncSharedState{numRows}, chunks{std::move(chunks)}, currentChunk{0} {}
2727

2828
ArrowArrayWrapper* getNextChunk();
29+
30+
// TableFunctionCall::copy() (used by the physical-plan cache on prepared
31+
// statement re-execution) shares the same sharedState instance, so the
32+
// chunk cursor must be rewound here or a re-executed scan yields 0 rows.
33+
void resetState() override { currentChunk = 0; }
2934
};
3035

3136
struct PyArrowTableScanFunctionData final : public function::TableFuncBindData {

‎test/test_scan_pandas_pyarrow.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,18 @@ def test_pyarrow_primitive(conn_db_empty: ConnDB) -> None:
141141
pyarrow_test_helper(establish_connection, sf, thread)
142142

143143

144+
def test_pyarrow_scan_repeated_execution(conn_db_in_mem: ConnDB) -> None:
145+
# Regression test: re-executing the same (implicitly cached) prepared
146+
# statement that scans a python object returned 0 rows on the second run.
147+
# The cached physical plan shares the scan's shared state, whose chunk
148+
# cursor was never rewound between executions.
149+
conn, _ = conn_db_in_mem
150+
df = pd.DataFrame({"a": pd.Series([1, 2, 3, None], dtype="int32[pyarrow]")})
151+
for _ in range(3):
152+
result = conn.execute("LOAD FROM df RETURN count(*)")
153+
assert result.get_next()[0] == 4
154+
155+
144156
def test_pyarrow_time(conn_db_readonly: ConnDB) -> None:
145157
conn, _ = conn_db_readonly
146158
col1 = pa.array([1000123, 2000123, 3000123], type=pa.duration("s"))

0 commit comments

Comments
 (0)