Skip to main content

mito2/sst/index/indexer/
finish.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use common_telemetry::{debug, warn};
16use index::target::IndexTarget;
17use puffin::puffin_manager::{PuffinManager, PuffinWriter};
18use store_api::storage::ColumnId;
19
20use crate::sst::file::{RegionFileId, RegionIndexId};
21use crate::sst::index::puffin_manager::SstPuffinWriter;
22use crate::sst::index::statistics::{ByteCount, RowCount};
23use crate::sst::index::{
24    BloomFilterOutput, FulltextIndexOutput, IndexOutput, Indexer, InvertedIndexOutput,
25};
26
27impl Indexer {
28    pub(crate) async fn do_finish(&mut self) -> IndexOutput {
29        self.dense_pk_decoder = None;
30        let mut output = IndexOutput::default();
31
32        let Some(mut writer) = self.build_puffin_writer().await else {
33            self.do_abort().await;
34            return output;
35        };
36
37        let success = self
38            .do_finish_inverted_index(&mut writer, &mut output)
39            .await;
40        if !success {
41            self.do_abort().await;
42            return IndexOutput::default();
43        }
44
45        let success = self
46            .do_finish_fulltext_index(&mut writer, &mut output)
47            .await;
48        if !success {
49            self.do_abort().await;
50            return IndexOutput::default();
51        }
52
53        let success = self.do_finish_bloom_filter(&mut writer, &mut output).await;
54        if !success {
55            self.do_abort().await;
56            return IndexOutput::default();
57        }
58
59        self.do_prune_intm_sst_dir().await;
60        output.file_size = self.do_finish_puffin_writer(writer).await;
61        // A failed puffin file must not advertise the indexes written into it. Not covered by
62        // tests: `do_finish_puffin_writer` panics on failure in test builds.
63        if output.file_size == 0 {
64            return IndexOutput::default();
65        }
66        output.version = self.index_version;
67        output
68    }
69
70    async fn build_puffin_writer(&mut self) -> Option<SstPuffinWriter> {
71        let puffin_manager = self.puffin_manager.clone()?;
72
73        let err = match puffin_manager
74            .writer(&RegionIndexId::new(
75                RegionFileId::new(self.physical_region_id, self.file_id),
76                self.index_version,
77            ))
78            .await
79        {
80            Ok(writer) => return Some(writer),
81            Err(err) => err,
82        };
83
84        if cfg!(any(test, feature = "test")) {
85            panic!(
86                "Failed to create puffin writer, region_id: {}, file_id: {}, err: {:?}",
87                self.region_id, self.file_id, err
88            );
89        } else {
90            warn!(
91                err; "Failed to create puffin writer, region_id: {}, file_id: {}",
92                self.region_id, self.file_id,
93            );
94        }
95
96        None
97    }
98
99    async fn do_finish_puffin_writer(&mut self, writer: SstPuffinWriter) -> ByteCount {
100        let err = match writer.finish().await {
101            Ok(size) => return size,
102            Err(err) => err,
103        };
104
105        if cfg!(any(test, feature = "test")) {
106            panic!(
107                "Failed to finish puffin writer, region_id: {}, file_id: {}, err: {:?}",
108                self.region_id, self.file_id, err
109            );
110        } else {
111            warn!(
112                err; "Failed to finish puffin writer, region_id: {}, file_id: {}",
113                self.region_id, self.file_id,
114            );
115        }
116
117        0
118    }
119
120    /// Returns false if the finish failed.
121    async fn do_finish_inverted_index(
122        &mut self,
123        puffin_writer: &mut SstPuffinWriter,
124        index_output: &mut IndexOutput,
125    ) -> bool {
126        let Some(mut indexer) = self.inverted_indexer.take() else {
127            return true;
128        };
129
130        let column_ids = indexer
131            .indexed_targets()
132            .filter_map(|target| match target {
133                IndexTarget::ColumnId(id) => Some(*id),
134                IndexTarget::JsonPath(_) => None,
135            })
136            .collect();
137        let err = match indexer.finish(puffin_writer).await {
138            Ok((row_count, byte_count)) => {
139                self.fill_inverted_index_output(
140                    &mut index_output.inverted_index,
141                    row_count,
142                    byte_count,
143                    column_ids,
144                );
145                return true;
146            }
147            Err(err) => err,
148        };
149
150        if cfg!(any(test, feature = "test")) {
151            panic!(
152                "Failed to finish inverted index, region_id: {}, file_id: {}, err: {:?}",
153                self.region_id, self.file_id, err
154            );
155        } else {
156            warn!(
157                err; "Failed to finish inverted index, region_id: {}, file_id: {}",
158                self.region_id, self.file_id,
159            );
160        }
161
162        false
163    }
164
165    async fn do_finish_fulltext_index(
166        &mut self,
167        puffin_writer: &mut SstPuffinWriter,
168        index_output: &mut IndexOutput,
169    ) -> bool {
170        let Some(mut indexer) = self.fulltext_indexer.take() else {
171            return true;
172        };
173
174        let column_ids = indexer
175            .indexed_targets()
176            .filter_map(|target| match target {
177                IndexTarget::ColumnId(id) => Some(*id),
178                IndexTarget::JsonPath(_) => None,
179            })
180            .collect();
181        index_output.fulltext_bloom_columns = indexer
182            .bloom_indexed_targets()
183            .filter_map(|target| match target {
184                IndexTarget::ColumnId(id) => Some(*id),
185                IndexTarget::JsonPath(_) => None,
186            })
187            .collect();
188        let err = match indexer.finish(puffin_writer).await {
189            Ok((row_count, byte_count)) => {
190                self.fill_fulltext_index_output(
191                    &mut index_output.fulltext_index,
192                    row_count,
193                    byte_count,
194                    column_ids,
195                );
196                return true;
197            }
198            Err(err) => err,
199        };
200
201        if cfg!(any(test, feature = "test")) {
202            panic!(
203                "Failed to finish full-text index, region_id: {}, file_id: {}, err: {:?}",
204                self.region_id, self.file_id, err
205            );
206        } else {
207            warn!(
208                err; "Failed to finish full-text index, region_id: {}, file_id: {}",
209                self.region_id, self.file_id,
210            );
211        }
212
213        false
214    }
215
216    async fn do_finish_bloom_filter(
217        &mut self,
218        puffin_writer: &mut SstPuffinWriter,
219        index_output: &mut IndexOutput,
220    ) -> bool {
221        let Some(mut indexer) = self.bloom_filter_indexer.take() else {
222            return true;
223        };
224
225        let column_ids = indexer
226            .indexed_targets()
227            .filter_map(|target| match target {
228                IndexTarget::ColumnId(id) => Some(*id),
229                IndexTarget::JsonPath(_) => None,
230            })
231            .collect();
232        let err = match indexer.finish(puffin_writer).await {
233            Ok((row_count, byte_count)) => {
234                self.fill_bloom_filter_output(
235                    &mut index_output.bloom_filter,
236                    row_count,
237                    byte_count,
238                    column_ids,
239                );
240                return true;
241            }
242            Err(err) => err,
243        };
244
245        if cfg!(any(test, feature = "test")) {
246            panic!(
247                "Failed to finish bloom filter, region_id: {}, file_id: {}, err: {:?}",
248                self.region_id, self.file_id, err
249            );
250        } else {
251            warn!(
252                err; "Failed to finish bloom filter, region_id: {}, file_id: {}",
253                self.region_id, self.file_id,
254            );
255        }
256
257        false
258    }
259
260    fn fill_inverted_index_output(
261        &mut self,
262        output: &mut InvertedIndexOutput,
263        row_count: RowCount,
264        byte_count: ByteCount,
265        column_ids: Vec<ColumnId>,
266    ) {
267        debug!(
268            "Inverted index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
269            self.region_id, self.file_id, byte_count, row_count, column_ids
270        );
271
272        output.index_size = byte_count;
273        output.row_count = row_count;
274        output.columns = column_ids;
275    }
276
277    fn fill_fulltext_index_output(
278        &mut self,
279        output: &mut FulltextIndexOutput,
280        row_count: RowCount,
281        byte_count: ByteCount,
282        column_ids: Vec<ColumnId>,
283    ) {
284        debug!(
285            "Full-text index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
286            self.region_id, self.file_id, byte_count, row_count, column_ids
287        );
288
289        output.index_size = byte_count;
290        output.row_count = row_count;
291        output.columns = column_ids;
292    }
293
294    fn fill_bloom_filter_output(
295        &mut self,
296        output: &mut BloomFilterOutput,
297        row_count: RowCount,
298        byte_count: ByteCount,
299        column_ids: Vec<ColumnId>,
300    ) {
301        debug!(
302            "Bloom filter created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
303            self.region_id, self.file_id, byte_count, row_count, column_ids
304        );
305
306        output.index_size = byte_count;
307        output.row_count = row_count;
308        output.columns = column_ids;
309    }
310
311    pub(crate) async fn do_prune_intm_sst_dir(&mut self) {
312        if let Some(manager) = self.intermediate_manager.take()
313            && let Err(e) = manager.prune_sst_dir(&self.region_id, &self.file_id).await
314        {
315            warn!(e; "Failed to prune intermediate SST directory, region_id: {}, file_id: {}", self.region_id, self.file_id);
316        }
317    }
318}