Refactor to use PyCapsule for ArrowArrayStream handling in tests and improve memory management

Author: kosiew

python/tests/test_dataframe.py

@@ -1606,7 +1606,7 @@ def fail_collect(self):  # pragma: no cover - failure path
 
 
 def test_arrow_c_stream_reader(df):
-    reader = pa.RecordBatchReader._import_from_c(df.__arrow_c_stream__())
+    reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
     assert isinstance(reader, pa.RecordBatchReader)
     table = pa.Table.from_batches(reader)
     expected = pa.Table.from_batches(df.collect())
@@ -2751,7 +2751,7 @@ def test_arrow_c_stream_interrupted():
         """
     )
 
-    reader = pa.RecordBatchReader._import_from_c(df.__arrow_c_stream__())
+    reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
 
     interrupted = False
     interrupt_error = None

python/tests/test_io.py

@@ -106,7 +106,7 @@ def test_arrow_c_stream_large_dataset(ctx):
     # Create a very large DataFrame using range; this would be terabytes if collected
     df = ctx.range(0, 1 << 40)
 
-    reader = pa.RecordBatchReader._import_from_c(df.__arrow_c_stream__())
+    reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
 
     # Track RSS before consuming batches
     psutil = pytest.importorskip("psutil")
@@ -126,7 +126,8 @@ def test_table_from_batches_stream(ctx, monkeypatch):
     df = ctx.range(0, 10)
 
     def fail_collect(self):  # pragma: no cover - failure path
-        raise AssertionError("collect should not be called")
+        msg = "collect should not be called"
+        raise AssertionError(msg)
 
     monkeypatch.setattr(DataFrame, "collect", fail_collect)
 

src/dataframe.rs

@@ -16,8 +16,8 @@
 // under the License.
 
 use std::collections::HashMap;
-use std::ffi::{c_void, CString};
-use std::sync::{Arc, OnceLock};
+use std::ffi::{c_void, CStr, CString};
+use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchReader};
 use arrow::compute::can_cast_types;
@@ -59,14 +59,15 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
-static ARROW_STREAM_NAME: OnceLock<CString> = OnceLock::new();
+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;
     }
-    let name = ARROW_STREAM_NAME.get_or_init(|| CString::new("arrow_array_stream").unwrap());
-    let stream_ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr()) as *mut FFI_ArrowArrayStream;
+    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));
     }
@@ -945,7 +946,7 @@ impl PyDataFrame {
         let stream = spawn_stream(py, async move { df.execute_stream().await })?;
 
         let mut schema: Schema = self.df.schema().to_owned().into();
-        let mut projection: Option<SchemaRef> = None;
+        let mut projection: Option<SchemaRef> = Some(Arc::new(schema.clone()));
 
         if let Some(schema_capsule) = requested_schema {
             validate_pycapsule(&schema_capsule, "arrow_schema")?;
@@ -957,7 +958,7 @@ impl PyDataFrame {
             projection = Some(Arc::new(schema.clone()));
         }
 
-        let schema_ref = projection.clone().unwrap_or_else(|| Arc::new(schema));
+        let schema_ref = Arc::new(schema.clone());
 
         let reader = DataFrameStreamReader {
             stream,
@@ -972,9 +973,12 @@ impl PyDataFrame {
             !stream_ptr.is_null(),
             "ArrowArrayStream pointer should never be null"
         );
-        let name = ARROW_STREAM_NAME.get_or_init(|| CString::new("arrow_array_stream").unwrap());
         let capsule = unsafe {
-            ffi::PyCapsule_New(stream_ptr as *mut c_void, name.as_ptr(), Some(drop_stream))
+            ffi::PyCapsule_New(
+                stream_ptr as *mut c_void,
+                ARROW_STREAM_NAME.as_ptr(),
+                Some(drop_stream),
+            )
         };
         if capsule.is_null() {
             unsafe { drop(Box::from_raw(stream_ptr)) };