Add memory limit test for record batch stream to ensure compliance with configured limits

Author: kosiew

python/tests/test_record_batch_stream.py

@@ -15,7 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import ctypes
+
+import pyarrow as pa
 import pytest
+from datafusion import RuntimeEnvBuilder, SessionContext
 
 
 def test_record_batch_stream_next(ctx):
@@ -65,3 +69,20 @@ def test_arrow_c_stream_partitioned(tmp_path, ctx):
         gc.collect()
 
     assert peak - baseline < rows_per_part * 8 * 2
+
+
+def test_record_batch_stream_memory_limit():
+    """Ensure streaming respects configured memory limits."""
+    runtime = RuntimeEnvBuilder().with_memory_limit(1024)
+    ctx = SessionContext(runtime=runtime)
+    df = ctx.sql("SELECT * FROM range(0, 1000000)")
+    capsule = df.__arrow_c_stream__()
+    ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
+    ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+    ptr = ctypes.pythonapi.PyCapsule_GetPointer(capsule, b"arrow_array_stream")
+    reader = pa.RecordBatchReader._import_from_c(ptr)
+
+    with pytest.raises(RuntimeError, match="(?i)memory"):
+        list(reader)
+
+    assert ctx.sql("SELECT 1").collect()[0].num_rows == 1