Revert "UNPICK"

This reverts commit 29d2faf.

Author: kosiew

docs/source/conf.py

@@ -72,6 +72,14 @@
 suppress_warnings = ["autoapi.python_import_resolution"]
 autoapi_python_class_content = "both"
 autoapi_keep_files = False  # set to True for debugging generated files
+autoapi_options = [
+    "members",
+    "undoc-members",
+    "special-members",
+    "show-inheritance",
+    "show-module-summary",
+    "imported-members",
+]
 
 
 def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool:  # noqa: ARG001

docs/source/user-guide/dataframe/index.rst

@@ -145,10 +145,41 @@ To materialize the results of your DataFrame operations:
     
     # Display results
     df.show()                         # Print tabular format to console
-    
+
     # Count rows
     count = df.count()
 
+PyArrow Streaming
+-----------------
+
+DataFusion DataFrames implement the ``__arrow_c_stream__`` protocol, enabling
+zero-copy streaming into libraries like `PyArrow <https://arrow.apache.org/>`_.
+Earlier versions eagerly converted the entire DataFrame when exporting to
+PyArrow, which could exhaust memory on large datasets. With streaming, batches
+are produced lazily so you can process arbitrarily large results without
+out-of-memory errors.
+
+.. code-block:: python
+
+    import pyarrow as pa
+
+    # Create a PyArrow RecordBatchReader without materializing all batches
+    reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
+    for batch in reader:
+        ...  # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch` objects
+that implement the Arrow C data interface. These batches can be consumed by
+libraries like PyArrow without copying:
+
+.. code-block:: python
+
+    for batch in df:
+        pa_batch = batch.to_pyarrow()  # optional conversion
+        ...  # process each batch as it is produced
+
+See :doc:`../io/arrow` for additional details on the Arrow interface.
+
 HTML Rendering
 --------------
 

python/datafusion/__init__.py

@@ -53,7 +53,7 @@
 )
 from .io import read_avro, read_csv, read_json, read_parquet
 from .plan import ExecutionPlan, LogicalPlan
-from .record_batch import RecordBatch, RecordBatchStream
+from .record_batch import RecordBatch, RecordBatchStream, to_record_batch_stream
 from .user_defined import (
     Accumulator,
     AggregateUDF,
@@ -107,6 +107,7 @@
     "read_json",
     "read_parquet",
     "substrait",
+    "to_record_batch_stream",
     "udaf",
     "udf",
     "udtf",

python/datafusion/dataframe.py

@@ -25,7 +25,9 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncIterator,
     Iterable,
+    Iterator,
     Literal,
     Optional,
     Union,
@@ -42,7 +44,11 @@
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
 from datafusion.expr import Expr, SortExpr, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
-from datafusion.record_batch import RecordBatchStream
+from datafusion.record_batch import (
+    RecordBatch,
+    RecordBatchStream,
+    to_record_batch_stream,
+)
 
 if TYPE_CHECKING:
     import pathlib
@@ -53,6 +59,7 @@
     import pyarrow as pa
 
     from datafusion._internal import expr as expr_internal
+    from datafusion.record_batch import RecordBatch
 
 from enum import Enum
 
@@ -289,6 +296,9 @@ def __init__(
 class DataFrame:
     """Two dimensional table representation of data.
 
+    DataFrame objects are iterable; iterating over a DataFrame yields
+    :class:`pyarrow.RecordBatch` instances lazily.
+
     See :ref:`user_guide_concepts` in the online documentation for more information.
     """
 
@@ -1018,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.
 
@@ -1098,21 +1124,41 @@ def unnest_columns(self, *columns: str, preserve_nulls: bool = True) -> DataFram
         return DataFrame(self.df.unnest_columns(columns, preserve_nulls=preserve_nulls))
 
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
-        """Export an Arrow PyCapsule Stream.
+        """Export the DataFrame as an Arrow C Stream.
 
-        This will execute and collect the DataFrame. We will attempt to respect the
-        requested schema, but only trivial transformations will be applied such as only
-        returning the fields listed in the requested schema if their data types match
-        those in the DataFrame.
+        The DataFrame is executed using DataFusion's streaming APIs and exposed via
+        Arrow's C Stream interface. Record batches are produced incrementally, so the
+        full result set is never materialized in memory. When ``requested_schema`` is
+        provided, only straightforward projections such as column selection or
+        reordering are applied.
 
         Args:
             requested_schema: Attempt to provide the DataFrame using this schema.
 
         Returns:
-            Arrow PyCapsule object.
+            Arrow PyCapsule object representing an ``ArrowArrayStream``.
         """
+        # ``DataFrame.__arrow_c_stream__`` in the Rust extension leverages
+        # ``execute_stream_partitioned`` under the hood to stream batches while
+        # preserving the original partition order.
         return self.df.__arrow_c_stream__(requested_schema)
 
+    def __iter__(self) -> Iterator[RecordBatch]:
+        """Yield record batches from the DataFrame without materializing results.
+
+        This implementation delegates to :func:`to_record_batch_stream`, which
+        executes the DataFrame and returns a :class:`RecordBatchStream`.
+        """
+        return to_record_batch_stream(self).__iter__()
+
+    def __aiter__(self) -> AsyncIterator[RecordBatch]:
+        """Asynchronously yield record batches from the DataFrame.
+
+        This delegates to :func:`to_record_batch_stream` to obtain a
+        :class:`RecordBatchStream` and returns its asynchronous iterator.
+        """
+        return to_record_batch_stream(self).__aiter__()
+
     def transform(self, func: Callable[..., DataFrame], *args: Any) -> DataFrame:
         """Apply a function to the current DataFrame which returns another DataFrame.
 

python/datafusion/record_batch.py

@@ -25,11 +25,13 @@
 
 from typing import TYPE_CHECKING
 
+import datafusion._internal as df_internal
+
 if TYPE_CHECKING:
     import pyarrow as pa
     import typing_extensions
 
-    import datafusion._internal as df_internal
+    from datafusion.dataframe import DataFrame
 
 
 class RecordBatch:
@@ -67,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."""
@@ -79,3 +81,15 @@ def __aiter__(self) -> typing_extensions.Self:
     def __iter__(self) -> typing_extensions.Self:
         """Iterator function."""
         return self
+
+
+def to_record_batch_stream(df: DataFrame) -> RecordBatchStream:
+    """Convert a DataFrame into a RecordBatchStream.
+
+    Args:
+        df: DataFrame to convert.
+
+    Returns:
+        A RecordBatchStream representing the DataFrame.
+    """
+    return df.execute_stream()

python/tests/conftest.py

@@ -17,7 +17,7 @@
 
 import pyarrow as pa
 import pytest
-from datafusion import SessionContext
+from datafusion import DataFrame, SessionContext
 from pyarrow.csv import write_csv
 
 
@@ -49,3 +49,12 @@ def database(ctx, tmp_path):
         delimiter=",",
         schema_infer_max_records=10,
     )
+
+
+@pytest.fixture
+def fail_collect(monkeypatch):
+    def _fail_collect(self, *args, **kwargs):  # pragma: no cover - failure path
+        msg = "collect should not be called"
+        raise AssertionError(msg)
+
+    monkeypatch.setattr(DataFrame, "collect", _fail_collect)