test: add context drop test for SessionContext to prevent segfaults

Author: kosiew

python/tests/test_dataframe.py

@@ -1832,6 +1832,22 @@ def test_arrow_c_stream_capsule_manual_destructor_noop(ctx):
     gc.collect()
 
 
+def test_arrow_c_stream_context_drop_no_segfault():
+    """Repeatedly create/drop SessionContext after __arrow_c_stream__."""
+    for _ in range(5):
+        ctx = SessionContext()
+        df = ctx.sql("SELECT 1 AS a")
+        capsule = df.__arrow_c_stream__()
+        del df
+        del ctx
+        reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+        del capsule
+        table = reader.read_all()
+        assert table.num_rows == 1
+        del reader
+        gc.collect()
+
+
 def test_arrow_stream_to_pylist(df):
     capsule = df.__arrow_c_stream__()
     reader = pa.RecordBatchReader._import_from_c_capsule(capsule)

src/dataframe.rs

@@ -18,6 +18,7 @@
 use std::collections::HashMap;
 use std::ffi::CString;
 use std::sync::Arc;
+use tokio::sync::Mutex;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchIterator, RecordBatchReader};
 use arrow::compute::can_cast_types;
@@ -33,8 +34,8 @@ use datafusion::config::{CsvOptions, ParquetColumnOptions, ParquetOptions, Table
 use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
 use datafusion::datasource::TableProvider;
 use datafusion::error::DataFusionError;
-use datafusion::execution::SendableRecordBatchStream;
 use datafusion::execution::session_state::SessionState;
+use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::prelude::*;
 use datafusion_ffi::table_provider::FFI_TableProvider;
@@ -360,7 +361,7 @@ impl PyDataFrame {
 /// their original partition order. When a `projection` is set, each batch is
 /// converted via `record_batch_into_schema` to apply schema changes per batch.
 struct PartitionedDataFrameStreamReader {
-    streams: Vec<SendableRecordBatchStream>,
+    streams: Vec<Arc<Mutex<SendableRecordBatchStream>>>,
     // Hold a reference to the session state to keep the context alive
     _state: Arc<SessionState>,
     schema: SchemaRef,
@@ -373,12 +374,17 @@ impl Iterator for PartitionedDataFrameStreamReader {
 
     fn next(&mut self) -> Option<Self::Item> {
         while self.current < self.streams.len() {
-            let stream = &mut self.streams[self.current];
-            let fut = poll_next_batch(stream);
-            let result = Python::with_gil(|py| wait_for_future(py, fut));
+            let stream = self.streams[self.current].clone();
+
+            let result = Python::with_gil(|py| {
+                spawn_future(py, async move {
+                    let mut s = stream.lock().await;
+                    poll_next_batch(&mut s).await
+                })
+            });
 
             match result {
-                Ok(Ok(Some(batch))) => {
+                Ok(Some(batch)) => {
                     let batch = if let Some(ref schema) = self.projection {
                         match record_batch_into_schema(batch, schema.as_ref()) {
                             Ok(b) => b,
@@ -389,13 +395,10 @@ impl Iterator for PartitionedDataFrameStreamReader {
                     };
                     return Some(Ok(batch));
                 }
-                Ok(Ok(None)) => {
+                Ok(None) => {
                     self.current += 1;
                     continue;
                 }
-                Ok(Err(e)) => {
-                    return Some(Err(ArrowError::ExternalError(Box::new(e))));
-                }
                 Err(e) => {
                     return Some(Err(ArrowError::ExternalError(Box::new(e))));
                 }
@@ -956,6 +959,10 @@ impl PyDataFrame {
         let df = self.df.as_ref().clone();
         let state = df.session_state().clone();
         let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
+        let streams = streams
+            .into_iter()
+            .map(|s| Arc::new(Mutex::new(s)))
+            .collect();
 
         let mut schema: Schema = self.df.schema().to_owned().into();
         let mut projection: Option<SchemaRef> = None;
@@ -999,12 +1006,10 @@ impl PyDataFrame {
         let df = self.df.as_ref().clone();
         let state = df.session_state().clone();
         let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
-        Ok(
-            streams
-                .into_iter()
-                .map(|stream| PyRecordBatchStream::new(stream, state.clone()))
-                .collect(),
-        )
+        Ok(streams
+            .into_iter()
+            .map(|stream| PyRecordBatchStream::new(stream, state.clone()))
+            .collect())
     }
 
     /// Convert to pandas dataframe with pyarrow