feat: add test for arrow C stream schema selection in DataFrame

Author: kosiew

python/tests/test_dataframe.py

@@ -1637,6 +1637,37 @@ def test_arrow_c_stream_reader(df):
     assert table.equals(expected)
 
 
+def test_arrow_c_stream_schema_selection(fail_collect):
+    ctx = SessionContext()
+
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array([1, 2]),
+            pa.array([3, 4]),
+            pa.array([5, 6]),
+        ],
+        names=["a", "b", "c"],
+    )
+    df = ctx.create_dataframe([[batch]])
+
+    requested_schema = pa.schema([("c", pa.int64()), ("a", pa.int64())])
+    schema_capsule = requested_schema._export_to_c()
+
+    reader = pa.RecordBatchReader._import_from_c_capsule(
+        df.__arrow_c_stream__(schema_capsule)
+    )
+
+    assert reader.schema == requested_schema
+
+    batches = list(reader)
+
+    assert len(batches) == 1
+    expected_batch = pa.record_batch(
+        [pa.array([5, 6]), pa.array([1, 2])], names=["c", "a"]
+    )
+    assert batches[0].equals(expected_batch)
+
+
 def test_to_pylist(df):
     # Convert datafusion dataframe to Python list
     pylist = df.to_pylist()