feat: enhance DataFrame and RecordBatchStream iteration to yield pyarrow.RecordBatch objects

Author: kosiew

python/datafusion/dataframe.py

@@ -1028,6 +1028,22 @@ def to_arrow_table(self) -> pa.Table:
         """
         return self.df.to_arrow_table()
 
+    def __iter__(self) -> Iterator[pa.RecordBatch]:
+        """Iterate over :py:class:`pyarrow.RecordBatch` objects.
+
+        This executes the DataFrame and yields each partition as a native
+        :py:class:`pyarrow.RecordBatch`.
+
+        Yields:
+            pyarrow.RecordBatch: the next batch in the result stream.
+        """
+        for batch in self.execute_stream():
+            # ``execute_stream`` yields batches that may be ``RecordBatch``
+            # wrappers or ``pyarrow.RecordBatch`` objects directly. Convert
+            # to native PyArrow batches when necessary to provide a consistent
+            # iterator interface.
+            yield batch.to_pyarrow() if hasattr(batch, "to_pyarrow") else batch
+
     def execute_stream(self) -> RecordBatchStream:
         """Executes this DataFrame and returns a stream over a single partition.
 

python/datafusion/record_batch.py

@@ -69,10 +69,10 @@ async def __anext__(self) -> RecordBatch:
         next_batch = await self.rbs.__anext__()
         return RecordBatch(next_batch)
 
-    def __next__(self) -> RecordBatch:
+    def __next__(self) -> pa.RecordBatch:
         """Iterator function."""
         next_batch = next(self.rbs)
-        return RecordBatch(next_batch)
+        return next_batch.to_pyarrow()
 
     def __aiter__(self) -> typing_extensions.Self:
         """Async iterator function."""

python/tests/test_dataframe.py

@@ -1505,7 +1505,8 @@ def test_to_arrow_table(df):
 
 def test_execute_stream(df):
     stream = df.execute_stream()
-    assert all(batch is not None for batch in stream)
+    batches = list(stream)
+    assert all(isinstance(batch, pa.RecordBatch) for batch in batches)
     assert not list(stream)  # after one iteration the generator must be exhausted
 
 
@@ -1526,30 +1527,9 @@ def test_execute_stream_to_arrow_table(df, schema):
     stream = df.execute_stream()
 
     if schema:
-        pyarrow_table = pa.Table.from_batches(
-            (batch.to_pyarrow() for batch in stream), schema=df.schema()
-        )
+        pyarrow_table = pa.Table.from_batches(stream, schema=df.schema())
     else:
-        pyarrow_table = pa.Table.from_batches(batch.to_pyarrow() for batch in stream)
-
-    assert isinstance(pyarrow_table, pa.Table)
-    assert pyarrow_table.shape == (3, 3)
-    assert set(pyarrow_table.column_names) == {"a", "b", "c"}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("schema", [True, False])
-async def test_execute_stream_to_arrow_table_async(df, schema):
-    stream = df.execute_stream()
-
-    if schema:
-        pyarrow_table = pa.Table.from_batches(
-            [batch.to_pyarrow() async for batch in stream], schema=df.schema()
-        )
-    else:
-        pyarrow_table = pa.Table.from_batches(
-            [batch.to_pyarrow() async for batch in stream]
-        )
+        pyarrow_table = pa.Table.from_batches(stream)
 
     assert isinstance(pyarrow_table, pa.Table)
     assert pyarrow_table.shape == (3, 3)
@@ -1558,7 +1538,9 @@ async def test_execute_stream_to_arrow_table_async(df, schema):
 
 def test_execute_stream_partitioned(df):
     streams = df.execute_stream_partitioned()
-    assert all(batch is not None for stream in streams for batch in stream)
+    assert all(
+        isinstance(batch, pa.RecordBatch) for stream in streams for batch in stream
+    )
     assert all(
         not list(stream) for stream in streams
     )  # after one iteration all generators must be exhausted