feat(transaction): support transform-based sort orders

crates/iceberg/public-api.txt

@@ -3249,7 +3249,9 @@ pub fn iceberg::transaction::FastAppendAction::with_check_duplicate(self, v: boo
 pub struct iceberg::transaction::ReplaceSortOrderAction
 impl iceberg::transaction::ReplaceSortOrderAction
 pub fn iceberg::transaction::ReplaceSortOrderAction::asc(self, name: &str, null_order: iceberg::spec::NullOrder) -> Self
+pub fn iceberg::transaction::ReplaceSortOrderAction::asc_with_transform(self, name: &str, transform: iceberg::spec::Transform, null_order: iceberg::spec::NullOrder) -> Self
 pub fn iceberg::transaction::ReplaceSortOrderAction::desc(self, name: &str, null_order: iceberg::spec::NullOrder) -> Self
+pub fn iceberg::transaction::ReplaceSortOrderAction::desc_with_transform(self, name: &str, transform: iceberg::spec::Transform, null_order: iceberg::spec::NullOrder) -> Self
 pub struct iceberg::transaction::Transaction
 impl iceberg::transaction::Transaction
 pub async fn iceberg::transaction::Transaction::commit(self, catalog: &dyn iceberg::Catalog) -> iceberg::Result<iceberg::table::Table>

crates/iceberg/src/spec/sort.rs

@@ -180,7 +180,7 @@ impl SortOrderBuilder {
             match schema.field_by_id(sort_field.source_id) {
                 None => {
                     return Err(Error::new(
-                        ErrorKind::Unexpected,
+                        ErrorKind::DataInvalid,
                         format!("Cannot find source column for sort field: {sort_field}"),
                     ));
                 }
@@ -189,20 +189,13 @@ impl SortOrderBuilder {
 
                     if !source_type.is_primitive() {
                         return Err(Error::new(
-                            ErrorKind::Unexpected,
+                            ErrorKind::DataInvalid,
                             format!("Cannot sort by non-primitive source field: {source_type}"),
                         ));
                     }
 
                     let field_transform = sort_field.transform;
-                    if field_transform.result_type(source_type).is_err() {
-                        return Err(Error::new(
-                            ErrorKind::Unexpected,
-                            format!(
-                                "Invalid source type {source_type} for transform {field_transform}"
-                            ),
-                        ));
-                    }
+                    field_transform.result_type(source_type)?;
                 }
             }
         }
@@ -403,12 +396,12 @@ mod tests {
             )
             .build(&schema);
 
+        let err = sort_order_builder_result.expect_err("Expected an Err value");
+        assert_eq!(err.kind(), ErrorKind::DataInvalid);
         assert_eq!(
-            sort_order_builder_result
-                .expect_err("Expected an Err value")
-                .message(),
+            err.message(),
             "Cannot find source column for sort field: SortField { source_id: 2, transform: identity, direction: ascending, null_order: first }"
-        )
+        );
     }
 
     #[test]
@@ -444,12 +437,12 @@ mod tests {
             )
             .build(&schema);
 
+        let err = sort_order_builder_result.expect_err("Expected an Err value");
+        assert_eq!(err.kind(), ErrorKind::DataInvalid);
         assert_eq!(
-            sort_order_builder_result
-                .expect_err("Expected an Err value")
-                .message(),
+            err.message(),
             "Cannot sort by non-primitive source field: list"
-        )
+        );
     }
 
     #[test]
@@ -473,12 +466,12 @@ mod tests {
             )
             .build(&schema);
 
+        let err = sort_order_builder_result.expect_err("Expected an Err value");
+        assert_eq!(err.kind(), ErrorKind::DataInvalid);
         assert_eq!(
-            sort_order_builder_result
-                .expect_err("Expected an Err value")
-                .message(),
+            err.message(),
             "Cannot sort by non-primitive source field: variant"
-        )
+        );
     }
 
     #[test]
@@ -502,12 +495,12 @@ mod tests {
             )
             .build(&schema);
 
+        let err = sort_order_builder_result.expect_err("Expected an Err value");
+        assert_eq!(err.kind(), ErrorKind::DataInvalid);
         assert_eq!(
-            sort_order_builder_result
-                .expect_err("Expected an Err value")
-                .message(),
-            "Invalid source type int for transform year"
-        )
+            err.message(),
+            "int is not a valid input type of year transform"
+        );
     }
 
     #[test]

crates/iceberg/src/transaction/sort_order.rs

@@ -31,6 +31,7 @@ use crate::{Error, ErrorKind, TableRequirement, TableUpdate};
 #[derive(Debug, PartialEq, Eq, Clone)]
 struct PendingSortField {
     name: String,
+    transform: Transform,
     direction: SortDirection,
     null_order: NullOrder,
 }
@@ -44,9 +45,18 @@ impl PendingSortField {
             )
         })?;
 
+        // This action intentionally refuses to author all-null (void) sort keys;
+        // this policy does not restrict reading or preserving existing metadata.
+        if matches!(self.transform, Transform::Unknown | Transform::Void) {
+            return Err(Error::new(
+                ErrorKind::DataInvalid,
+                format!("Cannot sort by transform {}", self.transform),
+            ));
+        }
+
         Ok(SortField::builder()
             .source_id(field_id)
-            .transform(Transform::Identity)
+            .transform(self.transform)
             .direction(self.direction)
             .null_order(self.null_order)
             .build())
@@ -65,24 +75,64 @@ impl ReplaceSortOrderAction {
         }
     }
 
-    /// Adds a field for sorting in ascending order.
+    /// Adds a field for sorting in ascending order, sorting by the column's raw value
+    /// (an identity transform). To sort by a transform of the column instead (e.g.
+    /// `bucket[N]`, `year`, `truncate[W]`), use [`Self::asc_with_transform`].
     pub fn asc(self, name: &str, null_order: NullOrder) -> Self {
-        self.add_sort_field(name, SortDirection::Ascending, null_order)
+        self.asc_with_transform(name, Transform::Identity, null_order)
     }
 
-    /// Adds a field for sorting in descending order.
+    /// Adds a field for sorting in descending order, sorting by the column's raw value
+    /// (an identity transform). To sort by a transform of the column instead (e.g.
+    /// `bucket[N]`, `year`, `truncate[W]`), use [`Self::desc_with_transform`].
     pub fn desc(self, name: &str, null_order: NullOrder) -> Self {
-        self.add_sort_field(name, SortDirection::Descending, null_order)
+        self.desc_with_transform(name, Transform::Identity, null_order)
+    }
+
+    /// Adds a field for sorting in ascending order by a transform of the column's value
+    /// (e.g. `Transform::Bucket(16)`, `Transform::Year`, `Transform::Truncate(4)`).
+    ///
+    /// Whether the transform is valid for the column's type is checked at commit time,
+    /// once the table schema is available (mirroring Java's `SortOrder.Builder.build()`).
+    /// `Transform::Unknown` and `Transform::Void` are rejected at commit time with
+    /// [`ErrorKind::DataInvalid`]. Rejecting `Void` is an intentional authoring policy
+    /// for this action: an all-null sort key adds no ordering information.
+    ///
+    /// Note: `Term` is currently a plain column reference. Once it becomes
+    /// transform-carrying (#2665), sort-order declaration is expected to converge on
+    /// Term-based `asc`/`desc` (as in Java's `SortOrderBuilder`), at which point the
+    /// `_with_transform` variants can be deprecated in its favor.
+    pub fn asc_with_transform(
+        self,
+        name: &str,
+        transform: Transform,
+        null_order: NullOrder,
+    ) -> Self {
+        self.add_sort_field(name, transform, SortDirection::Ascending, null_order)
+    }
+
+    /// Adds a field for sorting in descending order by a transform of the column's value.
+    ///
+    /// See [`Self::asc_with_transform`] for validation and future API direction.
+    pub fn desc_with_transform(
+        self,
+        name: &str,
+        transform: Transform,
+        null_order: NullOrder,
+    ) -> Self {
+        self.add_sort_field(name, transform, SortDirection::Descending, null_order)
     }
 
     fn add_sort_field(
         mut self,
         name: &str,
+        transform: Transform,
         sort_direction: SortDirection,
         null_order: NullOrder,
     ) -> Self {
         self.pending_sort_fields.push(PendingSortField {
             name: name.to_string(),
+            transform,
             direction: sort_direction,
             null_order,
         });
@@ -127,12 +177,17 @@ impl TransactionAction for ReplaceSortOrderAction {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use as_any::Downcast;
 
-    use crate::spec::{NullOrder, SortDirection};
+    use crate::catalog::Catalog;
+    use crate::memory::tests::new_memory_catalog;
+    use crate::spec::{NullOrder, SortDirection, Transform};
     use crate::transaction::sort_order::{PendingSortField, ReplaceSortOrderAction};
-    use crate::transaction::tests::make_v2_table;
-    use crate::transaction::{ApplyTransactionAction, Transaction};
+    use crate::transaction::tests::{make_v2_table, make_v3_minimal_table_in_catalog};
+    use crate::transaction::{ApplyTransactionAction, Transaction, TransactionAction};
+    use crate::{ErrorKind, TableUpdate};
 
     #[test]
     fn test_replace_sort_order() {
@@ -153,14 +208,139 @@ mod tests {
         assert_eq!(replace_sort_order.pending_sort_fields, vec![
             PendingSortField {
                 name: String::from("x"),
+                transform: Transform::Identity,
                 direction: SortDirection::Ascending,
                 null_order: NullOrder::First,
             },
             PendingSortField {
                 name: String::from("y"),
+                transform: Transform::Identity,
                 direction: SortDirection::Descending,
                 null_order: NullOrder::Last,
             }
         ]);
     }
+
+    #[test]
+    fn test_replace_sort_order_with_transform() {
+        let table = make_v2_table();
+        let tx = Transaction::new(&table);
+        let replace_sort_order = tx.replace_sort_order();
+
+        let tx = replace_sort_order
+            .asc_with_transform("x", Transform::Bucket(16), NullOrder::First)
+            .desc_with_transform("y", Transform::Truncate(4), NullOrder::Last)
+            .apply(tx)
+            .unwrap();
+
+        let replace_sort_order = (*tx.actions[0])
+            .downcast_ref::<ReplaceSortOrderAction>()
+            .unwrap();
+
+        assert_eq!(replace_sort_order.pending_sort_fields, vec![
+            PendingSortField {
+                name: String::from("x"),
+                transform: Transform::Bucket(16),
+                direction: SortDirection::Ascending,
+                null_order: NullOrder::First,
+            },
+            PendingSortField {
+                name: String::from("y"),
+                transform: Transform::Truncate(4),
+                direction: SortDirection::Descending,
+                null_order: NullOrder::Last,
+            }
+        ]);
+    }
+
+    #[tokio::test]
+    async fn test_replace_sort_order_with_transform_commits() {
+        let table = make_v2_table();
+        let action = Arc::new(ReplaceSortOrderAction::new().asc_with_transform(
+            "x",
+            Transform::Bucket(16),
+            NullOrder::First,
+        ));
+
+        let mut action_commit = TransactionAction::commit(action, &table).await.unwrap();
+        let updates = action_commit.take_updates();
+
+        assert_eq!(updates.len(), 2);
+        let TableUpdate::AddSortOrder { sort_order } = &updates[0] else {
+            panic!("expected AddSortOrder, got {:?}", updates[0]);
+        };
+        assert_eq!(sort_order.fields.len(), 1);
+        assert_eq!(sort_order.fields[0].transform, Transform::Bucket(16));
+    }
+
+    #[tokio::test]
+    async fn test_replace_sort_order_rejects_incompatible_transform() {
+        let table = make_v2_table();
+        // `x` is a `long` column; `year` only accepts date/timestamp types.
+        let action = Arc::new(ReplaceSortOrderAction::new().asc_with_transform(
+            "x",
+            Transform::Year,
+            NullOrder::First,
+        ));
+
+        let err = TransactionAction::commit(action, &table)
+            .await
+            .err()
+            .expect("year transform on a long column should be rejected");
+        assert_eq!(err.kind(), ErrorKind::DataInvalid);
+    }
+
+    #[tokio::test]
+    async fn test_replace_sort_order_rejects_unknown_and_void_transforms() {
+        let table = make_v2_table();
+        for transform in [Transform::Unknown, Transform::Void] {
+            for action in [
+                ReplaceSortOrderAction::new().asc_with_transform("x", transform, NullOrder::First),
+                ReplaceSortOrderAction::new().desc_with_transform("x", transform, NullOrder::Last),
+            ] {
+                let err = TransactionAction::commit(Arc::new(action), &table)
+                    .await
+                    .err()
+                    .expect("unknown and void sort transforms should be rejected");
+                assert_eq!(err.kind(), ErrorKind::DataInvalid);
+                assert_eq!(
+                    err.message(),
+                    format!("Cannot sort by transform {transform}")
+                );
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_sort_order_transform_survives_metadata_json_round_trip() {
+        // Commit a transform-based sort order through a real catalog: the memory
+        // catalog serializes the updated table metadata to a metadata.json file
+        // (`TableMetadata::write_to`), and `load_table` reads that file back and
+        // parses it (`TableMetadata::read_from`). This exercises the full
+        // JSON round-trip of the transform (e.g. `"bucket[16]"`), not just the
+        // in-memory `TableUpdate`.
+        let catalog = new_memory_catalog().await;
+        let table = make_v3_minimal_table_in_catalog(&catalog).await;
+
+        let tx = Transaction::new(&table);
+        let tx = tx
+            .replace_sort_order()
+            .asc_with_transform("x", Transform::Bucket(16), NullOrder::First)
+            .desc_with_transform("y", Transform::Truncate(4), NullOrder::Last)
+            .apply(tx)
+            .unwrap();
+        let committed = tx.commit(&catalog).await.unwrap();
+
+        // Reload from the catalog: this parses the metadata.json written above.
+        let reloaded = catalog.load_table(committed.identifier()).await.unwrap();
+        let sort_order = reloaded.metadata().default_sort_order();
+
+        assert_eq!(sort_order.fields.len(), 2);
+        assert_eq!(sort_order.fields[0].transform, Transform::Bucket(16));
+        assert_eq!(sort_order.fields[0].direction, SortDirection::Ascending);
+        assert_eq!(sort_order.fields[0].null_order, NullOrder::First);
+        assert_eq!(sort_order.fields[1].transform, Transform::Truncate(4));
+        assert_eq!(sort_order.fields[1].direction, SortDirection::Descending);
+        assert_eq!(sort_order.fields[1].null_order, NullOrder::Last);
+    }
 }