iceberg/transaction/
sort_order.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use async_trait::async_trait;
19
20use crate::error::{Result, invalid_data};
21use crate::spec::{NullOrder, SchemaRef, SortDirection, SortField, SortOrder, Transform};
22use crate::table::Table;
23use crate::transaction::{ActionCommit, StatelessAction};
24use crate::{TableRequirement, TableUpdate};
25
26/// Represents a sort field whose construction and validation are deferred until commit time.
27/// This avoids the need to pass a `Table` reference into methods like `asc` or `desc` when
28/// adding sort orders.
29#[derive(Debug, PartialEq, Eq, Clone)]
30struct PendingSortField {
31    name: String,
32    transform: Transform,
33    direction: SortDirection,
34    null_order: NullOrder,
35}
36
37impl PendingSortField {
38    fn to_sort_field(&self, schema: &SchemaRef) -> Result<SortField> {
39        let field_id = schema
40            .field_id_by_name(self.name.as_str())
41            .ok_or_else(|| invalid_data!("Cannot find field {} in table schema", self.name))?;
42
43        // This action intentionally refuses to author all-null (void) sort keys;
44        // this policy does not restrict reading or preserving existing metadata.
45        if matches!(self.transform, Transform::Unknown | Transform::Void) {
46            return Err(invalid_data!("Cannot sort by transform {}", self.transform));
47        }
48
49        Ok(SortField::builder()
50            .source_id(field_id)
51            .transform(self.transform)
52            .direction(self.direction)
53            .null_order(self.null_order)
54            .build())
55    }
56}
57
58/// Transaction action for replacing sort order.
59#[derive(Clone)]
60pub struct ReplaceSortOrderAction {
61    pending_sort_fields: Vec<PendingSortField>,
62}
63
64impl ReplaceSortOrderAction {
65    pub(crate) fn new() -> Self {
66        ReplaceSortOrderAction {
67            pending_sort_fields: vec![],
68        }
69    }
70
71    /// Adds a field for sorting in ascending order, sorting by the column's raw value
72    /// (an identity transform). To sort by a transform of the column instead (e.g.
73    /// `bucket[N]`, `year`, `truncate[W]`), use [`Self::asc_with_transform`].
74    pub fn asc(self, name: &str, null_order: NullOrder) -> Self {
75        self.asc_with_transform(name, Transform::Identity, null_order)
76    }
77
78    /// Adds a field for sorting in descending order, sorting by the column's raw value
79    /// (an identity transform). To sort by a transform of the column instead (e.g.
80    /// `bucket[N]`, `year`, `truncate[W]`), use [`Self::desc_with_transform`].
81    pub fn desc(self, name: &str, null_order: NullOrder) -> Self {
82        self.desc_with_transform(name, Transform::Identity, null_order)
83    }
84
85    /// Adds a field for sorting in ascending order by a transform of the column's value
86    /// (e.g. `Transform::Bucket(16)`, `Transform::Year`, `Transform::Truncate(4)`).
87    ///
88    /// Whether the transform is valid for the column's type is checked at commit time,
89    /// once the table schema is available (mirroring Java's `SortOrder.Builder.build()`).
90    /// `Transform::Unknown` and `Transform::Void` are rejected at commit time with
91    /// [`ErrorKind::DataInvalid`](crate::ErrorKind::DataInvalid). Rejecting `Void` is an
92    /// intentional authoring policy for this action: an all-null sort key adds no
93    /// ordering information.
94    ///
95    /// Note: `Term` is currently a plain column reference. Once it becomes
96    /// transform-carrying (#2665), sort-order declaration is expected to converge on
97    /// Term-based `asc`/`desc` (as in Java's `SortOrderBuilder`), at which point the
98    /// `_with_transform` variants can be deprecated in its favor.
99    pub fn asc_with_transform(
100        self,
101        name: &str,
102        transform: Transform,
103        null_order: NullOrder,
104    ) -> Self {
105        self.add_sort_field(name, transform, SortDirection::Ascending, null_order)
106    }
107
108    /// Adds a field for sorting in descending order by a transform of the column's value.
109    ///
110    /// See [`Self::asc_with_transform`] for validation and future API direction.
111    pub fn desc_with_transform(
112        self,
113        name: &str,
114        transform: Transform,
115        null_order: NullOrder,
116    ) -> Self {
117        self.add_sort_field(name, transform, SortDirection::Descending, null_order)
118    }
119
120    fn add_sort_field(
121        mut self,
122        name: &str,
123        transform: Transform,
124        sort_direction: SortDirection,
125        null_order: NullOrder,
126    ) -> Self {
127        self.pending_sort_fields.push(PendingSortField {
128            name: name.to_string(),
129            transform,
130            direction: sort_direction,
131            null_order,
132        });
133
134        self
135    }
136}
137
138#[async_trait]
139impl StatelessAction for ReplaceSortOrderAction {
140    async fn commit(&self, table: &Table) -> Result<ActionCommit> {
141        let current_schema = table.metadata().current_schema();
142        let sort_fields: Result<Vec<SortField>> = self
143            .pending_sort_fields
144            .iter()
145            .map(|p| p.to_sort_field(current_schema))
146            .collect();
147
148        let bound_sort_order = SortOrder::builder()
149            .with_fields(sort_fields?)
150            .build(current_schema)?;
151
152        let updates = vec![
153            TableUpdate::AddSortOrder {
154                sort_order: bound_sort_order,
155            },
156            TableUpdate::SetDefaultSortOrder { sort_order_id: -1 },
157        ];
158
159        let requirements = vec![
160            TableRequirement::CurrentSchemaIdMatch {
161                current_schema_id: current_schema.schema_id(),
162            },
163            TableRequirement::DefaultSortOrderIdMatch {
164                default_sort_order_id: table.metadata().default_sort_order().order_id,
165            },
166        ];
167
168        Ok(ActionCommit::new(updates, requirements))
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use as_any::Downcast;
175
176    use crate::catalog::Catalog;
177    use crate::memory::tests::new_memory_catalog;
178    use crate::spec::{NullOrder, SortDirection, Transform};
179    use crate::test_utils::commit_with_fresh_state;
180    use crate::transaction::action::ActionEntry;
181    use crate::transaction::sort_order::{PendingSortField, ReplaceSortOrderAction};
182    use crate::transaction::tests::{make_v2_table, make_v3_minimal_table_in_catalog};
183    use crate::transaction::{ApplyTransactionAction, Transaction};
184    use crate::{ErrorKind, TableUpdate};
185
186    #[test]
187    fn test_replace_sort_order() {
188        let table = make_v2_table();
189        let tx = Transaction::new(&table);
190        let replace_sort_order = tx.replace_sort_order();
191
192        let tx = replace_sort_order
193            .asc("x", NullOrder::First)
194            .desc("y", NullOrder::Last)
195            .apply(tx)
196            .unwrap();
197
198        let replace_sort_order = (*tx.actions[0])
199            .downcast_ref::<ActionEntry<ReplaceSortOrderAction>>()
200            .unwrap()
201            .action();
202
203        assert_eq!(replace_sort_order.pending_sort_fields, vec![
204            PendingSortField {
205                name: String::from("x"),
206                transform: Transform::Identity,
207                direction: SortDirection::Ascending,
208                null_order: NullOrder::First,
209            },
210            PendingSortField {
211                name: String::from("y"),
212                transform: Transform::Identity,
213                direction: SortDirection::Descending,
214                null_order: NullOrder::Last,
215            }
216        ]);
217    }
218
219    #[test]
220    fn test_replace_sort_order_with_transform() {
221        let table = make_v2_table();
222        let tx = Transaction::new(&table);
223        let replace_sort_order = tx.replace_sort_order();
224
225        let tx = replace_sort_order
226            .asc_with_transform("x", Transform::Bucket(16), NullOrder::First)
227            .desc_with_transform("y", Transform::Truncate(4), NullOrder::Last)
228            .apply(tx)
229            .unwrap();
230
231        let replace_sort_order = (*tx.actions[0])
232            .downcast_ref::<ActionEntry<ReplaceSortOrderAction>>()
233            .unwrap()
234            .action();
235
236        assert_eq!(replace_sort_order.pending_sort_fields, vec![
237            PendingSortField {
238                name: String::from("x"),
239                transform: Transform::Bucket(16),
240                direction: SortDirection::Ascending,
241                null_order: NullOrder::First,
242            },
243            PendingSortField {
244                name: String::from("y"),
245                transform: Transform::Truncate(4),
246                direction: SortDirection::Descending,
247                null_order: NullOrder::Last,
248            }
249        ]);
250    }
251
252    #[tokio::test]
253    async fn test_replace_sort_order_with_transform_commits() {
254        let table = make_v2_table();
255        let action = ReplaceSortOrderAction::new().asc_with_transform(
256            "x",
257            Transform::Bucket(16),
258            NullOrder::First,
259        );
260
261        let mut action_commit = commit_with_fresh_state(action, &table).await.unwrap();
262        let updates = action_commit.take_updates();
263
264        assert_eq!(updates.len(), 2);
265        let TableUpdate::AddSortOrder { sort_order } = &updates[0] else {
266            panic!("expected AddSortOrder, got {:?}", updates[0]);
267        };
268        assert_eq!(sort_order.fields.len(), 1);
269        assert_eq!(sort_order.fields[0].transform, Transform::Bucket(16));
270    }
271
272    #[tokio::test]
273    async fn test_replace_sort_order_rejects_incompatible_transform() {
274        let table = make_v2_table();
275        // `x` is a `long` column; `year` only accepts date/timestamp types.
276        let action = ReplaceSortOrderAction::new().asc_with_transform(
277            "x",
278            Transform::Year,
279            NullOrder::First,
280        );
281
282        let err = commit_with_fresh_state(action, &table)
283            .await
284            .err()
285            .expect("year transform on a long column should be rejected");
286        assert_eq!(err.kind(), ErrorKind::DataInvalid);
287    }
288
289    #[tokio::test]
290    async fn test_replace_sort_order_rejects_unknown_and_void_transforms() {
291        let table = make_v2_table();
292        for transform in [Transform::Unknown, Transform::Void] {
293            for action in [
294                ReplaceSortOrderAction::new().asc_with_transform("x", transform, NullOrder::First),
295                ReplaceSortOrderAction::new().desc_with_transform("x", transform, NullOrder::Last),
296            ] {
297                let err = commit_with_fresh_state(action, &table)
298                    .await
299                    .err()
300                    .expect("unknown and void sort transforms should be rejected");
301                assert_eq!(err.kind(), ErrorKind::DataInvalid);
302                assert_eq!(
303                    err.message(),
304                    format!("Cannot sort by transform {transform}")
305                );
306            }
307        }
308    }
309
310    #[tokio::test]
311    async fn test_sort_order_transform_survives_metadata_json_round_trip() {
312        // Commit a transform-based sort order through a real catalog: the memory
313        // catalog serializes the updated table metadata to a metadata.json file
314        // (`TableMetadata::write_to`), and `load_table` reads that file back and
315        // parses it (`TableMetadata::read_from`). This exercises the full
316        // JSON round-trip of the transform (e.g. `"bucket[16]"`), not just the
317        // in-memory `TableUpdate`.
318        let catalog = new_memory_catalog().await;
319        let table = make_v3_minimal_table_in_catalog(&catalog).await;
320
321        let tx = Transaction::new(&table);
322        let tx = tx
323            .replace_sort_order()
324            .asc_with_transform("x", Transform::Bucket(16), NullOrder::First)
325            .desc_with_transform("y", Transform::Truncate(4), NullOrder::Last)
326            .apply(tx)
327            .unwrap();
328        let committed = tx.commit(&catalog).await.unwrap();
329
330        // Reload from the catalog: this parses the metadata.json written above.
331        let reloaded = catalog.load_table(committed.identifier()).await.unwrap();
332        let sort_order = reloaded.metadata().default_sort_order();
333
334        assert_eq!(sort_order.fields.len(), 2);
335        assert_eq!(sort_order.fields[0].transform, Transform::Bucket(16));
336        assert_eq!(sort_order.fields[0].direction, SortDirection::Ascending);
337        assert_eq!(sort_order.fields[0].null_order, NullOrder::First);
338        assert_eq!(sort_order.fields[1].transform, Transform::Truncate(4));
339        assert_eq!(sort_order.fields[1].direction, SortDirection::Descending);
340        assert_eq!(sort_order.fields[1].null_order, NullOrder::Last);
341    }
342}
iceberg/transaction/sort_order.rs

iceberg/transaction/
sort_order.rs
