test: add cleanup test for iterating and streaming with arrow C

Author: kosiew

python/tests/test_dataframe.py

@@ -1866,6 +1866,27 @@ def test_arrow_stream_to_pandas(df):
     assert pdf.shape == (3, 3)
 
 
+def test_iter_then_arrow_c_stream_gc():
+    """Iterate via __iter__ and __arrow_c_stream__ repeatedly to ensure cleanup."""
+
+    for _ in range(3):
+        ctx = SessionContext()
+        df = ctx.from_pydict({"a": [1, 2]})
+
+        batches = list(df)
+        assert len(batches) == 1
+        assert batches[0].column(0).to_pylist() == [1, 2]
+
+        reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
+        table = pa.Table.from_batches(reader)
+        assert table.column("a").to_pylist() == [1, 2]
+
+        del reader
+        del df
+        del ctx
+        gc.collect()
+
+
 def test_to_pylist(df):
     # Convert datafusion dataframe to Python list
     pylist = df.to_pylist()