feat: support arrow pycapsule streams - #3447
Conversation
|
This pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that's incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions. |
271f50d to
38e61ea
Compare
Adopt the Arrow PyCapsule interface on both sides of the read/write boundary. append/overwrite accept any object implementing __arrow_c_stream__ (coerced to a streaming RecordBatchReader), and Table plus every scan expose __arrow_c_stream__ via BaseScan (not just DataScan) so any Arrow consumer can ingest them.
38e61ea to
6b1fbd8
Compare
|
Still active - happy to address any feedback. |
|
|
||
| def _coerce_arrow_input(df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable) -> pa.Table | pa.RecordBatchReader: | ||
| """Normalize Arrow write input to a pa.Table or pa.RecordBatchReader. | ||
|
|
||
| Native pyarrow inputs pass through unchanged; any object implementing the | ||
| Arrow PyCapsule stream interface (``__arrow_c_stream__``) is imported as a | ||
| streaming RecordBatchReader. | ||
| """ | ||
| if isinstance(df, (pa.Table, pa.RecordBatchReader)): | ||
| return df | ||
|
|
||
| # Any object implementing the Arrow PyCapsule stream interface. | ||
| if hasattr(df, "__arrow_c_stream__"): | ||
| return pa.RecordBatchReader.from_stream(df) | ||
|
|
||
| raise ValueError( | ||
| f"Expected pa.Table, pa.RecordBatchReader, or an object implementing the " | ||
| f"Arrow PyCapsule interface (__arrow_c_stream__), got: {df!r}" | ||
| ) |
There was a problem hiding this comment.
This looks to be the core change in this PR and looks valid 👍
|
This pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that's incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions. |
|
Still active - happy to address any feedback. |
|
This pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that's incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions. |
|
still active |
| return pc.struct_field(field_array, path_parts[1:]) | ||
|
|
||
|
|
||
| def _coerce_arrow_input(df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable) -> pa.Table | pa.RecordBatchReader: |
There was a problem hiding this comment.
We're importing this all over the place. Might as well make it public.
| try: | ||
| return data.nbytes | ||
| except pyarrow.lib.ArrowTypeError: | ||
| return data.get_total_buffer_size() |
There was a problem hiding this comment.
This is an overestimation. Is there a way for us to get more precise?
(I recognize this is just the case for the stream case, so an overestimation is probably fine. It's not changing the existing code paths.)
|
|
||
| def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: | ||
| """Export this Table as an Arrow C stream (PyCapsule interface).""" | ||
| return self.scan().to_arrow_batch_reader().__arrow_c_stream__(requested_schema) |
There was a problem hiding this comment.
This is the part I'm most concerned about and something I'd like the community's input on.
By turning a Table into a PyCapsule interface, we can now append or overwrite a Table with another Table. We explicitly blocked this earlier. Is this something we're okay enabling?
3 participants
Closes #2680
Closes #1655
Rationale for this change
PyIceberg is coupled to PyArrow at its read/write boundary:
append/overwritereject anything that isn't apa.Table/pa.RecordBatchReader, and external Arrow consumers can't read a table/scan withoutto_arrow(). Users of other Arrow-native libraries (polars, arro3, nanoarrow, …) therefore have to convert to PyArrow explicitly.This PR adopts the Arrow PyCapsule interface on both sides:
append/overwriteaccept any object implementing__arrow_c_stream__, in addition to PyArrow types.TableandDataScanimplement__arrow_c_stream__, so they can be handed to any Arrow consumer.PyArrow inputs are unchanged; other producers are wrapped as a streaming
RecordBatchReader. PyArrow remains an internal write dependency, only the caller-side requirement is removed.Side effect: bin-packing now falls back to referenced buffer size for Arrow view types (e.g.
string_view) that PyArrow can't size vianbytes, since recent Polars exports produce them.Not in scope:
upsert/dynamic_partition_overwritestill require a materializedpa.Table(they do random access/joins, not streaming).append/overwriteon a partitioned table still raisesNotImplementedError, the same restriction aspa.RecordBatchReadertoday.pa.Tablewrites are unaffected either way.Are these changes tested?
Yes.
tests/table/test_arrow_capsule.py(no Docker) covers coercion,append/overwriteacross all input forms, the partitioned regression, and round-trips throughpa.table().tests/io/test_pyarrow.pycovers thestring_viewbin-packing fallback.Are there any user-facing changes?
Yes, additive and backwards compatible.
append/overwriteaccept Arrow PyCapsule producers;Table/DataScanimplement__arrow_c_stream__. No change for existing PyArrow inputs.