Skip to content

Commit

Permalink
fix(python/adbc_driver_manager): return 'real' reader in fetch_record…
Browse files Browse the repository at this point in the history
…_batch

We wrap RecordBatchReader to try to provide extra error metadata from
ADBC when using this API.  But PyArrow may try to extract the underlying
C++ arrow::RecordBatch pointer from the reader, which won't exist in our
wrapped reader.  For the public API, return the unwrapped reader
instead.

Fixes apache#1523.
  • Loading branch information
lidavidm committed Feb 8, 2024
1 parent f71e398 commit b01551d
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 1 deletion.
5 changes: 4 additions & 1 deletion python/adbc_driver_manager/adbc_driver_manager/dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -1089,7 +1089,10 @@ def fetch_record_batch(self) -> pyarrow.RecordBatchReader:
"Cannot fetch_record_batch() before execute()",
status_code=_lib.AdbcStatusCode.INVALID_STATE,
)
return self._results._reader
# XXX(https://github.com/apache/arrow-adbc/issues/1523): return the
# "real" PyArrow reader since PyArrow may try to poke the internal C++
# reader pointer
return self._results._reader._reader


# ----------------------------------------------------------
Expand Down
10 changes: 10 additions & 0 deletions python/adbc_driver_manager/tests/test_dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import pandas
import pyarrow
import pyarrow.dataset
import pytest
from pandas.testing import assert_frame_equal

Expand Down Expand Up @@ -351,6 +352,15 @@ def test_fetch_empty(sqlite):
assert cur.fetchall() == []


@pytest.mark.sqlite
def test_reader(sqlite, tmp_path) -> None:
# Regression test for https://github.com/apache/arrow-adbc/issues/1523
with sqlite.cursor() as cur:
cur.execute("SELECT 1")
reader = cur.fetch_record_batch()
pyarrow.dataset.write_dataset(reader, tmp_path, format="parquet")


@pytest.mark.sqlite
def test_prepare(sqlite):
with sqlite.cursor() as cur:
Expand Down

0 comments on commit b01551d

Please sign in to comment.