refactor: remove drop_stream function and update capsule management for ArrowArrayStream

Author: kosiew

src/dataframe.rs

@@ -63,31 +63,6 @@ use crate::{
 static ARROW_STREAM_NAME: &CStr =
     unsafe { CStr::from_bytes_with_nul_unchecked(b"arrow_array_stream\0") };
 
-unsafe extern "C" fn drop_stream(capsule: *mut ffi::PyObject) {
-    if capsule.is_null() {
-        return;
-    }
-
-    // When PyArrow imports this capsule it steals the raw stream pointer and
-    // sets the capsule's internal pointer to NULL. In that case
-    // `PyCapsule_IsValid` returns 0 and this destructor must not drop the
-    // stream as ownership has been transferred to PyArrow. If the capsule was
-    // never imported, the pointer remains valid and we are responsible for
-    // freeing the stream here.
-    if ffi::PyCapsule_IsValid(capsule, ARROW_STREAM_NAME.as_ptr()) == 1 {
-        let stream_ptr = ffi::PyCapsule_GetPointer(capsule, ARROW_STREAM_NAME.as_ptr())
-            as *mut FFI_ArrowArrayStream;
-        if !stream_ptr.is_null() {
-            drop(Box::from_raw(stream_ptr));
-        }
-    }
-
-    // `PyCapsule_GetPointer` sets a Python error on failure. Clear it only
-    // after the stream has been released (or determined to be owned
-    // elsewhere).
-    ffi::PyErr_Clear();
-}
-
 // 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
@@ -998,16 +973,11 @@ impl PyDataFrame {
             !stream_ptr.is_null(),
             "ArrowArrayStream pointer should never be null",
         );
-        // 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 `drop_stream` knows not to
-        // free it.
+        // 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(),
-                Some(drop_stream),
-            )
+            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)) };