1. fix tests 2. avoid clone 3. only return a single batch

Author: wangxiaoying

connectorx-python/connectorx/__init__.py

@@ -411,9 +411,27 @@ def read_sql(
             dd = try_import_module("dask.dataframe")
             df = dd.from_pandas(df, npartitions=1)
 
-    elif return_type in {"arrow", "polars", "arrow_record_batches"}:
+    elif return_type in {"arrow", "polars"}:
         try_import_module("pyarrow")
 
+        result = _read_sql(
+            conn,
+            "arrow",
+            queries=queries,
+            protocol=protocol,
+            partition_query=partition_query,
+            pre_execution_queries=pre_execution_queries,
+        )
+
+        df = reconstruct_arrow(result)
+        if return_type in {"polars"}:
+            pl = try_import_module("polars")
+            try:
+                df = pl.from_arrow(df)
+            except AttributeError:
+                # previous polars api (< 0.8.*) was pl.DataFrame.from_arrow
+                df = pl.DataFrame.from_arrow(df)
+    elif return_type in {"arrow_record_batches"}:
         record_batch_size = int(kwargs.get("record_batch_size", 10000))
         result = _read_sql(
             conn,
@@ -425,17 +443,7 @@ def read_sql(
             record_batch_size=record_batch_size
         )
 
-        if return_type == "arrow_record_batches":
-            df = reconstruct_arrow_rb(result)
-        else:
-            df = reconstruct_arrow(result)
-            if return_type in {"polars"}:
-                pl = try_import_module("polars")
-                try:
-                    df = pl.from_arrow(df)
-                except AttributeError:
-                    # previous polars api (< 0.8.*) was pl.DataFrame.from_arrow
-                    df = pl.DataFrame.from_arrow(df)
+        df = reconstruct_arrow_rb(result)
     else:
         raise ValueError(return_type)
 
@@ -455,11 +463,10 @@ def reconstruct_arrow_rb(results) -> pa.RecordBatchReader:
 
     def generate_batches(iterator) -> Iterator[pa.RecordBatch]:
         for rb_ptrs in iterator:
-            names, chunk_ptrs_list = rb_ptrs.to_ptrs()
-            for chunk_ptrs in chunk_ptrs_list:
-                yield pa.RecordBatch.from_arrays(
-                    [pa.Array._import_from_c(*col_ptr) for col_ptr in chunk_ptrs], names
-                )
+            chunk_ptrs = rb_ptrs.to_ptrs()
+            yield pa.RecordBatch.from_arrays(
+                [pa.Array._import_from_c(*col_ptr) for col_ptr in chunk_ptrs], names
+            )
 
     return pa.RecordBatchReader.from_batches(schema=schema, batches=generate_batches(results))
 

connectorx-python/src/arrow.rs

@@ -1,4 +1,5 @@
 use crate::errors::ConnectorXPythonError;
+use anyhow::anyhow;
 use arrow::record_batch::RecordBatch;
 use connectorx::source_router::SourceConn;
 use connectorx::{prelude::*, sql::CXQuery};
@@ -12,7 +13,7 @@ use std::sync::Arc;
 
 /// Python-exposed RecordBatch wrapper
 #[pyclass]
-pub struct PyRecordBatch(RecordBatch);
+pub struct PyRecordBatch(Option<RecordBatch>);
 
 /// Python-exposed iterator over RecordBatches
 #[pyclass(module = "connectorx")]
@@ -21,20 +22,22 @@ pub struct PyRecordBatchIterator(Box<dyn RecordBatchIterator>);
 #[pymethods]
 impl PyRecordBatch {
     pub fn num_rows(&self) -> usize {
-        self.0.num_rows()
+        self.0.as_ref().map_or(0, |rb| rb.num_rows())
     }
 
     pub fn num_columns(&self) -> usize {
-        self.0.num_columns()
+        self.0.as_ref().map_or(0, |rb| rb.num_columns())
     }
 
     #[throws(ConnectorXPythonError)]
-    pub fn to_ptrs<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> {
+    pub fn to_ptrs<'py>(&mut self, py: Python<'py>) -> Bound<'py, PyAny> {
+        // Convert the RecordBatch to a vector of pointers, once the RecordBatch is taken, it cannot be reached again.
+        let rb = self
+            .0
+            .take()
+            .ok_or_else(|| anyhow!("RecordBatch is None, cannot convert to pointers"))?;
         let ptrs = py.allow_threads(
-            || -> Result<(Vec<String>, Vec<Vec<(uintptr_t, uintptr_t)>>), ConnectorXPythonError> {
-                let rbs = vec![self.0.clone()];
-                Ok(to_ptrs(rbs))
-            },
+            || -> Result<Vec<(uintptr_t, uintptr_t)>, ConnectorXPythonError> { Ok(to_ptrs_rb(rb)) },
         )?;
         let obj: PyObject = ptrs.into_py(py);
         obj.into_bound(py)
@@ -65,7 +68,7 @@ impl PyRecordBatchIterator {
     ) -> PyResult<Option<Py<PyRecordBatch>>> {
         match slf.0.next_batch() {
             Some(rb) => {
-                let wrapped = PyRecordBatch(rb);
+                let wrapped = PyRecordBatch(Some(rb));
                 let py_obj = Py::new(py, wrapped)?;
                 Ok(Some(py_obj))
             }
@@ -118,6 +121,24 @@ pub fn get_arrow_rb_iter<'py>(
     obj.into_bound(py)
 }
 
+pub fn to_ptrs_rb(rb: RecordBatch) -> Vec<(uintptr_t, uintptr_t)> {
+    let mut cols = vec![];
+
+    for array in rb.columns().into_iter() {
+        let data = array.to_data();
+        let array_ptr = Arc::new(arrow::ffi::FFI_ArrowArray::new(&data));
+        let schema_ptr = Arc::new(
+            arrow::ffi::FFI_ArrowSchema::try_from(data.data_type()).expect("export schema c"),
+        );
+        cols.push((
+            Arc::into_raw(array_ptr) as uintptr_t,
+            Arc::into_raw(schema_ptr) as uintptr_t,
+        ));
+    }
+
+    cols
+}
+
 pub fn to_ptrs(rbs: Vec<RecordBatch>) -> (Vec<String>, Vec<Vec<(uintptr_t, uintptr_t)>>) {
     if rbs.is_empty() {
         return (vec![], vec![]);
@@ -132,21 +153,7 @@ pub fn to_ptrs(rbs: Vec<RecordBatch>) -> (Vec<String>, Vec<Vec<(uintptr_t, uintp
         .collect();
 
     for rb in rbs.into_iter() {
-        let mut cols = vec![];
-
-        for array in rb.columns().into_iter() {
-            let data = array.to_data();
-            let array_ptr = Arc::new(arrow::ffi::FFI_ArrowArray::new(&data));
-            let schema_ptr = Arc::new(
-                arrow::ffi::FFI_ArrowSchema::try_from(data.data_type()).expect("export schema c"),
-            );
-            cols.push((
-                Arc::into_raw(array_ptr) as uintptr_t,
-                Arc::into_raw(schema_ptr) as uintptr_t,
-            ));
-        }
-
-        result.push(cols);
+        result.push(to_ptrs_rb(rb));
     }
     (names, result)
 }