From dc4e88c84613191b0ae030cc5ede8609776af030 Mon Sep 17 00:00:00 2001 From: adon-metcalfe Date: Sat, 22 Aug 2026 07:47:17 +0800 Subject: [PATCH] Honor identity sort orders for Arrow table writes --- pyiceberg/io/pyarrow.py | 59 ++++++++++++++++++-- tests/integration/test_writes/test_writes.py | 23 +++++++- tests/io/test_pyarrow.py | 26 ++++++++- 3 files changed, 100 insertions(+), 8 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..4d06a9ee72 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -152,6 +152,7 @@ from pyiceberg.table.metadata import TableMetadata from pyiceberg.table.name_mapping import NameMapping, apply_name_mapping from pyiceberg.table.puffin import PuffinFile +from pyiceberg.table.sorting import NullOrder, SortDirection from pyiceberg.transforms import IdentityTransform, TruncateTransform from pyiceberg.typedef import EMPTY_DICT, Properties, Record, TableVersion from pyiceberg.types import ( @@ -2763,10 +2764,7 @@ def write_data_file(task: WriteTask) -> DataFile: file_format=file_format, partition=task.partition_key.partition if task.partition_key else Record(), file_size_in_bytes=len(fo), - # After this has been fixed: - # https://github.com/apache/iceberg-python/issues/271 - # sort_order_id=task.sort_order_id, - sort_order_id=None, + sort_order_id=task.sort_order_id, # Just copy these from the table for now spec_id=table_metadata.default_spec_id, equality_ids=None, @@ -2998,6 +2996,7 @@ def _dataframe_to_data_files( downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, format_version=table_metadata.format_version, ) + sort_order = table_metadata.sort_order() if isinstance(df, pa.RecordBatchReader): if not table_metadata.spec().is_unpartitioned(): @@ -3006,6 +3005,11 @@ def _dataframe_to_data_files( "Materialise the reader as a pa.Table first, or follow " "https://github.com/apache/iceberg-python/issues/2152 for partitioned streaming support." ) + if not sort_order.is_unsorted: + warnings.warn( + "Sort order is not applied to streaming RecordBatchReader writes; data files are marked unsorted", + stacklevel=2, + ) yield from write_file( io=io, table_metadata=table_metadata, @@ -3017,11 +3021,18 @@ def _dataframe_to_data_files( return if table_metadata.spec().is_unpartitioned(): + df, sort_order_id = _sort_table_for_write(table_metadata, df) yield from write_file( io=io, table_metadata=table_metadata, tasks=( - WriteTask(write_uuid=write_uuid, task_id=next(counter), record_batches=batches, schema=task_schema) + WriteTask( + write_uuid=write_uuid, + task_id=next(counter), + record_batches=batches, + schema=task_schema, + sort_order_id=sort_order_id, + ) for batches in bin_pack_arrow_table(df, target_file_size) ), ) @@ -3037,13 +3048,49 @@ def _dataframe_to_data_files( record_batches=batches, partition_key=partition.partition_key, schema=task_schema, + sort_order_id=sort_order_id, ) for partition in partitions - for batches in bin_pack_arrow_table(partition.arrow_table_partition, target_file_size) + for sorted_partition, sort_order_id in [_sort_table_for_write(table_metadata, partition.arrow_table_partition)] + for batches in bin_pack_arrow_table(sorted_partition, target_file_size) ), ) +def _sort_table_for_write(table_metadata: TableMetadata, table: pa.Table) -> tuple[pa.Table, int | None]: + """Apply a supported Iceberg sort order and return its file sort-order id.""" + sort_order = table_metadata.sort_order() + if sort_order.is_unsorted: + return table, None + + sort_keys: list[tuple[str, str]] = [] + null_orders = set() + for field in sort_order.fields: + if not isinstance(field.transform, IdentityTransform): + warnings.warn( + f"Unsupported sort transform {field.transform}; data files are marked unsorted", + stacklevel=2, + ) + return table, None + name = table_metadata.schema().find_column_name(field.source_id) + if name is None or name not in table.column_names: + warnings.warn( + f"Unsupported nested or missing sort field id {field.source_id}; data files are marked unsorted", + stacklevel=2, + ) + return table, None + direction = "descending" if field.direction == SortDirection.DESC else "ascending" + sort_keys.append((name, direction)) + null_orders.add(field.null_order) + + if len(null_orders) != 1: + warnings.warn("Mixed null ordering is not supported; data files are marked unsorted", stacklevel=2) + return table, None + null_placement = "at_start" if null_orders == {NullOrder.NULLS_FIRST} else "at_end" + indices = pc.sort_indices(table, sort_keys=sort_keys, null_placement=null_placement) + return table.take(indices), sort_order.order_id + + @dataclass(frozen=True) class _TablePartition: partition_key: PartitionKey diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py index 30fdd76ab7..80de784b6f 100644 --- a/tests/integration/test_writes/test_writes.py +++ b/tests/integration/test_writes/test_writes.py @@ -52,7 +52,7 @@ from pyiceberg.schema import Schema from pyiceberg.table import TableProperties from pyiceberg.table.refs import MAIN_BRANCH -from pyiceberg.table.sorting import SortDirection, SortField, SortOrder +from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder from pyiceberg.transforms import DayTransform, HourTransform, IdentityTransform, Transform from pyiceberg.types import ( DateType, @@ -1069,6 +1069,27 @@ def test_create_table_with_non_default_values(catalog: Catalog, table_schema_wit assert tbl.sort_orders() == tbl_ref.sort_orders() +@pytest.mark.integration +def test_write_identity_sort_order(session_catalog: Catalog) -> None: + identifier = "default.write_identity_sort_order" + schema = Schema( + NestedField(1, "id", LongType(), required=False), + NestedField(2, "value", StringType(), required=False), + ) + table = session_catalog.create_table( + identifier, + schema, + sort_order=SortOrder( + SortField(1, IdentityTransform(), SortDirection.ASC, NullOrder.NULLS_LAST), + ), + ) + + table.append(pa.table({"id": [2, None, 1], "value": ["b", "null", "a"]})) + + assert table.scan().to_arrow()["id"].to_pylist() == [1, 2, None] + assert table.inspect.data_files()["sort_order_id"].to_pylist() == [table.sort_order().order_id] + + @pytest.mark.integration @pytest.mark.parametrize("format_version", [1, 2]) def test_table_properties_int_value( diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..1ada85465d 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -76,6 +76,7 @@ _determine_partitions, _primitive_to_physical, _read_deletes, + _sort_table_for_write, _task_to_record_batches, _to_requested_schema, bin_pack_arrow_table, @@ -91,8 +92,9 @@ from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema, make_compatible_name, visit from pyiceberg.table import FileScanTask, TableProperties, WriteTask -from pyiceberg.table.metadata import TableMetadataV2 +from pyiceberg.table.metadata import TableMetadataV2, new_table_metadata from pyiceberg.table.name_mapping import create_mapping_from_schema +from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder from pyiceberg.transforms import HourTransform, IdentityTransform from pyiceberg.typedef import UTF8, Properties, Record, TableVersion from pyiceberg.types import ( @@ -127,6 +129,28 @@ ) +def test_sort_table_for_identity_sort_order() -> None: + schema = Schema( + NestedField(1, "id", LongType(), required=False), + NestedField(2, "value", StringType(), required=False), + ) + metadata = new_table_metadata( + schema=schema, + partition_spec=PartitionSpec(), + sort_order=SortOrder( + SortField(1, IdentityTransform(), SortDirection.ASC, NullOrder.NULLS_LAST), + ), + location="file:///tmp/sorted", + properties={}, + ) + table = pa.table({"id": [2, None, 1], "value": ["b", "null", "a"]}) + + sorted_table, sort_order_id = _sort_table_for_write(metadata, table) + + assert sorted_table["id"].to_pylist() == [1, 2, None] + assert sort_order_id == metadata.default_sort_order_id + + def test_pyarrow_infer_local_fs_from_path() -> None: """Test path with `file` scheme and no scheme both use LocalFileSystem""" assert isinstance(PyArrowFileIO().new_output("file://tmp/warehouse")._filesystem, LocalFileSystem)