(arrow-adbc) branch main updated: fix(python/adbc_driver_manager): return 'real' reader in fetch_record_batch (#1530)

[li...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git


The following commit(s) were added to refs/heads/main by this push:
     new 0d20d48f fix(python/adbc_driver_manager): return 'real' reader in fetch_record_batch (#1530)
0d20d48f is described below

commit 0d20d48f0df1800bb75ce077ad6a7be9cf754543
Author: David Li <li...@gmail.com>
AuthorDate: Thu Feb 8 09:57:51 2024 -0500

    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 |  5 ++++-
 python/adbc_driver_manager/tests/test_dbapi.py          | 10 ++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
index ee6f318d..c296a257 100644
--- a/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
+++ b/python/adbc_driver_manager/adbc_driver_manager/dbapi.py
@@ -1089,7 +1089,10 @@ class Cursor(_Closeable):
                 "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
 
 
 # ----------------------------------------------------------
diff --git a/python/adbc_driver_manager/tests/test_dbapi.py b/python/adbc_driver_manager/tests/test_dbapi.py
index 20990eff..3cad3f64 100644
--- a/python/adbc_driver_manager/tests/test_dbapi.py
+++ b/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: