Refactor ArrowArrayStream handling to use PyCapsule with destructor for improved memory management

Author: kosiew

python/datafusion/dataframe.py

@@ -1127,7 +1127,7 @@ def __iter__(self) -> Iterator[pa.RecordBatch]:
         """
         import pyarrow as pa
 
-        reader = pa.RecordBatchReader._import_from_c(self.__arrow_c_stream__())
+        reader = pa.RecordBatchReader._import_from_c_capsule(self.__arrow_c_stream__())
         yield from reader
 
     def transform(self, func: Callable[..., DataFrame], *args: Any) -> DataFrame:

src/dataframe.rs

@@ -17,7 +17,7 @@
 
 use std::collections::HashMap;
 use std::ffi::{c_void, CString};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
 use arrow::compute::can_cast_types;
@@ -39,6 +39,7 @@ use datafusion::prelude::*;
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use futures::{StreamExt, TryStreamExt};
 use pyo3::exceptions::PyValueError;
+use pyo3::ffi;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
@@ -58,6 +59,19 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
+static ARROW_STREAM_NAME: OnceLock<CString> = OnceLock::new();
+
+unsafe extern "C" fn drop_stream(capsule: *mut ffi::PyObject) {
+    if capsule.is_null() {
+        return;
+    }
+    let name = ARROW_STREAM_NAME.get_or_init(|| CString::new("arrow_array_stream").unwrap());
+    let stream_ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr()) as *mut FFI_ArrowArrayStream;
+    if !stream_ptr.is_null() {
+        drop(Box::from_raw(stream_ptr));
+    }
+}
+
 // https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
 // - we have not decided on the table_provider approach yet
 // this is an interim implementation
@@ -958,20 +972,17 @@ impl PyDataFrame {
             !stream_ptr.is_null(),
             "ArrowArrayStream pointer should never be null"
         );
-        let stream_capsule_name = CString::new("arrow_array_stream").unwrap();
-        unsafe {
-            PyCapsule::new_bound_with_destructor(
-                py,
-                stream_ptr,
-                Some(stream_capsule_name),
-                |ptr: *mut FFI_ArrowArrayStream, _| {
-                    if !ptr.is_null() {
-                        unsafe { Box::from_raw(ptr) };
-                    }
-                },
-            )
+        let name = ARROW_STREAM_NAME.get_or_init(|| CString::new("arrow_array_stream").unwrap());
+        let capsule = unsafe {
+            ffi::PyCapsule_New(stream_ptr as *mut c_void, name.as_ptr(), Some(drop_stream))
+        };
+        if capsule.is_null() {
+            unsafe { drop(Box::from_raw(stream_ptr)) };
+            Err(PyErr::fetch(py).into())
+        } else {
+            let any = unsafe { Bound::from_owned_ptr(py, capsule) };
+            Ok(any.downcast_into::<PyCapsule>().unwrap())
         }
-        .map_err(PyDataFusionError::from)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {