test: add test for Arrow C stream capsule release behavior

Author: kosiew

python/tests/test_dataframe.py

@@ -1698,6 +1698,37 @@ def test_arrow_c_stream_schema_mismatch(fail_collect):
         df.__arrow_c_stream__(bad_capsule)
 
 
+def test_arrow_c_stream_capsule_released(ctx):
+    df = ctx.from_pydict({"a": [1]})
+
+    capsule = df.__arrow_c_stream__()
+
+    get_destructor = ctypes.pythonapi.PyCapsule_GetDestructor
+    get_destructor.restype = ctypes.c_void_p
+    get_destructor.argtypes = [ctypes.py_object]
+
+    # The capsule should not have a registered destructor
+    assert get_destructor(capsule) == 0
+
+    get_ptr = ctypes.pythonapi.PyCapsule_GetPointer
+    get_ptr.restype = ctypes.c_void_p
+    get_ptr.argtypes = [ctypes.py_object, ctypes.c_char_p]
+    pyerr_clear = ctypes.pythonapi.PyErr_Clear
+    pyerr_clear.restype = None
+    pyerr_clear.argtypes = []
+
+    # Pointer is valid before importing to PyArrow
+    assert get_ptr(capsule, b"arrow_array_stream")
+    pyerr_clear()
+
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    reader.read_all()
+
+    # After import the capsule no longer exposes the pointer
+    with pytest.raises(ValueError):
+        get_ptr(capsule, b"arrow_array_stream")
+    pyerr_clear()
+
 def test_to_pylist(df):
     # Convert datafusion dataframe to Python list
     pylist = df.to_pylist()

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
@@ -999,14 +974,12 @@ impl PyDataFrame {
             "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.
+        // PyArrow imports the capsule it assumes ownership of the stream.
         let capsule = unsafe {
             ffi::PyCapsule_New(
                 stream_ptr as *mut c_void,
                 ARROW_STREAM_NAME.as_ptr(),
-                Some(drop_stream),
+                None,
             )
         };
         if capsule.is_null() {