test: add test for Arrow C stream capsule ownership

Author: kosiew

python/tests/test_dataframe.py

@@ -16,6 +16,7 @@
 # under the License.
 import ctypes
 import datetime
+import gc
 import os
 import re
 import threading
@@ -1620,6 +1621,22 @@ def test_arrow_c_stream_to_table_and_reader(fail_collect):
     assert reader_table.equals(expected)
 
 
+def test_arrow_c_stream_capsule_ownership(fail_collect):
+    ctx = SessionContext()
+
+    batch = pa.record_batch([pa.array([1])], names=["a"])
+    df = ctx.create_dataframe([[batch]])
+
+    capsule = df.__arrow_c_stream__()
+    reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    del capsule
+    gc.collect()
+
+    table = pa.Table.from_batches(reader)
+    expected = pa.Table.from_batches([batch])
+    assert table.equals(expected)
+
+
 def test_arrow_c_stream_order():
     ctx = SessionContext()
 

src/dataframe.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use std::collections::HashMap;
-use std::ffi::{c_void, CStr, CString};
+use std::ffi::{CStr, CString};
 use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
@@ -39,7 +39,6 @@ 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};
@@ -967,28 +966,12 @@ impl PyDataFrame {
         };
         let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
 
-        let stream = Box::new(FFI_ArrowArrayStream::new(reader));
-        let stream_ptr = Box::into_raw(stream);
-        debug_assert!(
-            !stream_ptr.is_null(),
-            "ArrowArrayStream pointer should never be null",
-        );
+        let stream = FFI_ArrowArrayStream::new(reader);
         // The returned capsule allows zero-copy hand-off to PyArrow. When
-        // 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(),
-                None,
-            )
-        };
-        if capsule.is_null() {
-            unsafe { drop(Box::from_raw(stream_ptr)) };
-            Err(PyErr::fetch(py).into())
-        } else {
-            let any = unsafe { Bound::from_owned_ptr(py, capsule) };
-            Ok(any.downcast_into::<PyCapsule>().unwrap())
-        }
+        // 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.
+        PyCapsule::new(py, stream, Some(ARROW_STREAM_NAME)).map_err(Into::into)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {