Improve drop_stream function to handle PyArrow ownership transfer correctly

Author: kosiew

src/dataframe.rs

@@ -67,13 +67,25 @@ unsafe extern "C" fn drop_stream(capsule: *mut ffi::PyObject) {
     if capsule.is_null() {
         return;
     }
-    let stream_ptr =
-        ffi::PyCapsule_GetPointer(capsule, ARROW_STREAM_NAME.as_ptr()) as *mut FFI_ArrowArrayStream;
-    if stream_ptr.is_null() {
-        ffi::PyErr_Clear();
-        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));
+        }
     }
-    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
@@ -976,6 +988,10 @@ 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.
         let capsule = unsafe {
             ffi::PyCapsule_New(
                 stream_ptr as *mut c_void,