You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "kosiew (via GitHub)" <gi...@apache.org> on 2025/09/02 11:16:03 UTC
[PR] DRAFT Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
kosiew opened a new pull request, #1222:
URL: https://github.com/apache/datafusion-python/pull/1222
## Which issue does this PR close?
* Closes #1206
> does ArrowStreamExportable to load the full data into memory or it is a recordbatch reader as I am getting OOM when used in smaller VM
## Rationale for this change
Exporting a DataFrame to PyArrow previously required materializing (collecting) the entire result set. For large datasets this caused high memory usage and could lead to Out Of Memory (OOM) failures. The goal of this PR is to provide a zero‑copy, incremental export path into PyArrow and to make `DataFrame` objects iterable so consumers can process record batches lazily. In addition, the PR improves Python/Rust integration for running async DataFusion operations while respecting Python signal handling (e.g. `KeyboardInterrupt`).
These changes make streaming the default, memory‑efficient path for Python users exporting to Arrow and provide robust behavior when interrupted.
## What changes are included in this PR?
**High level**
* Add first-class Arrow C stream support for `DataFrame.__arrow_c_stream__`, returning an ArrowArrayStream PyCapsule that PyArrow can import without copying or full materialization.
* Make `DataFrame` iterable in Python so `for batch in df:` yields `pyarrow.RecordBatch` lazily.
* Implement a partition-aware synchronous `RecordBatchReader` wrapper in Rust that drains partitioned `SendableRecordBatchStream`s while preserving original partition order and applying an optional projection per-batch.
* Add robust PyCapsule lifecycle handling that transfers ownership to PyArrow when it imports the capsule (avoid double-free) and clears errors set by `PyCapsule_GetPointer` safely.
* Introduce `spawn_future` utility to run futures on the global Tokio runtime and wait on them while correctly translating join errors and preserving Python signal handling.
* Replace previous `collect`/join-handle patterns in a few Rust methods with `spawn_future` to unify behavior.
**Files changed (non-exhaustive, representative):**
* Documentation
* `docs/source/user-guide/dataframe/index.rst` — add `PyArrow Streaming` section that documents `__arrow_c_stream__` and iterator behavior.
* Python package
* `python/datafusion/dataframe.py` — update docstrings for `__arrow_c_stream__`, add `__iter__` impl that imports the Arrow capsule and yields batches lazily.
* `python/tests/*` — many tests added/updated to validate streaming behavior, ordering, schema projection, schema mismatch errors, interruption, and memory usage.
* `python/tests/utils.py` — test helper `range_table` for generating very large ranges in tests.
* `python/tests/test_dataframe_iter_stream.py` — new tests to validate reader close/release semantics.
* `python/tests/test_io.py` — add `test_arrow_c_stream_large_dataset` which verifies memory does not grow substantially while reading a handful of batches.
* `python/tests/conftest.py` — add `fail_collect` fixture that asserts `collect` is not called in streaming paths.
* Rust extension
* `src/dataframe.rs` — implement `PartitionedDataFrameStreamReader`, return `ArrowArrayStream` PyCapsule with destructor `drop_stream`, preserve partition order, support optional projection per batch.
* `src/record_batch.rs` — add `poll_next_batch` helper that converts the streaming `Option<Result<_>>` into `Result<Option<_>>` for ergonomic handling.
* `src/utils.rs` — add `spawn_future` that runs a given future on the Tokio runtime and waits while mapping JoinErrors and preserving Python signal semantics.
* `src/context.rs` — use `spawn_future` for streaming/execution entry points.
**Behavioral notes:**
* The exported capsule is zero‑copy and PyArrow takes ownership when it imports the capsule; the destructor checks `PyCapsule_IsValid` to determine whether ownership was transferred and avoids freeing the stream in that case.
* Schema projection is applied per-record-batch (not via in-memory collect) so requested schema changes are cheap and incremental.
* Partition order is preserved: batches are drained partition-by-partition in original order.
* Long-running streaming operations respond correctly to `KeyboardInterrupt` (tested by injecting `PyThreadState_SetAsyncExc`).
## Are these changes tested?
Yes — the PR adds a comprehensive set of tests focused on streaming semantics and correctness. Notable tests:
* `test_iter_batches_dataframe` — verifies iteration yields record batches in order.
* `test_arrow_c_stream_to_table_and_reader` — validates creating a `pyarrow.Table` from the stream and comparing batches.
* `test_arrow_c_stream_order` — ensures partition and batch ordering is preserved.
* `test_arrow_c_stream_schema_selection` — verifies schema projection is respected and returned reader schema matches requested schema.
* `test_arrow_c_stream_schema_mismatch` — ensures an informative error is raised on schema mismatch during requested projection.
* `test_arrow_c_stream_interrupted` — ensures stream iteration stops promptly on `KeyboardInterrupt`.
* `test_arrow_c_stream_large_dataset` — constructs a (virtually) huge range and reads a handful of batches while asserting RSS doesn't increase substantially (uses `psutil`, skipped if unavailable).
* `test_iter_releases_reader` — monkeypatch-based test that ensures the underlying reader is closed when iteration stops.
Tests are added to `python/tests/*` and use `pytest.importorskip` where optional dependencies (`pyarrow.cffi`, `psutil`) are required. The `fail_collect` fixture checks that `collect` is not called accidentally in streaming code paths.
## Are there any user-facing changes?
Yes.
* **New:** `DataFrame` objects are iterable in Python: `for batch in df:` yields `pyarrow.RecordBatch` lazily.
* **New:** `DataFrame.__arrow_c_stream__(requested_schema=None)` now exposes an ArrowArrayStream PyCapsule that PyArrow can consume zero‑copy.
* **Docs:** Added documentation and examples in the DataFrame user guide (`PyArrow Streaming`) showing how to import the capsule into a `pa.RecordBatchReader` and iterate without materializing all batches.
**Compatibility / migration notes**
* This is an additive change and should be backward compatible for existing users who call `collect()` or `to_pylist()`.
* Code that previously relied on `__arrow_c_stream__` materializing results (for example implicitly via earlier implementations) should continue to work — but the streaming path will be used, and tests or mocks that expected `collect()` to be invoked may need adjustment. To help guard against regressions, tests in this PR include a `fail_collect` fixture to ensure streaming paths do not call `collect()`.
## Review notes
Focus review on these high‑leverage areas:
1. **Safety of capsule lifecycle handling** (`src/dataframe.rs`): ensure `drop_stream` correctly handles both stolen and owned pointers, and that Python error state is handled safely.
2. **spawn\_future** (`src/utils.rs`) and its usage: validate signal handling, JoinError mapping, and that error types are converted to Python exceptions cleanly.
3. **Partitioned streaming reader semantics** (`PartitionedDataFrameStreamReader` in `src/dataframe.rs`): confirm the ordering and projection behavior is correct and performant.
4. **Python iterator implementation** (`python/datafusion/dataframe.py`): ensure the reader is imported and closed properly (context manager usage via `closing`), and that `for batch in df:` behaves as expected in IPython and scripts.
5. **Tests**: run the new Python test suite locally (including those that require optional deps) and inspect `fail_collect` behavior.
Files worth opening first in the review UI:
* `src/dataframe.rs`
* `src/record_batch.rs`
* `src/utils.rs`
* `src/context.rs`
* `python/datafusion/dataframe.py`
* `python/tests/test_dataframe.py`
* `python/tests/test_io.py`
* `python/tests/test_dataframe_iter_stream.py`
* `python/tests/utils.py`
## Release notes / changelog suggestion
* Add a note under "Python bindings" saying: "DataFrame now supports zero‑copy Arrow streaming and iteration: `DataFrame.__arrow_c_stream__` returns an ArrowArrayStream PyCapsule and DataFrame objects are iterable (yielding `pyarrow.RecordBatch`). This enables exporting arbitrarily large results to PyArrow without materializing them in memory."
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3275815833
> > I would strongly advocate for less direct integration with pyarrow, not more. Pyarrow is a massive dependency, while the Arrow PyCapsule Interface should allow for better decentralized sharing of Arrow data.
>
> I will work on this as a follow-up task.
Amazing. Happy to review as well. I think it's fine to have _some_ methods that specifically convert to a pyarrow record batch/table. But we should call `import pyarrow` conditionally so that pyarrow doesn't need to be a required dependency all the time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer merged PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3273674991
@kylebarron
Thanks for the review and the many helpful links.
I believe I have addressed your PR feedback points.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2319619566
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,39 @@ 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:`pyarrow.RecordBatch` objects
+lazily so you can loop over results directly:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # process each batch as it is produced
Review Comment:
We already have our own `RecordBatch` class: https://datafusion.apache.org/python/autoapi/datafusion/record_batch/index.html#datafusion.record_batch.RecordBatch
Also, we should ensure that the dunder methods are rendered in the docs. It doesn't look like they are currently. (Or maybe the dunder methods on that `RecordBatch` aren't documented?)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3273888907
@kylebarron
> I would strongly advocate for less direct integration with pyarrow, not more. Pyarrow is a massive dependency, while the Arrow PyCapsule Interface should allow for better decentralized sharing of Arrow data.
I will work on this as a follow-up task.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346253588
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
+
+Asynchronous iteration is supported as well, allowing integration with
+``asyncio`` event loops:
+
+.. code-block:: python
+
+ async for batch in df:
+ ... # process each batch as it is produced
+
+To work with the stream directly, use
+``execute_stream()``, which returns a
+:class:`~datafusion.RecordBatchStream`:
+
+.. code-block:: python
+
+ stream = df.execute_stream()
+ for batch in stream:
+ ...
+
+Execute as Stream
+^^^^^^^^^^^^^^^^^
+
+For finer control over streaming execution, use
+:py:meth:`~datafusion.DataFrame.execute_stream` to obtain a
+:py:class:`pyarrow.RecordBatchReader`:
Review Comment:
You're right!
Corrected.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2319619566
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,39 @@ 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:`pyarrow.RecordBatch` objects
+lazily so you can loop over results directly:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # process each batch as it is produced
Review Comment:
We already have our own `RecordBatch` class: https://datafusion.apache.org/python/autoapi/datafusion/record_batch/index.html#datafusion.record_batch.RecordBatch
Also, we should ensure that the dunder methods are rendered in the docs. It doesn't look like they are currently.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2347314377
##########
python/datafusion/dataframe.py:
##########
@@ -1105,21 +1110,55 @@ 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.
+
+ 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.
- 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.
+ When ``requested_schema`` is provided, DataFusion applies only simple
+ projections such as selecting a subset of existing columns or reordering
+ them. Column renaming, computed expressions, or type coercion are not
+ supported through this interface.
Args:
- requested_schema: Attempt to provide the DataFrame using this schema.
+ requested_schema: Either a :py:class:`pyarrow.Schema` or an Arrow C
+ Schema capsule (``PyCapsule``) produced by
+ ``schema._export_to_c_capsule()``. The DataFrame will attempt to
+ align its output with the fields and order specified by this schema.
Returns:
- Arrow PyCapsule object.
+ Arrow ``PyCapsule`` object representing an ``ArrowArrayStream``.
+
+ Examples:
+ >>> schema = df.schema()
+ >>> stream = df.__arrow_c_stream__(schema)
+ >>> capsule = schema._export_to_c_capsule()
+ >>> stream = df.__arrow_c_stream__(capsule)
+
+ Notes:
+ The Arrow C Data Interface PyCapsule details are documented by Apache
+ Arrow and can be found at:
+ https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
"""
+ # ``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]:
+ """Return an iterator over this DataFrame's record batches."""
+ return iter(self.execute_stream())
+
+ def __aiter__(self) -> AsyncIterator[RecordBatch]:
+ """Return an async iterator over this DataFrame's record batches.
+
+ `RecordBatchStream` implements ``__aiter__``, so return the stream
+ directly to remain compatible with Python < 3.10 (this project
+ supports Python >= 3.6).
Review Comment:
I think our minimum supported version is 3.9 right now
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3250106005
In https://github.com/apache/datafusion-python/issues/1227 I explicitly suggested removing the pyarrow dependency altogether. I thought I had created an issue before, but apparently not.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346183373
##########
python/tests/test_io.py:
##########
@@ -92,3 +96,38 @@ def test_read_avro():
path = Path.cwd() / "testing/data/avro/alltypes_plain.avro"
avro_df = read_avro(path=path)
assert avro_df is not None
+
+
+def test_arrow_c_stream_large_dataset(ctx):
+ """DataFrame.__arrow_c_stream__ yields batches incrementally.
+
+ This test constructs a DataFrame that would be far larger than available
+ memory if materialized. The ``__arrow_c_stream__`` method should expose a
+ stream of record batches without collecting the full dataset, so reading a
+ handful of batches should not exhaust process memory.
+ """
+ # Create a very large DataFrame using range; this would be terabytes if collected
+ df = range_table(ctx, 0, 1 << 40)
+
+ reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
+
+ # Track RSS before consuming batches
+ psutil = pytest.importorskip("psutil")
+ process = psutil.Process()
+ start_rss = process.memory_info().rss
+
+ for _ in range(5):
+ batch = reader.read_next_batch()
+ assert batch is not None
+ assert len(batch) > 0
+ current_rss = process.memory_info().rss
+ # Ensure memory usage hasn't grown substantially (>50MB)
+ assert current_rss - start_rss < 50 * 1024 * 1024
+
+
+def test_table_from_batches_stream(ctx, fail_collect):
+ df = range_table(ctx, 0, 10)
+
+ table = pa.Table.from_batches(batch.to_pyarrow() for batch in df)
Review Comment:
Good point.
I also renamed the test to
test_table_from_arrow_c_stream
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3500290283
@timsaucer
Thanks for the review and feedback.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2469239131
##########
python/datafusion/dataframe.py:
##########
@@ -1131,21 +1136,54 @@ 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.
+
+ 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.
- 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.
+ When ``requested_schema`` is provided, DataFusion applies only simple
+ projections such as selecting a subset of existing columns or reordering
+ them. Column renaming, computed expressions, or type coercion are not
+ supported through this interface.
Args:
- requested_schema: Attempt to provide the DataFrame using this schema.
+ requested_schema: Either a :py:class:`pyarrow.Schema` or an Arrow C
+ Schema capsule (``PyCapsule``) produced by
+ ``schema._export_to_c_capsule()``. The DataFrame will attempt to
+ align its output with the fields and order specified by this schema.
Returns:
- Arrow PyCapsule object.
+ Arrow ``PyCapsule`` object representing an ``ArrowArrayStream``.
+
+ Examples:
+ >>> schema = df.schema()
+ >>> stream = df.__arrow_c_stream__(schema)
+ >>> capsule = schema._export_to_c_capsule()
+ >>> stream = df.__arrow_c_stream__(capsule)
Review Comment:
I think examples should be user facing. I don't think we expect users to ever interact directly with this function. Instead this function should make it seamless to share dataframe with other libraries that take arrow c stream. I recommend removing this example and instead just recommending the online documentation for usage examples.
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -195,10 +195,119 @@ To materialize the results of your DataFrame operations:
# Display results
df.show() # Print tabular format to console
-
+
# Count rows
count = df.count()
+Zero-copy streaming to Arrow-based Python libraries
+---------------------------------------------------
+
+DataFusion DataFrames implement the ``__arrow_c_stream__`` protocol, enabling
+zero-copy, lazy streaming into Arrow-based Python libraries. With the streaming
+protocol, batches are produced on demand so you can process arbitrarily large
+results without out-of-memory errors.
Review Comment:
This is true to a point. It is still possible that the batches can lead to out of memory problems depending on the size of the batches and what the user is doing with them. I recommend removing the part "so you can process arbitrarily large results without out-of-memory errors." Similar to some of my other feedback, this statement mostly is oriented towards people who have already had a memory problem.
##########
docs/source/user-guide/io/arrow.rst:
##########
@@ -60,14 +60,22 @@ Exporting from DataFusion
DataFusion DataFrames implement ``__arrow_c_stream__`` PyCapsule interface, so any
Python library that accepts these can import a DataFusion DataFrame directly.
-.. warning::
- It is important to note that this will cause the DataFrame execution to happen, which may be
- a time consuming task. That is, you will cause a
- :py:func:`datafusion.dataframe.DataFrame.collect` operation call to occur.
+.. note::
+ Invoking ``__arrow_c_stream__`` still triggers execution of the underlying
+ query, but batches are yielded incrementally rather than materialized all at
+ once in memory. Consumers can process the stream as it arrives, avoiding the
+ memory overhead of a full
+ :py:func:`datafusion.dataframe.DataFrame.collect`.
+
+ For an example of this streamed execution and its memory safety, see the
+ ``test_arrow_c_stream_large_dataset`` unit test in
+ :mod:`python.tests.test_io`.
Review Comment:
This documentation is written from the perspective of a developer or user who has already encountered the issue of collection. For new users it probably is not particularly helpful. Instead I would remove the `note` annotator and remove the portion ", avoiding the memory overhead of a full :py:func:`datafusion.dataframe.DataFrame.collect`."
nit: also remove the word "still" in the first line
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346163736
##########
src/record_batch.rs:
##########
@@ -84,15 +84,21 @@ impl PyRecordBatchStream {
}
}
+/// Polls the next batch from a `SendableRecordBatchStream`, converting the `Option<Result<_>>` form.
+pub(crate) async fn poll_next_batch(
+ stream: &mut SendableRecordBatchStream,
+) -> datafusion::error::Result<Option<RecordBatch>> {
+ stream.next().await.transpose()
+}
+
async fn next_stream(
stream: Arc<Mutex<SendableRecordBatchStream>>,
sync: bool,
) -> PyResult<PyRecordBatch> {
let mut stream = stream.lock().await;
- match stream.next().await {
- Some(Ok(batch)) => Ok(batch.into()),
- Some(Err(e)) => Err(PyDataFusionError::from(e))?,
- None => {
+ match poll_next_batch(&mut stream).await {
Review Comment:
`poll_next_batch` was created to encapsulate the recurring asynchronous polling pattern used when consuming the stream. Centralizing this behavior improves readability, prevents duplicated error-mapping logic, facilitates unit testing, and enables future enhancements without modifying multiple call sites.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2337355501
##########
Cargo.toml:
##########
Review Comment:
Are most of the changes here just formatting? Did you just add `cstr`?
##########
python/tests/test_io.py:
##########
@@ -92,3 +96,38 @@ def test_read_avro():
path = Path.cwd() / "testing/data/avro/alltypes_plain.avro"
avro_df = read_avro(path=path)
assert avro_df is not None
+
+
+def test_arrow_c_stream_large_dataset(ctx):
+ """DataFrame.__arrow_c_stream__ yields batches incrementally.
+
+ This test constructs a DataFrame that would be far larger than available
+ memory if materialized. The ``__arrow_c_stream__`` method should expose a
+ stream of record batches without collecting the full dataset, so reading a
+ handful of batches should not exhaust process memory.
+ """
+ # Create a very large DataFrame using range; this would be terabytes if collected
+ df = range_table(ctx, 0, 1 << 40)
+
+ reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
+
+ # Track RSS before consuming batches
+ psutil = pytest.importorskip("psutil")
+ process = psutil.Process()
+ start_rss = process.memory_info().rss
+
+ for _ in range(5):
+ batch = reader.read_next_batch()
+ assert batch is not None
+ assert len(batch) > 0
+ current_rss = process.memory_info().rss
+ # Ensure memory usage hasn't grown substantially (>50MB)
+ assert current_rss - start_rss < 50 * 1024 * 1024
+
+
+def test_table_from_batches_stream(ctx, fail_collect):
+ df = range_table(ctx, 0, 10)
+
+ table = pa.Table.from_batches(batch.to_pyarrow() for batch in df)
Review Comment:
Are you intentionally testing a Python iterator here instead of using `__arrow_c_stream__`?
If you call
```py
pa.table(df)
```
that will use the C Stream and materialize the data on the pyarrow side, and it'll be more efficient because it keeps everything as a C pointer instead of having to go through Python.
##########
python/tests/test_io.py:
##########
@@ -92,3 +96,38 @@ def test_read_avro():
path = Path.cwd() / "testing/data/avro/alltypes_plain.avro"
avro_df = read_avro(path=path)
assert avro_df is not None
+
+
+def test_arrow_c_stream_large_dataset(ctx):
+ """DataFrame.__arrow_c_stream__ yields batches incrementally.
+
+ This test constructs a DataFrame that would be far larger than available
+ memory if materialized. The ``__arrow_c_stream__`` method should expose a
+ stream of record batches without collecting the full dataset, so reading a
+ handful of batches should not exhaust process memory.
+ """
+ # Create a very large DataFrame using range; this would be terabytes if collected
+ df = range_table(ctx, 0, 1 << 40)
+
+ reader = pa.RecordBatchReader._import_from_c_capsule(df.__arrow_c_stream__())
Review Comment:
```suggestion
reader = pa.RecordBatchReader.from_stream(df)
```
##########
python/datafusion/record_batch.py:
##########
@@ -46,6 +46,26 @@ def to_pyarrow(self) -> pa.RecordBatch:
"""Convert to :py:class:`pa.RecordBatch`."""
return self.record_batch.to_pyarrow()
+ def __arrow_c_array__(
Review Comment:
👍 👍
##########
python/datafusion/dataframe.py:
##########
@@ -1105,21 +1122,33 @@ 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
Review Comment:
Might be good to have a link somewhere in the docstring to https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
##########
python/datafusion/dataframe.py:
##########
@@ -1105,21 +1122,33 @@ 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]:
+ """Return an iterator over this DataFrame's record batches."""
+ return iter(self.execute_stream())
+
+ def __aiter__(self) -> AsyncIterator[RecordBatch]:
+ """Return an async iterator over this DataFrame's record batches."""
+ return self.execute_stream().__aiter__()
Review Comment:
Nit, use the same approach for both of these; either use `iter` and `aiter` or call the dunder in each case
```suggestion
return aiter(self.execute_stream())
```
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ To materialize the results of your DataFrame operations:
# Display results
df.show() # Print tabular format to console
-
+
# Count rows
count = df.count()
+PyArrow Streaming
Review Comment:
The heading here is specific to PyArrow, but I think it would be good to provide a distinction here. Maybe "Zero-copy Streaming to other Arrow implementations"? Or something like that?
Then we can have a sub-section dedicated to pyarrow, but also explain that it works with any arrow-based python impl.
##########
src/record_batch.rs:
##########
@@ -84,15 +84,21 @@ impl PyRecordBatchStream {
}
}
+/// Polls the next batch from a `SendableRecordBatchStream`, converting the `Option<Result<_>>` form.
+pub(crate) async fn poll_next_batch(
+ stream: &mut SendableRecordBatchStream,
+) -> datafusion::error::Result<Option<RecordBatch>> {
+ stream.next().await.transpose()
+}
+
async fn next_stream(
stream: Arc<Mutex<SendableRecordBatchStream>>,
sync: bool,
) -> PyResult<PyRecordBatch> {
let mut stream = stream.lock().await;
- match stream.next().await {
- Some(Ok(batch)) => Ok(batch.into()),
- Some(Err(e)) => Err(PyDataFusionError::from(e))?,
- None => {
+ match poll_next_batch(&mut stream).await {
Review Comment:
I'm not a tokio expert; how does this materially change?
Did you have to make a new function here? Could you have just used
```rs
match stream.next().await.transpose()
```
##########
python/datafusion/dataframe.py:
##########
@@ -1042,6 +1047,18 @@ def execute_stream_partitioned(self) -> list[RecordBatchStream]:
streams = self.df.execute_stream_partitioned()
return [RecordBatchStream(rbs) for rbs in streams]
+ @deprecated("Use execute_stream() instead")
+ def to_record_batch_stream(self) -> RecordBatchStream:
Review Comment:
Is this a new method? Why are we creating a new method that's immediately deprecated?
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
+
+Asynchronous iteration is supported as well, allowing integration with
+``asyncio`` event loops:
+
+.. code-block:: python
+
+ async for batch in df:
+ ... # process each batch as it is produced
+
+To work with the stream directly, use
+``execute_stream()``, which returns a
+:class:`~datafusion.RecordBatchStream`:
+
+.. code-block:: python
+
+ stream = df.execute_stream()
+ for batch in stream:
+ ...
+
+Execute as Stream
+^^^^^^^^^^^^^^^^^
+
+For finer control over streaming execution, use
+:py:meth:`~datafusion.DataFrame.execute_stream` to obtain a
+:py:class:`pyarrow.RecordBatchReader`:
Review Comment:
`execute_stream` returns a _`pyarrow`_ `RecordBatchReader`? I thought it returned our own DataFusion-specific stream wrapper? I.e. `datafusion.RecordBatchStream`
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
Review Comment:
I think if we provide an example of collecting as a table, we should suggest `pa.table(df)` instead, which will be more efficient because it doesn't go through Python.
And
> without collecting everything eagerly
This isn't a good example for that, because `pa.Table` _necessarily_ materializes everything
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
+
+Asynchronous iteration is supported as well, allowing integration with
+``asyncio`` event loops:
+
+.. code-block:: python
+
+ async for batch in df:
+ ... # process each batch as it is produced
+
+To work with the stream directly, use
+``execute_stream()``, which returns a
+:class:`~datafusion.RecordBatchStream`:
+
+.. code-block:: python
+
+ stream = df.execute_stream()
+ for batch in stream:
+ ...
+
+Execute as Stream
+^^^^^^^^^^^^^^^^^
+
+For finer control over streaming execution, use
+:py:meth:`~datafusion.DataFrame.execute_stream` to obtain a
+:py:class:`pyarrow.RecordBatchReader`:
+
+.. code-block:: python
+
+ reader = df.execute_stream()
+ for batch in reader:
+ ... # process each batch as it is produced
+
+When partition boundaries are important,
+:py:meth:`~datafusion.DataFrame.execute_stream_partitioned`
+returns an iterable of :py:class:`pyarrow.RecordBatchReader` objects, one per
+partition:
+
+.. code-block:: python
+
+ for stream in df.execute_stream_partitioned():
+ for batch in stream:
Review Comment:
Interesting. Can these streams be polled concurrently? Can you do
```py
streams = list(df.execute_stream_partitioned())
```
and then concurrently iterate over all the streams, yielding whatever batch comes in first? I suppose that would just do in Python what `execute_stream` is doing in Rust?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3250108342
This would also close https://github.com/apache/datafusion-python/issues/1011
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2319615655
##########
python/datafusion/dataframe.py:
##########
@@ -1098,21 +1102,42 @@ 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[pa.RecordBatch]:
Review Comment:
`RecordBatchStream` already has `__iter__` and `__aiter__` methods https://datafusion.apache.org/python/autoapi/datafusion/record_batch/index.html#datafusion.record_batch.RecordBatchStream
Can we just have a method that converts a `DataFrame` into a `RecordBatchStream`? Then an `__iter__` on `DataFrame` would just convert to a `RecordBatchStream` under the hood.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346199959
##########
python/datafusion/dataframe.py:
##########
@@ -1105,21 +1122,33 @@ 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
Review Comment:
Added
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346258279
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
Review Comment:
corrected.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3245924248
I'm invested in this and plan to review this this afternoon!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346250200
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ 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.from_stream(df)
+ for batch in reader:
+ ... # process each batch as it is produced
+
+DataFrames are also iterable, yielding :class:`datafusion.RecordBatch`
+objects lazily so you can loop over results directly without importing
+PyArrow:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # each batch is a ``datafusion.RecordBatch``
+
+Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow
+table without collecting everything eagerly:
+
+.. code-block:: python
+
+ import pyarrow as pa
+ table = pa.Table.from_batches(b.to_pyarrow() for b in df)
+
+Asynchronous iteration is supported as well, allowing integration with
+``asyncio`` event loops:
+
+.. code-block:: python
+
+ async for batch in df:
+ ... # process each batch as it is produced
+
+To work with the stream directly, use
+``execute_stream()``, which returns a
+:class:`~datafusion.RecordBatchStream`:
+
+.. code-block:: python
+
+ stream = df.execute_stream()
+ for batch in stream:
+ ...
+
+Execute as Stream
+^^^^^^^^^^^^^^^^^
+
+For finer control over streaming execution, use
+:py:meth:`~datafusion.DataFrame.execute_stream` to obtain a
+:py:class:`pyarrow.RecordBatchReader`:
+
+.. code-block:: python
+
+ reader = df.execute_stream()
+ for batch in reader:
+ ... # process each batch as it is produced
+
+When partition boundaries are important,
+:py:meth:`~datafusion.DataFrame.execute_stream_partitioned`
+returns an iterable of :py:class:`pyarrow.RecordBatchReader` objects, one per
+partition:
+
+.. code-block:: python
+
+ for stream in df.execute_stream_partitioned():
+ for batch in stream:
Review Comment:
Good question!
I added a concurrent iteration example in the same document to clarify this.
```
To process partitions concurrently, first collect the streams into a list
and then poll each one in a separate ``asyncio`` task:
.. code-block:: python
import asyncio
async def consume(stream):
async for batch in stream:
...
streams = list(df.execute_stream_partitioned())
await asyncio.gather(*(consume(s) for s in streams))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2347912854
##########
python/datafusion/dataframe.py:
##########
@@ -1105,21 +1110,55 @@ 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.
+
+ 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.
- 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.
+ When ``requested_schema`` is provided, DataFusion applies only simple
+ projections such as selecting a subset of existing columns or reordering
+ them. Column renaming, computed expressions, or type coercion are not
+ supported through this interface.
Args:
- requested_schema: Attempt to provide the DataFrame using this schema.
+ requested_schema: Either a :py:class:`pyarrow.Schema` or an Arrow C
+ Schema capsule (``PyCapsule``) produced by
+ ``schema._export_to_c_capsule()``. The DataFrame will attempt to
+ align its output with the fields and order specified by this schema.
Returns:
- Arrow PyCapsule object.
+ Arrow ``PyCapsule`` object representing an ``ArrowArrayStream``.
+
+ Examples:
+ >>> schema = df.schema()
+ >>> stream = df.__arrow_c_stream__(schema)
+ >>> capsule = schema._export_to_c_capsule()
+ >>> stream = df.__arrow_c_stream__(capsule)
+
+ Notes:
+ The Arrow C Data Interface PyCapsule details are documented by Apache
+ Arrow and can be found at:
+ https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
"""
+ # ``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]:
+ """Return an iterator over this DataFrame's record batches."""
+ return iter(self.execute_stream())
+
+ def __aiter__(self) -> AsyncIterator[RecordBatch]:
+ """Return an async iterator over this DataFrame's record batches.
+
+ `RecordBatchStream` implements ``__aiter__``, so return the stream
+ directly to remain compatible with Python < 3.10 (this project
+ supports Python >= 3.6).
Review Comment:
I will amend the comment
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3499167366
Thank you for all the work on this. As soon as CI goes green, I'll hit the merge button.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2483157734
##########
python/datafusion/dataframe.py:
##########
@@ -1131,21 +1136,54 @@ 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.
+
+ 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.
- 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.
+ When ``requested_schema`` is provided, DataFusion applies only simple
+ projections such as selecting a subset of existing columns or reordering
+ them. Column renaming, computed expressions, or type coercion are not
+ supported through this interface.
Args:
- requested_schema: Attempt to provide the DataFrame using this schema.
+ requested_schema: Either a :py:class:`pyarrow.Schema` or an Arrow C
+ Schema capsule (``PyCapsule``) produced by
+ ``schema._export_to_c_capsule()``. The DataFrame will attempt to
+ align its output with the fields and order specified by this schema.
Returns:
- Arrow PyCapsule object.
+ Arrow ``PyCapsule`` object representing an ``ArrowArrayStream``.
+
+ Examples:
+ >>> schema = df.schema()
+ >>> stream = df.__arrow_c_stream__(schema)
+ >>> capsule = schema._export_to_c_capsule()
+ >>> stream = df.__arrow_c_stream__(capsule)
Review Comment:
Amended
##########
docs/source/user-guide/io/arrow.rst:
##########
@@ -60,14 +60,22 @@ Exporting from DataFusion
DataFusion DataFrames implement ``__arrow_c_stream__`` PyCapsule interface, so any
Python library that accepts these can import a DataFusion DataFrame directly.
-.. warning::
- It is important to note that this will cause the DataFrame execution to happen, which may be
- a time consuming task. That is, you will cause a
- :py:func:`datafusion.dataframe.DataFrame.collect` operation call to occur.
+.. note::
+ Invoking ``__arrow_c_stream__`` still triggers execution of the underlying
+ query, but batches are yielded incrementally rather than materialized all at
+ once in memory. Consumers can process the stream as it arrives, avoiding the
+ memory overhead of a full
+ :py:func:`datafusion.dataframe.DataFrame.collect`.
+
+ For an example of this streamed execution and its memory safety, see the
+ ``test_arrow_c_stream_large_dataset`` unit test in
+ :mod:`python.tests.test_io`.
Review Comment:
Amended
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Expose Arrow C stream and DataFrame iterator (zero‑copy streaming to PyArrow) [datafusion-python]
Posted by "kylebarron (via GitHub)" <gi...@apache.org>.
kylebarron commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2319535613
##########
python/datafusion/dataframe.py:
##########
@@ -1098,21 +1102,42 @@ 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[pa.RecordBatch]:
Review Comment:
I don't really think there's a good rationale for having this method, especially as it reuses the exact same mechanism as the PyCapsule Interface. If anything, we might want to have an `__aiter__` method that has a custom async connection to the DataFusion context.
##########
src/dataframe.rs:
##########
@@ -39,25 +39,55 @@ use datafusion::prelude::*;
use datafusion_ffi::table_provider::FFI_TableProvider;
use futures::{StreamExt, TryStreamExt};
use pyo3::exceptions::PyValueError;
+use pyo3::ffi;
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedStr;
use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
-use tokio::task::JoinHandle;
use crate::catalog::PyTable;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError};
+use crate::errors::{py_datafusion_err, PyDataFusionError};
use crate::expr::sort_expr::to_sort_expressions;
use crate::physical_plan::PyExecutionPlan;
-use crate::record_batch::PyRecordBatchStream;
+use crate::record_batch::{poll_next_batch, PyRecordBatchStream};
use crate::sql::logical::PyLogicalPlan;
use crate::utils::{
- get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
+ get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_future, validate_pycapsule,
+ wait_for_future,
};
use crate::{
errors::PyDataFusionResult,
expr::{sort_expr::PySortExpr, PyExpr},
};
+#[allow(clippy::manual_c_str_literals)]
+static ARROW_STREAM_NAME: &CStr =
+ unsafe { CStr::from_bytes_with_nul_unchecked(b"arrow_array_stream\0") };
Review Comment:
As suggested by the linter, can we just use `c"arrow_array_stream"`?
##########
src/dataframe.rs:
##########
@@ -39,25 +39,55 @@ use datafusion::prelude::*;
use datafusion_ffi::table_provider::FFI_TableProvider;
use futures::{StreamExt, TryStreamExt};
use pyo3::exceptions::PyValueError;
+use pyo3::ffi;
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedStr;
use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
-use tokio::task::JoinHandle;
use crate::catalog::PyTable;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError};
+use crate::errors::{py_datafusion_err, PyDataFusionError};
use crate::expr::sort_expr::to_sort_expressions;
use crate::physical_plan::PyExecutionPlan;
-use crate::record_batch::PyRecordBatchStream;
+use crate::record_batch::{poll_next_batch, PyRecordBatchStream};
use crate::sql::logical::PyLogicalPlan;
use crate::utils::{
- get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, validate_pycapsule, wait_for_future,
+ get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_future, validate_pycapsule,
+ wait_for_future,
};
use crate::{
errors::PyDataFusionResult,
expr::{sort_expr::PySortExpr, PyExpr},
};
+#[allow(clippy::manual_c_str_literals)]
+static ARROW_STREAM_NAME: &CStr =
+ unsafe { CStr::from_bytes_with_nul_unchecked(b"arrow_array_stream\0") };
+
+unsafe extern "C" fn drop_stream(capsule: *mut ffi::PyObject) {
+ if capsule.is_null() {
+ return;
+ }
+
+ // When PyArrow imports this capsule it steals the raw stream pointer and
+ // sets the capsule's internal pointer to NULL. In that case
+ // `PyCapsule_IsValid` returns 0 and this destructor must not drop the
+ // stream as ownership has been transferred to PyArrow. If the capsule was
+ // never imported, the pointer remains valid and we are responsible for
+ // freeing the stream here.
+ if ffi::PyCapsule_IsValid(capsule, ARROW_STREAM_NAME.as_ptr()) == 1 {
+ let stream_ptr = ffi::PyCapsule_GetPointer(capsule, ARROW_STREAM_NAME.as_ptr())
+ as *mut FFI_ArrowArrayStream;
+ if !stream_ptr.is_null() {
+ drop(Box::from_raw(stream_ptr));
+ }
+ }
+
+ // `PyCapsule_GetPointer` sets a Python error on failure. Clear it only
+ // after the stream has been released (or determined to be owned
+ // elsewhere).
+ ffi::PyErr_Clear();
+}
Review Comment:
We shouldn't need to do any of this, according to upstream discussion https://github.com/apache/arrow-rs/pull/5070#discussion_r1391489292
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,39 @@ 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:`pyarrow.RecordBatch` objects
+lazily so you can loop over results directly:
+
+.. code-block:: python
+
+ for batch in df:
+ ... # process each batch as it is produced
Review Comment:
Because the user can iterate over the stream accessed by the target library, I don't think we should define our own custom integration here, and if we do, then the yielded object should not be a pyarrow RecordBatch, but rather an opaque, minimal Python class that just exposes `__arrow_c_array__` so that the user can choose what Arrow library they want to use to work with the batch.
##########
src/dataframe.rs:
##########
@@ -354,6 +384,63 @@ impl PyDataFrame {
}
}
+/// Synchronous wrapper around partitioned [`SendableRecordBatchStream`]s used
+/// for the `__arrow_c_stream__` implementation.
+///
+/// It drains each partition's stream sequentially, yielding record batches in
+/// their original partition order. When a `projection` is set, each batch is
+/// converted via `record_batch_into_schema` to apply schema changes per batch.
+struct PartitionedDataFrameStreamReader {
+ streams: Vec<SendableRecordBatchStream>,
+ schema: SchemaRef,
+ projection: Option<SchemaRef>,
+ current: usize,
+}
+
+impl Iterator for PartitionedDataFrameStreamReader {
+ type Item = Result<RecordBatch, ArrowError>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ while self.current < self.streams.len() {
+ let stream = &mut self.streams[self.current];
+ let fut = poll_next_batch(stream);
+ let result = Python::with_gil(|py| wait_for_future(py, fut));
+
+ match result {
+ Ok(Ok(Some(batch))) => {
+ let batch = if let Some(ref schema) = self.projection {
+ match record_batch_into_schema(batch, schema.as_ref()) {
+ Ok(b) => b,
+ Err(e) => return Some(Err(e)),
+ }
+ } else {
+ batch
+ };
+ return Some(Ok(batch));
+ }
+ Ok(Ok(None)) => {
+ self.current += 1;
+ continue;
+ }
+ Ok(Err(e)) => {
+ return Some(Err(ArrowError::ExternalError(Box::new(e))));
+ }
+ Err(e) => {
+ return Some(Err(ArrowError::ExternalError(Box::new(e))));
+ }
+ }
+ }
+
+ None
+ }
+}
+
+impl RecordBatchReader for PartitionedDataFrameStreamReader {
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
+
#[pymethods]
impl PyDataFrame {
Review Comment:
Essentially this changes the `DataFrame` construct to always be "lazy"? Previously a `DataFrame` was always materialized in memory, whereas now it's just a representation of future batches?
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,39 @@ 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__())
Review Comment:
```suggestion
reader = pa.RecordBatchReader.from_stream(df)
```
##########
src/dataframe.rs:
##########
@@ -889,44 +979,55 @@ impl PyDataFrame {
let desired_schema = Schema::try_from(schema_ptr)?;
schema = project_schema(schema, desired_schema)?;
-
- batches = batches
- .into_iter()
- .map(|record_batch| record_batch_into_schema(record_batch, &schema))
- .collect::<Result<Vec<RecordBatch>, ArrowError>>()?;
+ projection = Some(Arc::new(schema.clone()));
}
- let batches_wrapped = batches.into_iter().map(Ok);
+ let schema_ref = Arc::new(schema.clone());
- let reader = RecordBatchIterator::new(batches_wrapped, Arc::new(schema));
+ let reader = PartitionedDataFrameStreamReader {
+ streams,
+ schema: schema_ref,
+ projection,
+ current: 0,
+ };
let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
- let ffi_stream = FFI_ArrowArrayStream::new(reader);
- let stream_capsule_name = CString::new("arrow_array_stream").unwrap();
- PyCapsule::new(py, ffi_stream, Some(stream_capsule_name)).map_err(PyDataFusionError::from)
+ let stream = Box::new(FFI_ArrowArrayStream::new(reader));
Review Comment:
If you have an `FFI_ArrowArrayStream` you should be able to just pass that to `PyCapsule::new` without touching any unsafe: https://github.com/kylebarron/arro3/blob/cb2453bf022d0d8704e56e81a324ab5a772e0247/pyo3-arrow/src/ffi/to_python/utils.rs#L93-L94
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#issuecomment-3274544611
I'm sorry I've not been following this one closely, but I hope to look into this tomorrow.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346212351
##########
python/datafusion/dataframe.py:
##########
@@ -1042,6 +1047,18 @@ def execute_stream_partitioned(self) -> list[RecordBatchStream]:
streams = self.df.execute_stream_partitioned()
return [RecordBatchStream(rbs) for rbs in streams]
+ @deprecated("Use execute_stream() instead")
+ def to_record_batch_stream(self) -> RecordBatchStream:
Review Comment:
🤦
Removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2346273569
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -145,10 +145,90 @@ To materialize the results of your DataFrame operations:
# Display results
df.show() # Print tabular format to console
-
+
# Count rows
count = df.count()
+PyArrow Streaming
Review Comment:
Good point!
Implemented your suggestion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [datafusion-python]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #1222:
URL: https://github.com/apache/datafusion-python/pull/1222#discussion_r2483148456
##########
docs/source/user-guide/dataframe/index.rst:
##########
@@ -195,10 +195,119 @@ To materialize the results of your DataFrame operations:
# Display results
df.show() # Print tabular format to console
-
+
# Count rows
count = df.count()
+Zero-copy streaming to Arrow-based Python libraries
+---------------------------------------------------
+
+DataFusion DataFrames implement the ``__arrow_c_stream__`` protocol, enabling
+zero-copy, lazy streaming into Arrow-based Python libraries. With the streaming
+protocol, batches are produced on demand so you can process arbitrarily large
+results without out-of-memory errors.
Review Comment:
Amended
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org