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. |
2 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.