-
Notifications
You must be signed in to change notification settings - Fork 575
feat: support arrow pycapsule streams #3447
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -150,7 +150,7 @@ | |
| from pyiceberg.table.name_mapping import NameMapping, apply_name_mapping | ||
| from pyiceberg.table.puffin import PuffinFile | ||
| from pyiceberg.transforms import IdentityTransform, TruncateTransform | ||
| from pyiceberg.typedef import EMPTY_DICT, Properties, Record, TableVersion | ||
| from pyiceberg.typedef import EMPTY_DICT, ArrowStreamExportable, Properties, Record, TableVersion | ||
| from pyiceberg.types import ( | ||
| BinaryType, | ||
| BooleanType, | ||
|
|
@@ -2690,30 +2690,45 @@ def bin_pack_arrow_table(tbl: pa.Table, target_file_size: int) -> Iterator[list[ | |
| """Bin-pack ``tbl`` into groups of RecordBatches, each ~``target_file_size``. | ||
|
|
||
| Note: | ||
| ``target_file_size`` is measured in **uncompressed in-memory** Arrow bytes | ||
| (``Table.nbytes`` / ``RecordBatch.nbytes``), not compressed on-disk Parquet | ||
| bytes. The resulting Parquet file after compression (zstd by default, | ||
| plus dictionary/RLE encoding) is typically 3-10× smaller than | ||
| ``target_file_size``. This is a coarse proxy for the spec-defined | ||
| ``target_file_size`` is measured in **uncompressed in-memory** Arrow | ||
| bytes, not compressed on-disk Parquet bytes. The size estimate uses | ||
| ``nbytes`` when available and falls back to referenced buffer size for | ||
| Arrow view types that do not support ``nbytes``. The resulting Parquet | ||
| file after compression (zstd by default, plus dictionary/RLE encoding) | ||
| is typically 3-10× smaller than ``target_file_size``. This is a coarse | ||
| proxy for the spec-defined | ||
| ``write.target-file-size-bytes`` and will be tightened to true on-disk | ||
| bytes once the writer is switched to a rolling-``ParquetWriter`` with | ||
| ``OutputStream.tell()`` (#2998). | ||
| """ | ||
| from pyiceberg.utils.bin_packing import PackingIterator | ||
|
|
||
| avg_row_size_bytes = tbl.nbytes / tbl.num_rows | ||
| avg_row_size_bytes = _arrow_data_size(tbl) / tbl.num_rows | ||
| target_rows_per_file = max(1, int(target_file_size / avg_row_size_bytes)) | ||
| batches = tbl.to_batches(max_chunksize=target_rows_per_file) | ||
| bin_packed_record_batches = PackingIterator( | ||
| items=batches, | ||
| target_weight=target_file_size, | ||
| lookback=len(batches), # ignore lookback | ||
| weight_func=lambda x: x.nbytes, | ||
| weight_func=_arrow_data_size, | ||
| largest_bin_first=False, | ||
| ) | ||
| return bin_packed_record_batches | ||
|
|
||
|
|
||
| def _arrow_data_size(data: pa.Table | pa.RecordBatch) -> int: | ||
| """Estimate Arrow data size for writer bin-packing. | ||
|
|
||
| ``nbytes`` is the better logical-size estimate, but PyArrow can raise for | ||
| view types such as ``string_view`` exported by libraries like Polars. Fall | ||
| back to total referenced buffer size so those streams can still be written. | ||
| """ | ||
| try: | ||
| return data.nbytes | ||
| except pyarrow.lib.ArrowTypeError: | ||
| return data.get_total_buffer_size() | ||
|
|
||
|
|
||
| def bin_pack_record_batches(batches: Iterable[pa.RecordBatch], target_file_size: int) -> Iterator[list[pa.RecordBatch]]: | ||
| """Microbatch a single-pass stream of RecordBatches into target-sized groups. | ||
|
|
||
|
|
@@ -2729,9 +2744,11 @@ def bin_pack_record_batches(batches: Iterable[pa.RecordBatch], target_file_size: | |
|
|
||
| Note: | ||
| ``target_file_size`` is measured in **uncompressed in-memory** Arrow | ||
| bytes (``RecordBatch.nbytes``), not compressed on-disk Parquet bytes. | ||
| The resulting Parquet file after compression is typically 3-10× | ||
| smaller than ``target_file_size``. Matches the existing | ||
| bytes, not compressed on-disk Parquet bytes. The size estimate uses | ||
| ``nbytes`` when available and falls back to referenced buffer size for | ||
| Arrow view types that do not support ``nbytes``. The resulting Parquet | ||
| file after compression is typically 3-10× smaller than | ||
| ``target_file_size``. Matches the existing | ||
| :func:`bin_pack_arrow_table` semantics; both will be tightened to true | ||
| on-disk bytes once the writer is switched to a rolling- | ||
| ``ParquetWriter`` with ``OutputStream.tell()`` (#2998). | ||
|
|
@@ -2740,7 +2757,7 @@ def bin_pack_record_batches(batches: Iterable[pa.RecordBatch], target_file_size: | |
| buffer_bytes = 0 | ||
| for batch in batches: | ||
| buffer.append(batch) | ||
| buffer_bytes += batch.nbytes | ||
| buffer_bytes += _arrow_data_size(batch) | ||
| if buffer_bytes >= target_file_size: | ||
| yield buffer | ||
| buffer = [] | ||
|
|
@@ -3043,3 +3060,23 @@ def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Ar | |
| field_array = arrow_table[path_parts[0]] | ||
| # Navigate into the struct using the remaining path parts | ||
| return pc.struct_field(field_array, path_parts[1:]) | ||
|
|
||
|
|
||
| def _coerce_arrow_input(df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable) -> pa.Table | pa.RecordBatchReader: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We're importing this all over the place. Might as well make it public. |
||
| """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}" | ||
| ) | ||
|
Comment on lines
+3064
to
+3082
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This looks to be the core change in this PR and looks valid 👍 |
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -89,6 +89,7 @@ | |
| from pyiceberg.transforms import IdentityTransform | ||
| from pyiceberg.typedef import ( | ||
| EMPTY_DICT, | ||
| ArrowStreamExportable, | ||
| IcebergBaseModel, | ||
| IcebergRootModel, | ||
| Identifier, | ||
|
|
@@ -459,7 +460,7 @@ def update_statistics(self) -> UpdateStatistics: | |
|
|
||
| def append( | ||
| self, | ||
| df: pa.Table | pa.RecordBatchReader, | ||
| df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable, | ||
| snapshot_properties: dict[str, str] = EMPTY_DICT, | ||
| branch: str | None = MAIN_BRANCH, | ||
| ) -> None: | ||
|
|
@@ -512,10 +513,9 @@ def append( | |
| except ModuleNotFoundError as e: | ||
| raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e | ||
|
|
||
| from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _dataframe_to_data_files | ||
| from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _coerce_arrow_input, _dataframe_to_data_files | ||
|
|
||
| if not isinstance(df, (pa.Table, pa.RecordBatchReader)): | ||
| raise ValueError(f"Expected pa.Table or pa.RecordBatchReader, got: {df}") | ||
| df = _coerce_arrow_input(df) | ||
|
|
||
| downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False | ||
| _check_pyarrow_schema_compatible( | ||
|
|
@@ -605,7 +605,7 @@ def dynamic_partition_overwrite( | |
|
|
||
| def overwrite( | ||
| self, | ||
| df: pa.Table | pa.RecordBatchReader, | ||
| df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable, | ||
| overwrite_filter: BooleanExpression | str = ALWAYS_TRUE, | ||
| snapshot_properties: dict[str, str] = EMPTY_DICT, | ||
| case_sensitive: bool = True, | ||
|
|
@@ -669,10 +669,9 @@ def overwrite( | |
| except ModuleNotFoundError as e: | ||
| raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e | ||
|
|
||
| from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _dataframe_to_data_files | ||
| from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, _coerce_arrow_input, _dataframe_to_data_files | ||
|
|
||
| if not isinstance(df, (pa.Table, pa.RecordBatchReader)): | ||
| raise ValueError(f"Expected pa.Table or pa.RecordBatchReader, got: {df}") | ||
| df = _coerce_arrow_input(df) | ||
|
|
||
| downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False | ||
| _check_pyarrow_schema_compatible( | ||
|
|
@@ -1534,7 +1533,7 @@ def upsert( | |
|
|
||
| def append( | ||
| self, | ||
| df: pa.Table | pa.RecordBatchReader, | ||
| df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable, | ||
| snapshot_properties: dict[str, str] = EMPTY_DICT, | ||
| branch: str | None = MAIN_BRANCH, | ||
| ) -> None: | ||
|
|
@@ -1569,7 +1568,7 @@ def dynamic_partition_overwrite( | |
|
|
||
| def overwrite( | ||
| self, | ||
| df: pa.Table | pa.RecordBatchReader, | ||
| df: pa.Table | pa.RecordBatchReader | ArrowStreamExportable, | ||
| overwrite_filter: BooleanExpression | str = ALWAYS_TRUE, | ||
| snapshot_properties: dict[str, str] = EMPTY_DICT, | ||
| case_sensitive: bool = True, | ||
|
|
@@ -1778,6 +1777,10 @@ def __datafusion_table_provider__(self, session: Any | None = None) -> IcebergDa | |
| ).__datafusion_table_provider__ | ||
| return provider(session) | ||
|
|
||
| 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) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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 |
||
|
|
||
|
|
||
| class StaticTable(Table): | ||
| """Load a table directly from a metadata file (i.e., without using a catalog).""" | ||
|
|
@@ -1908,6 +1911,13 @@ def plan_files(self) -> Iterable[ScanTask]: ... | |
| @abstractmethod | ||
| def to_arrow(self) -> pa.Table: ... | ||
|
|
||
| @abstractmethod | ||
| def to_arrow_batch_reader(self) -> pa.RecordBatchReader: ... | ||
|
|
||
| def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: | ||
| """Export this scan's result as an Arrow C stream (PyCapsule interface).""" | ||
| return self.to_arrow_batch_reader().__arrow_c_stream__(requested_schema) | ||
|
|
||
| def update(self: A, **overrides: Any) -> A: | ||
| """Create a copy of this table scan with updated fields.""" | ||
| from inspect import signature | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.)