Revert "Enhance ArrowArrayStream handling by adding ARROW_STREAM_NAME and improving memory management in drop_stream function"

This reverts commit cd99b1d.

Author: kosiew

src/dataframe.rs

@@ -23,7 +23,7 @@ use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
 use arrow::compute::can_cast_types;
 use arrow::error::ArrowError;
 use arrow::ffi::FFI_ArrowSchema;
-use arrow::ffi_stream::{FFI_ArrowArrayStream, ARROW_STREAM_NAME};
+use arrow::ffi_stream::FFI_ArrowArrayStream;
 use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::{Schema, SchemaRef};
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
@@ -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};
@@ -59,17 +58,6 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
-unsafe extern "C" fn drop_stream(capsule: *mut ffi::PyObject) {
-    if capsule.is_null() {
-        return;
-    }
-    let name = CString::new(ARROW_STREAM_NAME).unwrap();
-    let stream_ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr()) as *mut FFI_ArrowArrayStream;
-    if !stream_ptr.is_null() {
-        drop(Box::from_raw(stream_ptr));
-    }
-}
-
 // 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
@@ -970,20 +958,20 @@ impl PyDataFrame {
             !stream_ptr.is_null(),
             "ArrowArrayStream pointer should never be null"
         );
-        let stream_name = CString::new(ARROW_STREAM_NAME).unwrap();
-        let capsule = unsafe {
-            ffi::PyCapsule_New(
-                stream_ptr as *mut c_void,
-                stream_name.as_ptr(),
-                Some(drop_stream),
+        let stream_capsule_name = CString::new("arrow_array_stream").unwrap();
+        unsafe {
+            PyCapsule::new_bound_with_destructor(
+                py,
+                stream_ptr,
+                Some(stream_capsule_name),
+                |ptr: *mut FFI_ArrowArrayStream, _| {
+                    if !ptr.is_null() {
+                        unsafe { Box::from_raw(ptr) };
+                    }
+                },
             )
-        };
-        if capsule.is_null() {
-            unsafe { drop(Box::from_raw(stream_ptr)) };
-            Err(PyErr::fetch(py).into())
-        } else {
-            Ok(unsafe { Bound::from_owned_ptr(py, capsule) })
         }
+        .map_err(PyDataFusionError::from)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {