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 puffin::puffin_manager::{PuffinManager, PuffinWriter};
17use store_api::storage::ColumnId;
18
19use crate::sst::file::{RegionFileId, RegionIndexId};
20use crate::sst::index::puffin_manager::SstPuffinWriter;
21use crate::sst::index::statistics::{ByteCount, RowCount};
22use crate::sst::index::{
23    BloomFilterOutput, FulltextIndexOutput, IndexOutput, Indexer, InvertedIndexOutput,
24};
25
26impl Indexer {
27    pub(crate) async fn do_finish(&mut self) -> IndexOutput {
28        self.dense_pk_decoder = None;
29        let mut output = IndexOutput::default();
30
31        let Some(mut writer) = self.build_puffin_writer().await else {
32            self.do_abort().await;
33            return output;
34        };
35
36        let success = self
37            .do_finish_inverted_index(&mut writer, &mut output)
38            .await;
39        if !success {
40            self.do_abort().await;
41            return IndexOutput::default();
42        }
43
44        let success = self
45            .do_finish_fulltext_index(&mut writer, &mut output)
46            .await;
47        if !success {
48            self.do_abort().await;
49            return IndexOutput::default();
50        }
51
52        let success = self.do_finish_bloom_filter(&mut writer, &mut output).await;
53        if !success {
54            self.do_abort().await;
55            return IndexOutput::default();
56        }
57
58        self.do_prune_intm_sst_dir().await;
59        output.file_size = self.do_finish_puffin_writer(writer).await;
60        output.version = self.index_version;
61        output
62    }
63
64    async fn build_puffin_writer(&mut self) -> Option<SstPuffinWriter> {
65        let puffin_manager = self.puffin_manager.clone()?;
66
67        let err = match puffin_manager
68            .writer(&RegionIndexId::new(
69                RegionFileId::new(self.physical_region_id, self.file_id),
70                self.index_version,
71            ))
72            .await
73        {
74            Ok(writer) => return Some(writer),
75            Err(err) => err,
76        };
77
78        if cfg!(any(test, feature = "test")) {
79            panic!(
80                "Failed to create puffin writer, region_id: {}, file_id: {}, err: {:?}",
81                self.region_id, self.file_id, err
82            );
83        } else {
84            warn!(
85                err; "Failed to create puffin writer, region_id: {}, file_id: {}",
86                self.region_id, self.file_id,
87            );
88        }
89
90        None
91    }
92
93    async fn do_finish_puffin_writer(&mut self, writer: SstPuffinWriter) -> ByteCount {
94        let err = match writer.finish().await {
95            Ok(size) => return size,
96            Err(err) => err,
97        };
98
99        if cfg!(any(test, feature = "test")) {
100            panic!(
101                "Failed to finish puffin writer, region_id: {}, file_id: {}, err: {:?}",
102                self.region_id, self.file_id, err
103            );
104        } else {
105            warn!(
106                err; "Failed to finish puffin writer, region_id: {}, file_id: {}",
107                self.region_id, self.file_id,
108            );
109        }
110
111        0
112    }
113
114    /// Returns false if the finish failed.
115    async fn do_finish_inverted_index(
116        &mut self,
117        puffin_writer: &mut SstPuffinWriter,
118        index_output: &mut IndexOutput,
119    ) -> bool {
120        let Some(mut indexer) = self.inverted_indexer.take() else {
121            return true;
122        };
123
124        let column_ids = indexer.column_ids().collect();
125        let err = match indexer.finish(puffin_writer).await {
126            Ok((row_count, byte_count)) => {
127                self.fill_inverted_index_output(
128                    &mut index_output.inverted_index,
129                    row_count,
130                    byte_count,
131                    column_ids,
132                );
133                return true;
134            }
135            Err(err) => err,
136        };
137
138        if cfg!(any(test, feature = "test")) {
139            panic!(
140                "Failed to finish inverted index, region_id: {}, file_id: {}, err: {:?}",
141                self.region_id, self.file_id, err
142            );
143        } else {
144            warn!(
145                err; "Failed to finish inverted index, region_id: {}, file_id: {}",
146                self.region_id, self.file_id,
147            );
148        }
149
150        false
151    }
152
153    async fn do_finish_fulltext_index(
154        &mut self,
155        puffin_writer: &mut SstPuffinWriter,
156        index_output: &mut IndexOutput,
157    ) -> bool {
158        let Some(mut indexer) = self.fulltext_indexer.take() else {
159            return true;
160        };
161
162        let column_ids = indexer.column_ids().collect();
163        index_output.fulltext_bloom_columns = indexer.bloom_column_ids().collect();
164        let err = match indexer.finish(puffin_writer).await {
165            Ok((row_count, byte_count)) => {
166                self.fill_fulltext_index_output(
167                    &mut index_output.fulltext_index,
168                    row_count,
169                    byte_count,
170                    column_ids,
171                );
172                return true;
173            }
174            Err(err) => err,
175        };
176
177        if cfg!(any(test, feature = "test")) {
178            panic!(
179                "Failed to finish full-text index, region_id: {}, file_id: {}, err: {:?}",
180                self.region_id, self.file_id, err
181            );
182        } else {
183            warn!(
184                err; "Failed to finish full-text index, region_id: {}, file_id: {}",
185                self.region_id, self.file_id,
186            );
187        }
188
189        false
190    }
191
192    async fn do_finish_bloom_filter(
193        &mut self,
194        puffin_writer: &mut SstPuffinWriter,
195        index_output: &mut IndexOutput,
196    ) -> bool {
197        let Some(mut indexer) = self.bloom_filter_indexer.take() else {
198            return true;
199        };
200
201        let column_ids = indexer.column_ids().collect();
202        let err = match indexer.finish(puffin_writer).await {
203            Ok((row_count, byte_count)) => {
204                self.fill_bloom_filter_output(
205                    &mut index_output.bloom_filter,
206                    row_count,
207                    byte_count,
208                    column_ids,
209                );
210                return true;
211            }
212            Err(err) => err,
213        };
214
215        if cfg!(any(test, feature = "test")) {
216            panic!(
217                "Failed to finish bloom filter, region_id: {}, file_id: {}, err: {:?}",
218                self.region_id, self.file_id, err
219            );
220        } else {
221            warn!(
222                err; "Failed to finish bloom filter, region_id: {}, file_id: {}",
223                self.region_id, self.file_id,
224            );
225        }
226
227        false
228    }
229
230    fn fill_inverted_index_output(
231        &mut self,
232        output: &mut InvertedIndexOutput,
233        row_count: RowCount,
234        byte_count: ByteCount,
235        column_ids: Vec<ColumnId>,
236    ) {
237        debug!(
238            "Inverted index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
239            self.region_id, self.file_id, byte_count, row_count, column_ids
240        );
241
242        output.index_size = byte_count;
243        output.row_count = row_count;
244        output.columns = column_ids;
245    }
246
247    fn fill_fulltext_index_output(
248        &mut self,
249        output: &mut FulltextIndexOutput,
250        row_count: RowCount,
251        byte_count: ByteCount,
252        column_ids: Vec<ColumnId>,
253    ) {
254        debug!(
255            "Full-text index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
256            self.region_id, self.file_id, byte_count, row_count, column_ids
257        );
258
259        output.index_size = byte_count;
260        output.row_count = row_count;
261        output.columns = column_ids;
262    }
263
264    fn fill_bloom_filter_output(
265        &mut self,
266        output: &mut BloomFilterOutput,
267        row_count: RowCount,
268        byte_count: ByteCount,
269        column_ids: Vec<ColumnId>,
270    ) {
271        debug!(
272            "Bloom filter created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}",
273            self.region_id, self.file_id, byte_count, row_count, column_ids
274        );
275
276        output.index_size = byte_count;
277        output.row_count = row_count;
278        output.columns = column_ids;
279    }
280
281    pub(crate) async fn do_prune_intm_sst_dir(&mut self) {
282        if let Some(manager) = self.intermediate_manager.take()
283            && let Err(e) = manager.prune_sst_dir(&self.region_id, &self.file_id).await
284        {
285            warn!(e; "Failed to prune intermediate SST directory, region_id: {}, file_id: {}", self.region_id, self.file_id);
286        }
287    }
288}