[WIP] Feat: replace sort order

pyiceberg/table/__init__.py

@@ -117,6 +117,7 @@
     UpdateSnapshot,
     _FastAppendFiles,
 )
+from pyiceberg.table.update.sorting import UpdateSortOrder
 from pyiceberg.table.update.spec import UpdateSpec
 from pyiceberg.table.update.statistics import UpdateStatistics
 from pyiceberg.transforms import IdentityTransform
@@ -413,6 +414,20 @@ def update_schema(self, allow_incompatible_changes: bool = False, case_sensitive
             name_mapping=self.table_metadata.name_mapping(),
         )
 
+    def update_sort_order(self, case_sensitive: bool = True) -> UpdateSortOrder:
+        """Create a new UpdateSortOrder to update the sort order of this table.
+
+        Args:
+            case_sensitive: If field names are case-sensitive.
+
+        Returns:
+            A new UpdateSortOrder.
+        """
+        return UpdateSortOrder(
+            self,
+            case_sensitive=case_sensitive,
+        )
+
     def update_snapshot(self, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> UpdateSnapshot:
         """Create a new UpdateSnapshot to produce a new snapshot for the table.
 
@@ -1080,6 +1095,14 @@ def update_schema(self, allow_incompatible_changes: bool = False, case_sensitive
             name_mapping=self.name_mapping(),
         )
 
+    def update_sort_order(self) -> UpdateSortOrder:
+        """Create a new UpdateSortOrder to update the sort order of this table.
+
+        Returns:
+            A new UpdateSortOrder.
+        """
+        return UpdateSortOrder(transaction=Transaction(self, autocommit=True), case_sensitive=True)
+
     def name_mapping(self) -> Optional[NameMapping]:
         """Return the table's field-id NameMapping."""
         return self.metadata.name_mapping()

pyiceberg/table/update/sorting.py

@@ -0,0 +1,115 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Tuple
+
+from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder
+from pyiceberg.table.update import (
+    AddSortOrderUpdate,
+    AssertDefaultSortOrderId,
+    SetDefaultSortOrderUpdate,
+    TableRequirement,
+    TableUpdate,
+    UpdatesAndRequirements,
+    UpdateTableMetadata,
+)
+from pyiceberg.transforms import Transform
+
+if TYPE_CHECKING:
+    from pyiceberg.table import Transaction
+
+
+class UpdateSortOrder(UpdateTableMetadata["UpdateSortOrder"]):
+    _transaction: Transaction
+    _last_assigned_order_id: int
+    _case_sensitive: bool
+    _fields: List[SortField]
+    _last_sort_order_id: int
+
+    def __init__(self, transaction: Transaction, case_sensitive: bool = True) -> None:
+        super().__init__(transaction)
+        self._fields: List[SortField] = []
+        self._case_sensitive: bool = case_sensitive
+        self._last_sort_order_id: int = transaction.table_metadata.default_sort_order_id
+
+    def _column_name_to_id(self, column_name: str) -> int:
+        """Map the column name to the column field id."""
+        return (
+            self._transaction.table_metadata.schema()
+            .find_field(
+                name_or_id=column_name,
+                case_sensitive=self._case_sensitive,
+            )
+            .field_id
+        )
+
+    def _add_sort_field(
+        self,
+        source_id: int,
+        transform: Transform[Any, Any],
+        direction: SortDirection,
+        null_order: NullOrder,
+    ) -> UpdateSortOrder:
+        """Add a sort field to the sort order list."""
+        self._fields.append(
+            SortField(
+                source_id=source_id,
+                transform=transform,
+                direction=direction,
+                null_order=null_order,
+            )
+        )
+        return self
+
+    def asc(self, source_column_name: str, transform: Transform[Any, Any], null_order: NullOrder) -> UpdateSortOrder:
+        """Add a sort field with ascending order."""
+        return self._add_sort_field(
+            source_id=self._column_name_to_id(source_column_name),
+            transform=transform,
+            direction=SortDirection.ASC,
+            null_order=null_order,
+        )
+
+    def desc(self, source_column_name: str, transform: Transform[Any, Any], null_order: NullOrder) -> UpdateSortOrder:
+        """Add a sort field with descending order."""
+        return self._add_sort_field(
+            source_id=self._column_name_to_id(source_column_name),
+            transform=transform,
+            direction=SortDirection.DESC,
+            null_order=null_order,
+        )
+
+    def _apply(self) -> SortOrder:
+        """Return the sort order."""
+        return SortOrder(*self._fields, order_id=self._last_sort_order_id + 1)
+
+    def _commit(self) -> UpdatesAndRequirements:
+        """Apply the pending changes and commit."""
+        new_sort_order = self._apply()
+        requirements: Tuple[TableRequirement, ...] = ()
+        updates: Tuple[TableUpdate, ...] = ()
+
+        if self._transaction.table_metadata.default_sort_order_id != new_sort_order.order_id:
+            updates = (AddSortOrderUpdate(sort_order=new_sort_order), SetDefaultSortOrderUpdate(sort_order_id=-1))
+        else:
+            updates = (SetDefaultSortOrderUpdate(sort_order_id=new_sort_order.order_id),)
+
+        required_last_assigned_sort_order_id = self._transaction.table_metadata.default_sort_order_id
+        requirements = (AssertDefaultSortOrderId(default_sort_order_id=required_last_assigned_sort_order_id),)
+
+        return updates, requirements

tests/integration/test_sort_order_update.py

@@ -0,0 +1,109 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint:disable=redefined-outer-name
+
+import pytest
+
+from pyiceberg.catalog import Catalog
+from pyiceberg.exceptions import NoSuchTableError
+from pyiceberg.schema import Schema
+from pyiceberg.table import Table
+from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder
+from pyiceberg.transforms import (
+    IdentityTransform,
+)
+from pyiceberg.types import (
+    LongType,
+    NestedField,
+    StringType,
+    TimestampType,
+)
+
+
+def _simple_table(catalog: Catalog, table_schema_simple: Schema) -> Table:
+    return _create_table_with_schema(catalog, table_schema_simple, "1")
+
+
+def _table(catalog: Catalog) -> Table:
+    schema_with_timestamp = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "event_ts", TimestampType(), required=False),
+        NestedField(3, "str", StringType(), required=False),
+    )
+    return _create_table_with_schema(catalog, schema_with_timestamp, "1")
+
+
+def _table_v2(catalog: Catalog) -> Table:
+    schema_with_timestamp = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "event_ts", TimestampType(), required=False),
+        NestedField(3, "str", StringType(), required=False),
+    )
+    return _create_table_with_schema(catalog, schema_with_timestamp, "2")
+
+
+def _create_table_with_schema(catalog: Catalog, schema: Schema, format_version: str) -> Table:
+    tbl_name = "default.test_schema_evolution"
+    try:
+        catalog.drop_table(tbl_name)
+    except NoSuchTableError:
+        pass
+    return catalog.create_table(identifier=tbl_name, schema=schema, properties={"format-version": format_version})
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog"), pytest.lazy_fixture("session_catalog_hive")])
+def test_map_column_name_to_id(catalog: Catalog, table_schema_simple: Schema) -> None:
+    simple_table = _simple_table(catalog, table_schema_simple)
+    for col_name, col_id in {"foo": 1, "bar": 2, "baz": 3}.items():
+        assert col_id == simple_table.update_sort_order()._column_name_to_id(col_name)
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog"), pytest.lazy_fixture("session_catalog_hive")])
+def test_replace_sort_order(catalog: Catalog, table_schema_simple: Schema) -> None:
+    simple_table = _simple_table(catalog, table_schema_simple)
+    simple_table.update_sort_order().asc("foo", IdentityTransform(), NullOrder.NULLS_FIRST).desc(
+        "bar", IdentityTransform(), NullOrder.NULLS_LAST
+    ).commit()
+    assert simple_table.sort_order() == SortOrder(
+        SortField(source_id=1, transform=IdentityTransform(), direction=SortDirection.ASC, null_order=NullOrder.NULLS_FIRST),
+        SortField(source_id=2, transform=IdentityTransform(), direction=SortDirection.DESC, null_order=NullOrder.NULLS_LAST),
+        order_id=1,
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog"), pytest.lazy_fixture("session_catalog_hive")])
+def test_replace_existing_sort_order(catalog: Catalog, table_schema_simple: Schema) -> None:
+    simple_table = _simple_table(catalog, table_schema_simple)
+    simple_table.update_sort_order().asc("foo", IdentityTransform(), NullOrder.NULLS_FIRST).commit()
+    assert simple_table.sort_order() == SortOrder(
+        SortField(source_id=1, transform=IdentityTransform(), direction=SortDirection.ASC, null_order=NullOrder.NULLS_FIRST),
+        order_id=1,
+    )
+    simple_table.update_sort_order().asc("foo", IdentityTransform(), NullOrder.NULLS_LAST).desc(
+        "bar", IdentityTransform(), NullOrder.NULLS_FIRST
+    ).commit()
+    assert (
+        len(simple_table.sort_orders()) == 3
+    )  # 0: empty sort order from creating tables, 1: first sort order, 2: second sort order
+    assert simple_table.sort_order() == SortOrder(
+        SortField(source_id=1, transform=IdentityTransform(), direction=SortDirection.ASC, null_order=NullOrder.NULLS_LAST),
+        SortField(source_id=2, transform=IdentityTransform(), direction=SortDirection.DESC, null_order=NullOrder.NULLS_FIRST),
+        order_id=2,
+    )