fix(python/adbc_driver_manager): return 'real' reader in fetch_record…

…_batch (#1530)

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 #1523.

python/adbc_driver_manager/adbc_driver_manager/dbapi.py

@@ -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
 
 
 # ----------------------------------------------------------

python/adbc_driver_manager/tests/test_dbapi.py

@@ -17,6 +17,7 @@
 
 import pandas
 import pyarrow
+import pyarrow.dataset
 import pytest
 from pandas.testing import assert_frame_equal
 
@@ -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: