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

index/inverted_index/search/
fst_values_mapper.rs
