feat: implement capsule destructor for FFI_ArrowArrayStream to prevent double free

Author: kosiew

src/dataframe.rs

@@ -17,6 +17,7 @@
 
 use std::collections::HashMap;
 use std::ffi::{CStr, CString};
+use std::os::raw::c_void;
 use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
@@ -39,6 +40,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};
@@ -62,6 +64,25 @@ use crate::{
 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
@@ -971,9 +992,18 @@ impl PyDataFrame {
         // 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.
-        // ARROW_STREAM_NAME is a `&CStr`; PyCapsule::new expects an
-        // `Option<CString>` for the name, so convert it here.
-        PyCapsule::new(py, stream, Some(ARROW_STREAM_NAME.into())).map_err(Into::into)
+        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())
+            }
+        }
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {