test: add garbage collection test for arrow C stream capsule

Author: kosiew

python/tests/test_dataframe.py

@@ -1833,6 +1833,20 @@ def test_arrow_c_stream_capsule_manual_destructor_noop(ctx):
     gc.collect()
 
 
+def test_arrow_c_stream_capsule_gc_without_consuming(ctx):
+    df = ctx.from_pydict({"a": [1]})
+
+    capsule = df.__arrow_c_stream__()
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+
+    del capsule
+    gc.collect()
+
+    table = reader.read_all()
+    expected = pa.table({"a": [1]})
+    assert table.equals(expected)
+
+
 def test_arrow_c_stream_context_drop_no_segfault():
     """Repeatedly create/drop SessionContext after __arrow_c_stream__."""
     for _ in range(5):

src/dataframe.rs

@@ -40,9 +40,11 @@ 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};
+use std::os::raw::c_void;
 
 use crate::catalog::PyTable;
 use crate::errors::{py_datafusion_err, PyDataFusionError};
@@ -376,7 +378,7 @@ struct PartitionedDataFrameStreamReader {
 /// interface.
 struct StreamWithContext {
     reader: PartitionedDataFrameStreamReader,
-    ctx: Arc<SessionContext>,
+    _ctx: Arc<SessionContext>,
 }
 
 impl Iterator for PartitionedDataFrameStreamReader {
@@ -452,6 +454,18 @@ impl Drop for StreamWithContext {
     }
 }
 
+unsafe extern "C" fn stream_capsule_destructor(capsule: *mut ffi::PyObject) {
+    let name = pyo3::ffi::c_str!("arrow_array_stream");
+    unsafe {
+        let ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr());
+        if ptr.is_null() {
+            ffi::PyCapsule_SetDestructor(capsule, None);
+            return;
+        }
+        drop(Box::from_raw(ptr.cast::<FFI_ArrowArrayStream>()));
+    }
+}
+
 #[pymethods]
 impl PyDataFrame {
     /// Enable selection for `df[col]`, `df[col1, col2, col3]`, and `df[[col1, col2, col3]]`
@@ -1022,14 +1036,25 @@ impl PyDataFrame {
             projection,
             current: 0,
         };
-        let reader = StreamWithContext { reader, ctx };
+        let reader = StreamWithContext { reader, _ctx: ctx };
         let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
 
         let stream = FFI_ArrowArrayStream::new(reader);
         let name = pyo3::ffi::c_str!("arrow_array_stream");
 
-        let capsule =
-            PyCapsule::new(py, stream, Some(name.to_owned())).map_err(py_datafusion_err)?;
+        let capsule_ptr = unsafe {
+            ffi::PyCapsule_New(
+                Box::into_raw(Box::new(stream)) as *mut c_void,
+                name.as_ptr(),
+                Some(stream_capsule_destructor),
+            )
+        };
+        if capsule_ptr.is_null() {
+            return Err(PyErr::fetch(py).into());
+        }
+        let capsule = unsafe {
+            Bound::from_owned_ptr(py, capsule_ptr).downcast_into_unchecked::<PyCapsule>()
+        };
         Ok(capsule)
     }