refactor: simplify capsule creation and remove custom destructor for FFI_ArrowArrayStream

Author: kosiew

src/dataframe.rs

@@ -16,8 +16,7 @@
 // under the License.
 
 use std::collections::HashMap;
-use std::ffi::{CStr, CString};
-use std::os::raw::c_void;
+use std::ffi::CString;
 use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
@@ -40,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};
@@ -60,29 +58,6 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
-#[allow(clippy::manual_c_str_literals)]
-static ARROW_STREAM_NAME: &CStr =
-    unsafe { CStr::from_bytes_with_nul_unchecked(b"arrow_array_stream\0") };
-
-/// Capsule destructor for [`FFI_ArrowArrayStream`].
-///
-/// PyArrow currently takes ownership of the stream and nulls out the capsule's
-/// pointer when importing. This destructor defensively checks for a null pointer
-/// to guard against double free until PyArrow is removed.
-unsafe extern "C" fn drop_arrow_array_stream_capsule(capsule: *mut ffi::PyObject) {
-    // Safety: `capsule` is expected to be a valid `PyCapsule` containing a
-    // `FFI_ArrowArrayStream` created by this module.
-    let ptr = unsafe { ffi::PyCapsule_GetPointer(capsule, ARROW_STREAM_NAME.as_ptr()) };
-    if ptr.is_null() {
-        return;
-    }
-    // Reconstruct and drop the stream then null out the capsule pointer.
-    unsafe {
-        drop(Box::from_raw(ptr.cast::<FFI_ArrowArrayStream>()));
-        ffi::PyCapsule_SetPointer(capsule, std::ptr::null_mut());
-    }
-}
-
 // 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
@@ -988,25 +963,9 @@ impl PyDataFrame {
         let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
 
         let stream = FFI_ArrowArrayStream::new(reader);
-        // This custom destructor is defensive to avoid a double free. It is
-        // a temporary workaround until DataFusion's Python bindings decouple
-        // from PyArrow.
-        // The returned capsule allows zero-copy hand-off to PyArrow. When
-        // PyArrow imports the capsule it assumes ownership of the stream and
-        // nulls out the capsule's internal pointer so the destructor does not
-        // free it twice.
-        unsafe {
-            let capsule = ffi::PyCapsule_New(
-                Box::into_raw(Box::new(stream)) as *mut c_void,
-                ARROW_STREAM_NAME.as_ptr(),
-                Some(drop_arrow_array_stream_capsule),
-            );
-            if capsule.is_null() {
-                Err(PyErr::fetch(py).into())
-            } else {
-                Ok(Bound::from_owned_ptr(py, capsule).downcast_into_unchecked())
-            }
-        }
+        let name = CString::new("arrow_array_stream").unwrap();
+        let capsule = PyCapsule::new(py, stream, Some(name)).map_err(py_datafusion_err)?;
+        Ok(capsule)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {