test: add test for arrow stream context drop to prevent segfaults

Author: kosiew

python/tests/test_dataframe.py

@@ -34,6 +34,7 @@
     WindowFrame,
     column,
     literal,
+    to_record_batch_stream,
 )
 from datafusion import (
     functions as f,
@@ -1848,6 +1849,31 @@ def test_arrow_c_stream_context_drop_no_segfault():
         gc.collect()
 
 
+def test_arrow_streams_context_drop_no_segfault():
+    """__arrow_c_stream__ and to_record_batch_stream survive context drop."""
+
+    for _ in range(5):
+        ctx = SessionContext()
+        df = ctx.sql("SELECT 1 AS a")
+        capsule = df.__arrow_c_stream__()
+        stream = to_record_batch_stream(df)
+        del df
+        del ctx
+
+        reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+        del capsule
+        table = reader.read_all()
+        assert table.column("a").to_pylist() == [1]
+        del reader
+
+        batches = list(stream)
+        assert len(batches) == 1
+        assert batches[0].column("a").to_pylist() == [1]
+        del batches
+        del stream
+        gc.collect()
+
+
 def test_arrow_stream_to_pylist(df):
     capsule = df.__arrow_c_stream__()
     reader = pa.RecordBatchReader._import_from_c_capsule(capsule)