iceberg/transaction/
update_statistics.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::collections::HashMap;
19
20use async_trait::async_trait;
21
22use crate::spec::StatisticsFile;
23use crate::table::Table;
24use crate::transaction::{ActionCommit, StatelessAction};
25use crate::{Result, TableUpdate};
26
27/// A transactional action for updating statistics files in a table
28#[derive(Clone)]
29pub struct UpdateStatisticsAction {
30    statistics_to_set: HashMap<i64, Option<StatisticsFile>>,
31}
32
33impl UpdateStatisticsAction {
34    pub(crate) fn new() -> Self {
35        Self {
36            statistics_to_set: HashMap::default(),
37        }
38    }
39
40    /// Set the table's statistics file for given snapshot, replacing the previous statistics file for
41    /// the snapshot if any exists. The snapshot id of the statistics file will be used.
42    ///
43    /// # Arguments
44    ///
45    /// * `statistics_file` - The [`StatisticsFile`] to associate with its corresponding snapshot ID.
46    ///
47    /// # Returns
48    ///
49    /// An updated [`UpdateStatisticsAction`] with the new statistics file applied.
50    pub fn set_statistics(mut self, statistics_file: StatisticsFile) -> Self {
51        self.statistics_to_set
52            .insert(statistics_file.snapshot_id, Some(statistics_file));
53        self
54    }
55
56    /// Remove the table's statistics file for given snapshot.
57    ///
58    /// # Arguments
59    ///
60    /// * `snapshot_id` - The ID of the snapshot whose statistics file should be removed.
61    ///
62    /// # Returns
63    ///
64    /// An updated [`UpdateStatisticsAction`] with the removal operation recorded.
65    pub fn remove_statistics(mut self, snapshot_id: i64) -> Self {
66        self.statistics_to_set.insert(snapshot_id, None);
67        self
68    }
69}
70
71#[async_trait]
72impl StatelessAction for UpdateStatisticsAction {
73    async fn commit(&self, _table: &Table) -> Result<ActionCommit> {
74        let mut updates: Vec<TableUpdate> = vec![];
75
76        self.statistics_to_set
77            .iter()
78            .for_each(|(snapshot_id, statistic_file)| {
79                if let Some(statistics) = statistic_file {
80                    updates.push(TableUpdate::SetStatistics {
81                        statistics: statistics.clone(),
82                    })
83                } else {
84                    updates.push(TableUpdate::RemoveStatistics {
85                        snapshot_id: *snapshot_id,
86                    })
87                }
88            });
89
90        Ok(ActionCommit::new(updates, vec![]))
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::collections::HashMap;
97
98    use as_any::Downcast;
99
100    use crate::spec::{BlobMetadata, StatisticsFile};
101    use crate::transaction::action::ActionEntry;
102    use crate::transaction::tests::make_v2_table;
103    use crate::transaction::update_statistics::UpdateStatisticsAction;
104    use crate::transaction::{ApplyTransactionAction, Transaction};
105
106    #[test]
107    fn test_update_statistics() {
108        let table = make_v2_table();
109        let tx = Transaction::new(&table);
110
111        let statistics_file_1 = StatisticsFile {
112            snapshot_id: 3055729675574597004i64,
113            statistics_path: "s3://a/b/stats.puffin".to_string(),
114            file_size_in_bytes: 413,
115            file_footer_size_in_bytes: 42,
116            key_metadata: None,
117            blob_metadata: vec![BlobMetadata {
118                r#type: "ndv".to_string(),
119                snapshot_id: 3055729675574597004i64,
120                sequence_number: 1,
121                fields: vec![1],
122                properties: HashMap::new(),
123            }],
124        };
125
126        let statistics_file_2 = StatisticsFile {
127            snapshot_id: 3366729675595277004i64,
128            statistics_path: "s3://a/b/stats.puffin".to_string(),
129            file_size_in_bytes: 413,
130            file_footer_size_in_bytes: 42,
131            key_metadata: None,
132            blob_metadata: vec![BlobMetadata {
133                r#type: "ndv".to_string(),
134                snapshot_id: 3366729675595277004i64,
135                sequence_number: 1,
136                fields: vec![1],
137                properties: HashMap::new(),
138            }],
139        };
140
141        // set stats1
142        let tx = tx
143            .update_statistics()
144            .set_statistics(statistics_file_1.clone())
145            .set_statistics(statistics_file_2.clone())
146            .remove_statistics(3055729675574597004i64) // remove stats1
147            .apply(tx)
148            .unwrap();
149
150        let action = (*tx.actions[0])
151            .downcast_ref::<ActionEntry<UpdateStatisticsAction>>()
152            .unwrap()
153            .action();
154        assert!(
155            action
156                .statistics_to_set
157                .get(&statistics_file_1.snapshot_id)
158                .unwrap()
159                .is_none()
160        ); // stats1 should have been removed
161        assert_eq!(
162            action
163                .statistics_to_set
164                .get(&statistics_file_2.snapshot_id)
165                .unwrap()
166                .clone(),
167            Some(statistics_file_2)
168        );
169    }
170
171    #[test]
172    fn test_set_single_statistics() {
173        let table = make_v2_table();
174        let tx = Transaction::new(&table);
175
176        let statistics_file = StatisticsFile {
177            snapshot_id: 1234567890i64,
178            statistics_path: "s3://a/b/stats1.puffin".to_string(),
179            file_size_in_bytes: 500,
180            file_footer_size_in_bytes: 50,
181            key_metadata: None,
182            blob_metadata: vec![],
183        };
184
185        // Set statistics
186        let tx = tx
187            .update_statistics()
188            .set_statistics(statistics_file.clone())
189            .apply(tx)
190            .unwrap();
191
192        let action = (*tx.actions[0])
193            .downcast_ref::<ActionEntry<UpdateStatisticsAction>>()
194            .unwrap()
195            .action();
196
197        // Verify that the statistics file is set correctly
198        assert_eq!(
199            action
200                .statistics_to_set
201                .get(&statistics_file.snapshot_id)
202                .unwrap()
203                .clone(),
204            Some(statistics_file)
205        );
206    }
207
208    #[test]
209    fn test_no_statistics_set() {
210        let table = make_v2_table();
211        let tx = Transaction::new(&table);
212
213        // No statistics are set or removed
214        let tx = tx.update_statistics().apply(tx).unwrap();
215
216        let action = (*tx.actions[0])
217            .downcast_ref::<ActionEntry<UpdateStatisticsAction>>()
218            .unwrap()
219            .action();
220
221        // Verify that no statistics are set
222        assert!(action.statistics_to_set.is_empty());
223    }
224}
iceberg/transaction/update_statistics.rs

iceberg/transaction/
update_statistics.rs
