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;
19use std::sync::Arc;
20
21use async_trait::async_trait;
22
23use crate::spec::StatisticsFile;
24use crate::table::Table;
25use crate::transaction::{ActionCommit, TransactionAction};
26use crate::{Result, TableUpdate};
27
28/// A transactional action for updating statistics files in a table
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 TransactionAction for UpdateStatisticsAction {
73    async fn commit(self: Arc<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::tests::make_v2_table;
102    use crate::transaction::update_statistics::UpdateStatisticsAction;
103    use crate::transaction::{ApplyTransactionAction, Transaction};
104
105    #[test]
106    fn test_update_statistics() {
107        let table = make_v2_table();
108        let tx = Transaction::new(&table);
109
110        let statistics_file_1 = StatisticsFile {
111            snapshot_id: 3055729675574597004i64,
112            statistics_path: "s3://a/b/stats.puffin".to_string(),
113            file_size_in_bytes: 413,
114            file_footer_size_in_bytes: 42,
115            key_metadata: None,
116            blob_metadata: vec![BlobMetadata {
117                r#type: "ndv".to_string(),
118                snapshot_id: 3055729675574597004i64,
119                sequence_number: 1,
120                fields: vec![1],
121                properties: HashMap::new(),
122            }],
123        };
124
125        let statistics_file_2 = StatisticsFile {
126            snapshot_id: 3366729675595277004i64,
127            statistics_path: "s3://a/b/stats.puffin".to_string(),
128            file_size_in_bytes: 413,
129            file_footer_size_in_bytes: 42,
130            key_metadata: None,
131            blob_metadata: vec![BlobMetadata {
132                r#type: "ndv".to_string(),
133                snapshot_id: 3366729675595277004i64,
134                sequence_number: 1,
135                fields: vec![1],
136                properties: HashMap::new(),
137            }],
138        };
139
140        // set stats1
141        let tx = tx
142            .update_statistics()
143            .set_statistics(statistics_file_1.clone())
144            .set_statistics(statistics_file_2.clone())
145            .remove_statistics(3055729675574597004i64) // remove stats1
146            .apply(tx)
147            .unwrap();
148
149        let action = (*tx.actions[0])
150            .downcast_ref::<UpdateStatisticsAction>()
151            .unwrap();
152        assert!(
153            action
154                .statistics_to_set
155                .get(&statistics_file_1.snapshot_id)
156                .unwrap()
157                .is_none()
158        ); // stats1 should have been removed
159        assert_eq!(
160            action
161                .statistics_to_set
162                .get(&statistics_file_2.snapshot_id)
163                .unwrap()
164                .clone(),
165            Some(statistics_file_2)
166        );
167    }
168
169    #[test]
170    fn test_set_single_statistics() {
171        let table = make_v2_table();
172        let tx = Transaction::new(&table);
173
174        let statistics_file = StatisticsFile {
175            snapshot_id: 1234567890i64,
176            statistics_path: "s3://a/b/stats1.puffin".to_string(),
177            file_size_in_bytes: 500,
178            file_footer_size_in_bytes: 50,
179            key_metadata: None,
180            blob_metadata: vec![],
181        };
182
183        // Set statistics
184        let tx = tx
185            .update_statistics()
186            .set_statistics(statistics_file.clone())
187            .apply(tx)
188            .unwrap();
189
190        let action = (*tx.actions[0])
191            .downcast_ref::<UpdateStatisticsAction>()
192            .unwrap();
193
194        // Verify that the statistics file is set correctly
195        assert_eq!(
196            action
197                .statistics_to_set
198                .get(&statistics_file.snapshot_id)
199                .unwrap()
200                .clone(),
201            Some(statistics_file)
202        );
203    }
204
205    #[test]
206    fn test_no_statistics_set() {
207        let table = make_v2_table();
208        let tx = Transaction::new(&table);
209
210        // No statistics are set or removed
211        let tx = tx.update_statistics().apply(tx).unwrap();
212
213        let action = (*tx.actions[0])
214            .downcast_ref::<UpdateStatisticsAction>()
215            .unwrap();
216
217        // Verify that no statistics are set
218        assert!(action.statistics_to_set.is_empty());
219    }
220}
iceberg/transaction/update_statistics.rs

iceberg/transaction/
update_statistics.rs
