refactor: simplify ArrowArrayStream handling in PyDataFrame implementation

Author: kosiew

src/dataframe.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use std::collections::HashMap;
-use std::ffi::{c_void, CStr, CString};
+use std::ffi::{CStr, CString};
 use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
@@ -39,7 +39,6 @@ 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};
@@ -967,25 +966,10 @@ impl PyDataFrame {
         };
         let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
 
-        let stream = Box::new(FFI_ArrowArrayStream::new(reader));
-        let stream_ptr = Box::into_raw(stream);
-        debug_assert!(
-            !stream_ptr.is_null(),
-            "ArrowArrayStream pointer should never be null",
-        );
-        // The returned capsule allows zero-copy hand-off to PyArrow. Once
-        // PyArrow imports the capsule it assumes ownership of the stream, so we
-        // rely on PyArrow or PyO3 to manage the stream's lifetime.
-        let capsule = unsafe {
-            ffi::PyCapsule_New(stream_ptr as *mut c_void, ARROW_STREAM_NAME.as_ptr(), None)
-        };
-        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())
-        }
+        let stream = FFI_ArrowArrayStream::new(reader);
+        // The returned capsule allows zero-copy hand-off to PyArrow.
+        let capsule = PyCapsule::new(py, stream, Some(ARROW_STREAM_NAME.into()))?;
+        Ok(capsule)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {