test: add tests for converting Arrow stream to Python list, dict, and pandas DataFrame

Author: kosiew

python/tests/test_dataframe.py

@@ -1795,6 +1795,40 @@ def test_arrow_c_stream_capsule_manual_destructor_noop(ctx):
     gc.collect()
 
 
+def test_arrow_stream_to_pylist(df):
+    capsule = df.__arrow_c_stream__()
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    reader.read_all()
+
+    pylist = df.to_pylist()
+    assert pylist == [
+        {"a": 1, "b": 4, "c": 8},
+        {"a": 2, "b": 5, "c": 5},
+        {"a": 3, "b": 6, "c": 8},
+    ]
+
+
+def test_arrow_stream_to_pydict(df):
+    capsule = df.__arrow_c_stream__()
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    reader.read_all()
+
+    pydict = df.to_pydict()
+    assert pydict == {"a": [1, 2, 3], "b": [4, 5, 6], "c": [8, 5, 8]}
+
+
+def test_arrow_stream_to_pandas(df):
+    pd = pytest.importorskip("pandas")
+
+    capsule = df.__arrow_c_stream__()
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    reader.read_all()
+
+    pdf = df.to_pandas()
+    assert isinstance(pdf, pd.DataFrame)
+    assert pdf.shape == (3, 3)
+
+
 def test_to_pylist(df):
     # Convert datafusion dataframe to Python list
     pylist = df.to_pylist()

src/dataframe.rs

@@ -17,6 +17,7 @@
 
 use std::collections::HashMap;
 use std::ffi::CString;
+use std::os::raw::c_void;
 use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
@@ -963,8 +964,23 @@ impl PyDataFrame {
         let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
 
         let stream = FFI_ArrowArrayStream::new(reader);
+        fn drop_stream(ptr: *mut FFI_ArrowArrayStream, _ctx: *mut c_void) {
+            if ptr.is_null() {
+                return;
+            }
+            unsafe {
+                drop(Box::from_raw(ptr));
+            }
+        }
+        let stream_ptr = Box::into_raw(Box::new(stream));
         let name = CString::new("arrow_array_stream").unwrap();
-        let capsule = PyCapsule::new(py, stream, Some(name)).map_err(py_datafusion_err)?;
+        let capsule = PyCapsule::new_bound_with_destructor(
+            py,
+            stream_ptr,
+            Some(name),
+            drop_stream as fn(*mut FFI_ArrowArrayStream, *mut c_void),
+        )
+        .map_err(py_datafusion_err)?;
         Ok(capsule)
     }