fix(python/adbc_driver_manager): don't consume result for description (#3554)

Closes #3543.

Author: lidavidm

python/adbc_driver_manager/adbc_driver_manager/_dbapi_backend.py

@@ -28,6 +28,9 @@
 
 from . import _lib
 
+if typing.TYPE_CHECKING:
+    from typing_extensions import CapsuleType
+
 
 class DbapiBackend(abc.ABC):
     """
@@ -87,6 +90,27 @@ def convert_executemany_parameters(
         """
         ...
 
+    @abc.abstractmethod
+    def convert_description(self, schema: "CapsuleType") -> typing.List[typing.Tuple]:
+        """Convert a schema capsule into a DB-API description.
+
+        Parameters
+        ----------
+        schema
+            A PyCapsule of type "arrow_schema".
+
+        Returns
+        -------
+        description : list[tuple]
+            A DB-API description, as a list of 7-item tuples.
+
+        See Also
+        --------
+        https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
+
+        """
+        ...
+
     @abc.abstractmethod
     def import_array_stream(self, handle: _lib.ArrowArrayStreamHandle) -> typing.Any:
         """Import an Arrow stream."""
@@ -120,6 +144,12 @@ def convert_executemany_parameters(
             status_code=_lib.AdbcStatusCode.INVALID_STATE,
         )
 
+    def convert_description(self, schema: "CapsuleType") -> typing.List[typing.Tuple]:
+        raise _lib.ProgrammingError(
+            "This API requires PyArrow or another suitable backend to be installed",
+            status_code=_lib.AdbcStatusCode.INVALID_STATE,
+        )
+
     def import_array_stream(
         self, handle: _lib.ArrowArrayStreamHandle
     ) -> _lib.ArrowArrayStreamHandle:
@@ -174,6 +204,11 @@ def convert_executemany_parameters(
             cols, bind_by_name = param_iterable_to_dict(parameters)
             return polars.DataFrame(cols), bind_by_name
 
+        def convert_description(
+            self, schema: "CapsuleType"
+        ) -> typing.List[typing.Tuple]:
+            raise _lib.NotSupportedError("Polars does not support __arrow_c_schema__")
+
         def import_array_stream(
             self, handle: _lib.ArrowArrayStreamHandle
         ) -> typing.Any:
@@ -207,6 +242,14 @@ def convert_executemany_parameters(
             cols, bind_by_name = param_iterable_to_dict(parameters)
             return pyarrow.RecordBatch.from_pydict(cols), bind_by_name
 
+        def convert_description(
+            self, schema: "CapsuleType"
+        ) -> typing.List[typing.Tuple]:
+            s = pyarrow.Schema._import_from_c_capsule(schema)
+            return [
+                (field.name, field.type, None, None, None, None, None) for field in s
+            ]
+
         def import_array_stream(
             self, handle: _lib.ArrowArrayStreamHandle
         ) -> pyarrow.RecordBatchReader:

python/adbc_driver_manager/adbc_driver_manager/_lib.pxd

@@ -32,9 +32,9 @@ cdef extern from "arrow-adbc/adbc.h" nogil:
     cdef struct CArrowArray"ArrowArray":
         CArrowArrayRelease release
 
-    ctypedef int (*CArrowArrayStreamGetLastError)(void*)
+    ctypedef char* (*CArrowArrayStreamGetLastError)(void*)
     ctypedef int (*CArrowArrayStreamGetNext)(void*, CArrowArray*)
-    ctypedef char* (*CArrowArrayStreamGetSchema)(void*, CArrowSchema*)
+    ctypedef int (*CArrowArrayStreamGetSchema)(void*, CArrowSchema*)
     ctypedef void (*CArrowArrayStreamRelease)(void*)
 
     cdef struct CArrowArrayStream"ArrowArrayStream":

python/adbc_driver_manager/adbc_driver_manager/_lib.pyx

@@ -479,6 +479,28 @@ cdef class ArrowArrayStreamHandle:
             self.stream.release(&self.stream)
             self.stream.release = NULL
 
+    def __arrow_c_schema__(self) -> object:
+        """Get a PyCapsule without consuming this object."""
+        cdef const char* err = NULL
+
+        if not self.is_valid:
+            raise ValueError("ArrowArrayStreamHandle already consumed")
+
+        cdef CArrowSchema* allocated = <CArrowSchema*> malloc(sizeof(CArrowSchema))
+        allocated.release = NULL
+        capsule = PyCapsule_New(
+            <void*>allocated, "arrow_schema", &pycapsule_schema_deleter,
+        )
+        rc = self.stream.get_schema(&self.stream, allocated)
+        if rc != 0:
+            err = self.stream.get_last_error(&self.stream)
+            if err == NULL:
+                raise RuntimeError(f"Failed to get schema: ({rc})")
+            else:
+                s = err.decode()
+                raise RuntimeError(f"Failed to get schema: ({rc}) {s}")
+        return capsule
+
     def __arrow_c_stream__(self, requested_schema=None) -> object:
         """Consume this object to get a PyCapsule."""
         if requested_schema is not None:

python/adbc_driver_manager/adbc_driver_manager/dbapi.py

@@ -370,7 +370,7 @@ def cursor(
         adbc_stmt_kwargs : dict, optional
           ADBC-specific options to pass to the underlying ADBC statement.
         """
-        return Cursor(self, adbc_stmt_kwargs)
+        return Cursor(self, adbc_stmt_kwargs, dbapi_backend=self._backend)
 
     def rollback(self) -> None:
         """Explicitly rollback."""
@@ -624,9 +624,12 @@ def __init__(
         self,
         conn: Connection,
         adbc_stmt_kwargs: Optional[Dict[str, Any]] = None,
+        *,
+        dbapi_backend: Optional[_dbapi_backend.DbapiBackend] = None,
     ) -> None:
         # Must be at top in case __init__ is interrupted and then __del__ is called
         self._closed = True
+        self._backend = dbapi_backend or _dbapi_backend.default_backend()
         self._conn = conn
         self._stmt = _lib.AdbcStatement(conn._conn)
         self._closed = False
@@ -772,7 +775,7 @@ def execute(self, operation: Union[bytes, str], parameters=None) -> "Self":
         handle, self._rowcount = _blocking_call(
             self._stmt.execute_query, (), {}, self._stmt.cancel
         )
-        self._results = _RowIterator(self._stmt, handle)
+        self._results = _RowIterator(self._stmt, handle, self._backend)
         return self
 
     def executemany(self, operation: Union[bytes, str], seq_of_parameters) -> None:
@@ -1141,7 +1144,7 @@ def adbc_read_partition(self, partition: bytes) -> None:
             self._conn._conn.read_partition, (partition,), {}, self._stmt.cancel
         )
         self._rowcount = -1
-        self._results = _RowIterator(self._stmt, handle)
+        self._results = _RowIterator(self._stmt, handle, self._backend)
 
     @property
     def adbc_statement(self) -> _lib.AdbcStatement:
@@ -1261,8 +1264,8 @@ def fetch_arrow(self) -> _lib.ArrowArrayStreamHandle:
         Fetch the result as an object implementing the Arrow PyCapsule interface.
 
         This can only be called once.  It must be called before any other
-        method that inspect the data (e.g. description, fetchone,
-        fetch_arrow_table, etc.).  Once this is called, other methods that
+        method that consume data (e.g. fetchone, fetch_arrow_table, etc.;
+        description is allowed).  Once this is called, other methods that
         inspect the data may not be called.
 
         Notes
@@ -1285,10 +1288,14 @@ class _RowIterator(_Closeable):
     """Track state needed to iterate over the result set."""
 
     def __init__(
-        self, stmt: _lib.AdbcStatement, handle: _lib.ArrowArrayStreamHandle
+        self,
+        stmt: _lib.AdbcStatement,
+        handle: _lib.ArrowArrayStreamHandle,
+        dbapi_backend: _dbapi_backend.DbapiBackend,
     ) -> None:
         self._stmt = stmt
         self._handle: Optional[_lib.ArrowArrayStreamHandle] = handle
+        self._backend = dbapi_backend
         self._reader: Optional["_reader.AdbcRecordBatchReader"] = None
         self._current_batch = None
         self._next_row = 0
@@ -1321,10 +1328,16 @@ def reader(self) -> "_reader.AdbcRecordBatchReader":
 
     @property
     def description(self) -> List[tuple]:
-        return [
-            (field.name, field.type, None, None, None, None, None)
-            for field in self.reader.schema
-        ]
+        if self._handle is None:
+            # Invalid state, or already imported into the reader
+            # (we assume PyArrow here for now)
+            return [
+                (field.name, field.type, None, None, None, None, None)
+                for field in self.reader.schema
+            ]
+        else:
+            # Not yet imported into the reader.  Do not force consumption
+            return self._backend.convert_description(self._handle.__arrow_c_schema__())
 
     def fetchone(self) -> Optional[tuple]:
         if self._current_batch is None or self._next_row >= len(self._current_batch):

python/adbc_driver_manager/tests/test_dbapi.py

@@ -243,6 +243,30 @@ def test_query_fetch_arrow(sqlite):
             cur.fetch_arrow()
 
 
+@pytest.mark.sqlite
+def test_query_fetch_arrow_3543(sqlite):
+    # Regression test for https://github.com/apache/arrow-adbc/issues/3543
+    with sqlite.cursor() as cur:
+        cur.execute("SELECT 1, 'foo' AS foo, 2.0")
+
+        # This should not consume the result
+        assert cur.description == [
+            ("1", dbapi.NUMBER, None, None, None, None, None),
+            ("foo", dbapi.STRING, None, None, None, None, None),
+            ("2.0", dbapi.NUMBER, None, None, None, None, None),
+        ]
+
+        capsule = cur.fetch_arrow().__arrow_c_stream__()
+        reader = pyarrow.RecordBatchReader._import_from_c_capsule(capsule)
+        assert reader.read_all() == pyarrow.table(
+            {
+                "1": [1],
+                "foo": ["foo"],
+                "2.0": [2.0],
+            }
+        )
+
+
 @pytest.mark.sqlite
 def test_query_fetch_arrow_table(sqlite):
     with sqlite.cursor() as cur: