index/inverted_index/search/
fst_values_mapper.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 greptime_proto::v1::index::{BitmapType, InvertedIndexMeta};
16
17use crate::bitmap::Bitmap;
18use crate::inverted_index::error::Result;
19use crate::inverted_index::format::reader::{InvertedIndexReadMetrics, InvertedIndexReader};
20
21/// `ParallelFstValuesMapper` enables parallel mapping of multiple FST value groups to their
22/// corresponding bitmaps within an inverted index.
23///
24/// This mapper processes multiple groups of FST values in parallel, where each group is associated
25/// with its own metadata. It optimizes bitmap retrieval by batching requests across all groups
26/// before combining them into separate result bitmaps.
27pub struct ParallelFstValuesMapper<'a> {
28    reader: &'a mut dyn InvertedIndexReader,
29}
30
31impl<'a> ParallelFstValuesMapper<'a> {
32    pub fn new(reader: &'a mut dyn InvertedIndexReader) -> Self {
33        Self { reader }
34    }
35
36    pub async fn map_values_vec(
37        &mut self,
38        value_and_meta_vec: &[(Vec<u64>, &InvertedIndexMeta)],
39        metrics: Option<&mut InvertedIndexReadMetrics>,
40    ) -> Result<Vec<Bitmap>> {
41        let groups = value_and_meta_vec
42            .iter()
43            .map(|(values, _)| values.len())
44            .collect::<Vec<_>>();
45        let len = groups.iter().sum::<usize>();
46        let mut fetch_ranges = Vec::with_capacity(len);
47
48        for (values, meta) in value_and_meta_vec {
49            for value in values {
50                // The higher 32 bits of each u64 value represent the
51                // bitmap offset and the lower 32 bits represent its size. This mapper uses these
52                // combined offset-size pairs to fetch and union multiple bitmaps into a single `BitVec`.
53                let [relative_offset, size] = bytemuck::cast::<u64, [u32; 2]>(*value);
54                let range = meta.base_offset + relative_offset as u64
55                    ..meta.base_offset + relative_offset as u64 + size as u64;
56                fetch_ranges.push((
57                    range,
58                    BitmapType::try_from(meta.bitmap_type).unwrap_or(BitmapType::BitVec),
59                ));
60            }
61        }
62
63        if fetch_ranges.is_empty() {
64            return Ok(vec![Bitmap::new_bitvec()]);
65        }
66
67        common_telemetry::debug!("fetch ranges: {:?}", fetch_ranges);
68        let mut bitmaps = self.reader.bitmap_deque(&fetch_ranges, metrics).await?;
69        let mut output = Vec::with_capacity(groups.len());
70
71        for counter in groups {
72            let mut bitmap = Bitmap::new_roaring();
73            for _ in 0..counter {
74                let bm = bitmaps.pop_front().unwrap();
75                bitmap.union(bm);
76            }
77
78            output.push(bitmap);
79        }
80
81        Ok(output)
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use std::collections::VecDeque;
88
89    use super::*;
90    use crate::inverted_index::format::reader::MockInvertedIndexReader;
91
92    fn value(offset: u32, size: u32) -> u64 {
93        bytemuck::cast::<[u32; 2], u64>([offset, size])
94    }
95
96    #[tokio::test]
97    async fn test_map_values_vec() {
98        let mut mock_reader = MockInvertedIndexReader::new();
99        mock_reader
100            .expect_bitmap_deque()
101            .returning(|ranges, _metrics| {
102                let mut output = VecDeque::new();
103                for (range, bitmap_type) in ranges {
104                    let offset = range.start;
105                    let size = range.end - range.start;
106                    match (offset, size, bitmap_type) {
107                        (1, 1, BitmapType::Roaring) => {
108                            output.push_back(Bitmap::from_lsb0_bytes(&[0b10101010], *bitmap_type))
109                        }
110                        (2, 1, BitmapType::Roaring) => {
111                            output.push_back(Bitmap::from_lsb0_bytes(&[0b01010101], *bitmap_type))
112                        }
113                        _ => unreachable!(),
114                    }
115                }
116                Ok(output)
117            });
118
119        let meta = InvertedIndexMeta {
120            bitmap_type: BitmapType::Roaring.into(),
121            ..Default::default()
122        };
123        let mut values_mapper = ParallelFstValuesMapper::new(&mut mock_reader);
124
125        let result = values_mapper
126            .map_values_vec(&[(vec![], &meta)], None)
127            .await
128            .unwrap();
129        assert_eq!(result[0].count_ones(), 0);
130
131        let result = values_mapper
132            .map_values_vec(&[(vec![value(1, 1)], &meta)], None)
133            .await
134            .unwrap();
135        assert_eq!(
136            result[0],
137            Bitmap::from_lsb0_bytes(&[0b10101010], BitmapType::Roaring)
138        );
139
140        let result = values_mapper
141            .map_values_vec(&[(vec![value(2, 1)], &meta)], None)
142            .await
143            .unwrap();
144        assert_eq!(
145            result[0],
146            Bitmap::from_lsb0_bytes(&[0b01010101], BitmapType::Roaring)
147        );
148
149        let result = values_mapper
150            .map_values_vec(&[(vec![value(1, 1), value(2, 1)], &meta)], None)
151            .await
152            .unwrap();
153        assert_eq!(
154            result[0],
155            Bitmap::from_lsb0_bytes(&[0b11111111], BitmapType::Roaring)
156        );
157
158        let result = values_mapper
159            .map_values_vec(&[(vec![value(2, 1), value(1, 1)], &meta)], None)
160            .await
161            .unwrap();
162        assert_eq!(
163            result[0],
164            Bitmap::from_lsb0_bytes(&[0b11111111], BitmapType::Roaring)
165        );
166
167        let result = values_mapper
168            .map_values_vec(
169                &[(vec![value(2, 1)], &meta), (vec![value(1, 1)], &meta)],
170                None,
171            )
172            .await
173            .unwrap();
174        assert_eq!(
175            result[0],
176            Bitmap::from_lsb0_bytes(&[0b01010101], BitmapType::Roaring)
177        );
178        assert_eq!(
179            result[1],
180            Bitmap::from_lsb0_bytes(&[0b10101010], BitmapType::Roaring)
181        );
182        let result = values_mapper
183            .map_values_vec(
184                &[
185                    (vec![value(2, 1), value(1, 1)], &meta),
186                    (vec![value(1, 1)], &meta),
187                ],
188                None,
189            )
190            .await
191            .unwrap();
192        assert_eq!(
193            result[0],
194            Bitmap::from_lsb0_bytes(&[0b11111111], BitmapType::Roaring)
195        );
196        assert_eq!(
197            result[1],
198            Bitmap::from_lsb0_bytes(&[0b10101010], BitmapType::Roaring)
199        );
200    }
201}
index/inverted_index/search/fst_values_mapper.rs

index/inverted_index/search/
fst_values_mapper.rs
