Fix 0-row result when re-executing a cached prepared statement scanning a python object - #54
Merged
Merged
Conversation
…ng 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
LOAD FROM dfqueries on one connection) returned 0 rows on every execution after the first. This broketest_pyarrow_primitive(its second, parameterizedLOAD FROM $dfexecution returned an empty table).TableFunctionCall::copy(), which shares the sameTableFuncSharedStateinstance between the template and its copies. The first execution advancedPyArrowTableScanSharedState::currentChunkto the end of the chunk list, andresetState()was the base-class no-op, soprepareForReuse()never rewound the cursor and subsequent executions scanned zero chunks.resetState()inPyArrowTableScanSharedStateto resetcurrentChunk = 0.Testing
test_pyarrow_scan_repeated_execution(executes the sameLOAD FROM df RETURN count(*)three times, asserts the full count each time). Verified it fails without the fix and passes with it.test_pyarrow_primitiveand the rest oftest_scan_pandas_pyarrow.pypass; CI will cover the full suite.CI note: this only touches the pybind extension header + tests; C++ change is one-line override in
src_cpp/include/pyarrow/pyarrow_scan.h.