index/inverted_index/create/
sort_create.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 std::collections::HashMap;
16use std::num::NonZeroUsize;
17
18use async_trait::async_trait;
19use snafu::ensure;
20
21use crate::BytesRef;
22use crate::bitmap::BitmapType;
23use crate::inverted_index::create::InvertedIndexCreator;
24use crate::inverted_index::create::sort::{SortOutput, Sorter};
25use crate::inverted_index::error::{InconsistentRowCountSnafu, Result};
26use crate::inverted_index::format::writer::InvertedIndexWriter;
27
28type IndexName = String;
29type SegmentRowCount = NonZeroUsize;
30
31/// Factory type to produce `Sorter` instances associated with an index name and segment row count
32pub type SorterFactory = Box<dyn Fn(IndexName, SegmentRowCount) -> Box<dyn Sorter> + Send>;
33
34/// `SortIndexCreator` orchestrates indexing by sorting input data for each named index
35/// and writing to an inverted index writer
36pub struct SortIndexCreator {
37    /// Factory for producing `Sorter` instances
38    sorter_factory: SorterFactory,
39
40    /// Map of index names to sorters
41    sorters: HashMap<IndexName, Box<dyn Sorter>>,
42
43    /// Number of rows in each segment, used to produce sorters
44    segment_row_count: NonZeroUsize,
45
46    /// Indexes whose sorter asked to spill its buffer.
47    pending_spills: Vec<IndexName>,
48}
49
50#[async_trait]
51impl InvertedIndexCreator for SortIndexCreator {
52    /// Inserts `n` values or nulls into the sorter for the specified index.
53    ///
54    /// If the index does not exist, a new index is created even if `n` is 0.
55    /// Caller may leverage this behavior to create indexes with no data.
56    fn push_with_name_n(
57        &mut self,
58        index_name: &str,
59        value: Option<BytesRef<'_>>,
60        n: usize,
61    ) -> bool {
62        let sorter = match self.sorters.get_mut(index_name) {
63            Some(sorter) => sorter,
64            None => {
65                let sorter = (self.sorter_factory)(index_name.to_string(), self.segment_row_count);
66                self.sorters.entry(index_name.to_string()).or_insert(sorter)
67            }
68        };
69        if sorter.push_n(value, n) {
70            self.pending_spills.push(index_name.to_string());
71            return true;
72        }
73        false
74    }
75
76    async fn spill(&mut self) -> Result<()> {
77        for index_name in std::mem::take(&mut self.pending_spills) {
78            if let Some(sorter) = self.sorters.get_mut(&index_name) {
79                sorter.spill().await?;
80            }
81        }
82        Ok(())
83    }
84
85    /// Finalizes the sorting for all indexes and writes them using the inverted index writer
86    async fn finish(
87        &mut self,
88        writer: &mut dyn InvertedIndexWriter,
89        bitmap_type: BitmapType,
90    ) -> Result<()> {
91        let mut output_row_count = None;
92        for (index_name, mut sorter) in self.sorters.drain() {
93            let SortOutput {
94                segment_null_bitmap,
95                sorted_stream,
96                total_row_count,
97            } = sorter.output().await?;
98
99            let expected_row_count = *output_row_count.get_or_insert(total_row_count);
100            ensure!(
101                expected_row_count == total_row_count,
102                InconsistentRowCountSnafu {
103                    index_name,
104                    total_row_count,
105                    expected_row_count,
106                }
107            );
108
109            writer
110                .add_index(index_name, segment_null_bitmap, sorted_stream, bitmap_type)
111                .await?;
112        }
113
114        let total_row_count = output_row_count.unwrap_or_default() as _;
115        let segment_row_count = self.segment_row_count as _;
116        writer.finish(total_row_count, segment_row_count).await
117    }
118}
119
120impl SortIndexCreator {
121    /// Creates a new `SortIndexCreator` with the given sorter factory and index writer
122    pub fn new(sorter_factory: SorterFactory, segment_row_count: NonZeroUsize) -> Self {
123        Self {
124            sorter_factory,
125            sorters: HashMap::new(),
126            segment_row_count,
127            pending_spills: Vec::new(),
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use std::collections::BTreeMap;
135    use std::sync::{Arc, Mutex};
136
137    use common_base::BitVec;
138    use futures::{StreamExt, stream};
139
140    use super::*;
141    use crate::Bytes;
142    use crate::bitmap::Bitmap;
143    use crate::inverted_index::error::Error;
144    use crate::inverted_index::format::writer::{MockInvertedIndexWriter, ValueStream};
145
146    #[tokio::test]
147    async fn test_sort_index_creator_basic() {
148        let mut creator =
149            SortIndexCreator::new(NaiveSorter::factory(), NonZeroUsize::new(1).unwrap());
150
151        let index_values = vec![
152            ("a", vec![b"3", b"2", b"1"]),
153            ("b", vec![b"6", b"5", b"4"]),
154            ("c", vec![b"1", b"2", b"3"]),
155        ];
156
157        for (index_name, values) in index_values {
158            for value in values {
159                assert!(!creator.push_with_name(index_name, Some(value)));
160            }
161        }
162
163        let mut mock_writer = MockInvertedIndexWriter::new();
164        mock_writer.expect_add_index().times(3).returning(
165            |name, null_bitmap, stream, bitmap_type| {
166                assert!(null_bitmap.is_empty());
167                assert_eq!(bitmap_type, BitmapType::Roaring);
168                match name.as_str() {
169                    "a" => assert_eq!(stream_to_values(stream), vec![b"1", b"2", b"3"]),
170                    "b" => assert_eq!(stream_to_values(stream), vec![b"4", b"5", b"6"]),
171                    "c" => assert_eq!(stream_to_values(stream), vec![b"1", b"2", b"3"]),
172                    _ => panic!("unexpected index name: {}", name),
173                }
174                Ok(())
175            },
176        );
177        mock_writer
178            .expect_finish()
179            .times(1)
180            .returning(|total_row_count, segment_row_count| {
181                assert_eq!(total_row_count, 3);
182                assert_eq!(segment_row_count.get(), 1);
183                Ok(())
184            });
185
186        creator
187            .finish(&mut mock_writer, BitmapType::Roaring)
188            .await
189            .unwrap();
190    }
191
192    #[tokio::test]
193    async fn test_sort_index_creator_inconsistent_row_count() {
194        let mut creator =
195            SortIndexCreator::new(NaiveSorter::factory(), NonZeroUsize::new(1).unwrap());
196
197        let index_values = vec![
198            ("a", vec![b"3", b"2", b"1"]),
199            ("b", vec![b"6", b"5", b"4"]),
200            ("c", vec![b"1", b"2"]),
201        ];
202
203        for (index_name, values) in index_values {
204            for value in values {
205                assert!(!creator.push_with_name(index_name, Some(value)));
206            }
207        }
208
209        let mut mock_writer = MockInvertedIndexWriter::new();
210        mock_writer
211            .expect_add_index()
212            .returning(|name, null_bitmap, stream, bitmap_type| {
213                assert!(null_bitmap.is_empty());
214                assert_eq!(bitmap_type, BitmapType::Roaring);
215                match name.as_str() {
216                    "a" => assert_eq!(stream_to_values(stream), vec![b"1", b"2", b"3"]),
217                    "b" => assert_eq!(stream_to_values(stream), vec![b"4", b"5", b"6"]),
218                    "c" => assert_eq!(stream_to_values(stream), vec![b"1", b"2"]),
219                    _ => panic!("unexpected index name: {}", name),
220                }
221                Ok(())
222            });
223        mock_writer.expect_finish().never();
224
225        let res = creator.finish(&mut mock_writer, BitmapType::Roaring).await;
226        assert!(matches!(res, Err(Error::InconsistentRowCount { .. })));
227    }
228
229    #[tokio::test]
230    async fn test_sort_index_creator_create_indexes_without_data() {
231        let mut creator =
232            SortIndexCreator::new(NaiveSorter::factory(), NonZeroUsize::new(1).unwrap());
233
234        assert!(!creator.push_with_name_n("a", None, 0));
235        assert!(!creator.push_with_name_n("b", None, 0));
236        assert!(!creator.push_with_name_n("c", None, 0));
237
238        let mut mock_writer = MockInvertedIndexWriter::new();
239        mock_writer
240            .expect_add_index()
241            .returning(|name, null_bitmap, stream, bitmap_type| {
242                assert!(null_bitmap.is_empty());
243                assert_eq!(bitmap_type, BitmapType::Roaring);
244                assert!(matches!(name.as_str(), "a" | "b" | "c"));
245                assert!(stream_to_values(stream).is_empty());
246                Ok(())
247            });
248        mock_writer
249            .expect_finish()
250            .times(1)
251            .returning(|total_row_count, segment_row_count| {
252                assert_eq!(total_row_count, 0);
253                assert_eq!(segment_row_count.get(), 1);
254                Ok(())
255            });
256
257        creator
258            .finish(&mut mock_writer, BitmapType::Roaring)
259            .await
260            .unwrap();
261    }
262
263    #[tokio::test]
264    async fn test_sort_index_creator_spills_requested_sorters() {
265        let spilled = Arc::new(Mutex::new(Vec::new()));
266        let factory: SorterFactory = {
267            let spilled = spilled.clone();
268            Box::new(move |index_name, _| {
269                Box::new(SpillRecordingSorter {
270                    index_name,
271                    spilled: spilled.clone(),
272                })
273            })
274        };
275        let mut creator = SortIndexCreator::new(factory, NonZeroUsize::new(1).unwrap());
276
277        assert!(creator.push_with_name("a", Some(b"1")));
278        assert!(!creator.push_with_name("b", Some(b"1")));
279        creator.spill().await.unwrap();
280        assert_eq!(*spilled.lock().unwrap(), vec!["a"]);
281
282        creator.spill().await.unwrap();
283        assert_eq!(*spilled.lock().unwrap(), vec!["a"]);
284    }
285
286    /// Requests a spill on every push to index `a` and records the spills it receives.
287    struct SpillRecordingSorter {
288        index_name: String,
289        spilled: Arc<Mutex<Vec<String>>>,
290    }
291
292    #[async_trait]
293    impl Sorter for SpillRecordingSorter {
294        fn push_n(&mut self, _value: Option<BytesRef<'_>>, _n: usize) -> bool {
295            self.index_name == "a"
296        }
297
298        async fn spill(&mut self) -> Result<()> {
299            self.spilled.lock().unwrap().push(self.index_name.clone());
300            Ok(())
301        }
302
303        async fn output(&mut self) -> Result<SortOutput> {
304            unreachable!()
305        }
306    }
307
308    fn set_bit(bit_vec: &mut BitVec, index: usize) {
309        if index >= bit_vec.len() {
310            bit_vec.resize(index + 1, false);
311        }
312        bit_vec.set(index, true);
313    }
314
315    struct NaiveSorter {
316        total_row_count: usize,
317        segment_row_count: NonZeroUsize,
318        values: BTreeMap<Option<Bytes>, BitVec>,
319    }
320
321    impl NaiveSorter {
322        fn factory() -> SorterFactory {
323            Box::new(|_index_name, segment_row_count| {
324                Box::new(NaiveSorter {
325                    total_row_count: 0,
326                    segment_row_count,
327                    values: BTreeMap::new(),
328                })
329            })
330        }
331    }
332
333    #[async_trait]
334    impl Sorter for NaiveSorter {
335        fn push_n(&mut self, value: Option<BytesRef<'_>>, n: usize) -> bool {
336            for _ in 0..n {
337                let segment_index = self.total_row_count / self.segment_row_count;
338                self.total_row_count += 1;
339
340                let bitmap = self.values.entry(value.map(Into::into)).or_default();
341                set_bit(bitmap, segment_index);
342            }
343            false
344        }
345
346        async fn spill(&mut self) -> Result<()> {
347            Ok(())
348        }
349
350        async fn output(&mut self) -> Result<SortOutput> {
351            let segment_null_bitmap = self.values.remove(&None).unwrap_or_default();
352            let segment_null_bitmap = Bitmap::BitVec(segment_null_bitmap);
353
354            Ok(SortOutput {
355                segment_null_bitmap,
356                sorted_stream: Box::new(stream::iter(
357                    std::mem::take(&mut self.values)
358                        .into_iter()
359                        .map(|(v, b)| Ok((v.unwrap(), Bitmap::BitVec(b)))),
360                )),
361                total_row_count: self.total_row_count,
362            })
363        }
364    }
365
366    fn stream_to_values(stream: ValueStream) -> Vec<Bytes> {
367        futures::executor::block_on(async {
368            stream.map(|r| r.unwrap().0).collect::<Vec<Bytes>>().await
369        })
370    }
371}
index/inverted_index/create/sort_create.rs

index/inverted_index/create/
sort_create.rs
