Honor identity sort orders for Arrow table writes

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

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(

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)