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

iceberg/transaction/
sort_order.rs
