refactor: add to_record_batch_stream method and improve iteration support in DataFrame

Author: kosiew

python/datafusion/dataframe.py

@@ -25,6 +25,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncIterator,
     Iterable,
     Iterator,
     Literal,
@@ -1043,6 +1044,15 @@ def execute_stream_partitioned(self) -> list[RecordBatchStream]:
         streams = self.df.execute_stream_partitioned()
         return [RecordBatchStream(rbs) for rbs in streams]
 
+    def to_record_batch_stream(self) -> RecordBatchStream:
+        """Return a :py:class:`RecordBatchStream` over this DataFrame's results.
+
+        Returns:
+            A ``RecordBatchStream`` representing the lazily generated record
+            batches for this DataFrame.
+        """
+        return self.execute_stream()
+
     def to_pandas(self) -> pd.DataFrame:
         """Execute the :py:class:`DataFrame` and convert it into a Pandas DataFrame.
 
@@ -1126,12 +1136,12 @@ def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
         return self.df.__arrow_c_stream__(requested_schema)
 
     def __iter__(self) -> Iterator[RecordBatch]:
-        """Yield :class:`RecordBatch` objects by streaming execution."""
-        yield from self.to_record_batch_stream()
+        """Return an iterator over this DataFrame's record batches."""
+        return iter(self.to_record_batch_stream())
 
-    async def __aiter__(self) -> RecordBatchStream:
-        """Return an asynchronous iterator over streamed ``RecordBatch`` objects."""
-        return await self.to_record_batch_stream().__aiter__()
+    def __aiter__(self) -> AsyncIterator[RecordBatch]:
+        """Return an async iterator over this DataFrame's record batches."""
+        return self.to_record_batch_stream().__aiter__()
 
     def transform(self, func: Callable[..., DataFrame], *args: Any) -> DataFrame:
         """Apply a function to the current DataFrame which returns another DataFrame.

python/tests/test_dataframe.py

@@ -29,6 +29,7 @@
     DataFrame,
     ParquetColumnOptions,
     ParquetWriterOptions,
+    RecordBatch,
     SessionContext,
     WindowFrame,
     column,
@@ -390,10 +391,23 @@ def test_iter_batches(df):
     assert len(batches) == 1
 
     batch = batches[0]
-    assert isinstance(batch, pa.RecordBatch)
-    assert batch.column(0).to_pylist() == [1, 2, 3]
-    assert batch.column(1).to_pylist() == [4, 5, 6]
-    assert batch.column(2).to_pylist() == [8, 5, 8]
+    assert isinstance(batch, RecordBatch)
+    pa_batch = batch.to_pyarrow()
+    assert pa_batch.column(0).to_pylist() == [1, 2, 3]
+    assert pa_batch.column(1).to_pylist() == [4, 5, 6]
+    assert pa_batch.column(2).to_pylist() == [8, 5, 8]
+
+
+def test_to_record_batch_stream(df):
+    stream = df.to_record_batch_stream()
+    batches = list(stream)
+
+    assert len(batches) == 1
+    assert isinstance(batches[0], RecordBatch)
+    pa_batch = batches[0].to_pyarrow()
+    assert pa_batch.column(0).to_pylist() == [1, 2, 3]
+    assert pa_batch.column(1).to_pylist() == [4, 5, 6]
+    assert pa_batch.column(2).to_pylist() == [8, 5, 8]
 
 
 def test_with_column_renamed(df):
@@ -1331,7 +1345,7 @@ def test_execution_plan(aggregate_df):
 @pytest.mark.asyncio
 async def test_async_iteration_of_df(aggregate_df):
     rows_returned = 0
-    async for batch in aggregate_df.execute_stream():
+    async for batch in aggregate_df:
         assert batch is not None
         rows_returned += len(batch.to_pyarrow()[0])