test: add garbage collection tests for arrow C stream reader

Author: kosiew

python/tests/test_dataframe.py

@@ -1874,6 +1874,55 @@ def test_arrow_streams_context_drop_no_segfault():
         gc.collect()
 
 
+def test_arrow_c_stream_reader_discard_gc():
+    """Capsule imported into a reader and immediately discarded."""
+
+    for _ in range(50):
+        ctx = SessionContext()
+        df = ctx.sql("SELECT 1 AS a")
+        reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
+        del reader
+        del df
+        del ctx
+        gc.collect()
+
+
+def test_arrow_c_stream_unconsumed_gc():
+    """Capsule left unconsumed and only garbage-collected."""
+
+    for _ in range(50):
+        ctx = SessionContext()
+        df = ctx.sql("SELECT 1 AS a")
+        df.__arrow_c_stream__()
+        del df
+        del ctx
+        gc.collect()
+
+
+def test_arrow_c_stream_concurrent_context_drop():
+    """Consume capsule while context is deleted concurrently."""
+
+    for _ in range(50):
+        ctx = SessionContext()
+        df = ctx.sql("SELECT 1 AS a")
+        capsule = df.__arrow_c_stream__()
+        reader = pa.RecordBatchReader._import_from_c_capsule(capsule)
+        result: list[pa.Table] = []
+
+        def consume(reader=reader, result=result) -> None:
+            result.append(reader.read_all())
+
+        thread = threading.Thread(target=consume)
+        thread.start()
+        del df
+        del ctx
+        thread.join()
+        assert result[0].column("a").to_pylist() == [1]
+        del reader
+        del capsule
+        gc.collect()
+
+
 def test_arrow_stream_to_pylist(df):
     capsule = df.__arrow_c_stream__()
     reader = pa.RecordBatchReader._import_from_c_capsule(capsule)