mito2/sst/index/inverted_index/applier/
builder.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
15mod between;
16mod comparison;
17mod eq_list;
18mod in_list;
19mod regex_match;
20
21use std::collections::{HashMap, HashSet};
22
23use common_telemetry::warn;
24use datafusion_common::ScalarValue;
25use datafusion_expr::{BinaryExpr, Expr, Operator};
26use datatypes::data_type::ConcreteDataType;
27use datatypes::value::Value;
28use index::inverted_index::search::index_apply::PredicatesIndexApplier;
29use index::inverted_index::search::predicate::Predicate;
30use object_store::ObjectStore;
31use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
32use snafu::{OptionExt, ResultExt};
33use store_api::metadata::RegionMetadata;
34use store_api::storage::ColumnId;
35
36use crate::cache::file_cache::FileCacheRef;
37use crate::cache::index::inverted_index::InvertedIndexCacheRef;
38use crate::error::{BuildIndexApplierSnafu, ColumnNotFoundSnafu, ConvertValueSnafu, Result};
39use crate::row_converter::SortField;
40use crate::sst::index::codec::IndexValueCodec;
41use crate::sst::index::inverted_index::applier::InvertedIndexApplier;
42use crate::sst::index::puffin_manager::PuffinManagerFactory;
43
44/// Constructs an [`InvertedIndexApplier`] which applies predicates to SST files during scan.
45pub(crate) struct InvertedIndexApplierBuilder<'a> {
46    /// Directory of the region, required argument for constructing [`InvertedIndexApplier`].
47    region_dir: String,
48
49    /// Object store, required argument for constructing [`InvertedIndexApplier`].
50    object_store: ObjectStore,
51
52    /// File cache, required argument for constructing [`InvertedIndexApplier`].
53    file_cache: Option<FileCacheRef>,
54
55    /// Metadata of the region, used to get metadata like column type.
56    metadata: &'a RegionMetadata,
57
58    /// Column ids of the columns that are indexed.
59    indexed_column_ids: HashSet<ColumnId>,
60
61    /// Stores predicates during traversal on the Expr tree.
62    output: HashMap<ColumnId, Vec<Predicate>>,
63
64    /// The puffin manager factory.
65    puffin_manager_factory: PuffinManagerFactory,
66
67    /// Cache for inverted index.
68    inverted_index_cache: Option<InvertedIndexCacheRef>,
69
70    /// Cache for puffin metadata.
71    puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
72}
73
74impl<'a> InvertedIndexApplierBuilder<'a> {
75    /// Creates a new [`InvertedIndexApplierBuilder`].
76    pub fn new(
77        region_dir: String,
78        object_store: ObjectStore,
79        metadata: &'a RegionMetadata,
80        indexed_column_ids: HashSet<ColumnId>,
81        puffin_manager_factory: PuffinManagerFactory,
82    ) -> Self {
83        Self {
84            region_dir,
85            object_store,
86            metadata,
87            indexed_column_ids,
88            output: HashMap::default(),
89            puffin_manager_factory,
90            file_cache: None,
91            inverted_index_cache: None,
92            puffin_metadata_cache: None,
93        }
94    }
95
96    /// Sets the file cache.
97    pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
98        self.file_cache = file_cache;
99        self
100    }
101
102    /// Sets the puffin metadata cache.
103    pub fn with_puffin_metadata_cache(
104        mut self,
105        puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
106    ) -> Self {
107        self.puffin_metadata_cache = puffin_metadata_cache;
108        self
109    }
110
111    /// Sets the inverted index cache.
112    pub fn with_inverted_index_cache(
113        mut self,
114        inverted_index_cache: Option<InvertedIndexCacheRef>,
115    ) -> Self {
116        self.inverted_index_cache = inverted_index_cache;
117        self
118    }
119
120    /// Consumes the builder to construct an [`InvertedIndexApplier`], optionally returned based on
121    /// the expressions provided. If no predicates match, returns `None`.
122    pub fn build(mut self, exprs: &[Expr]) -> Result<Option<InvertedIndexApplier>> {
123        for expr in exprs {
124            self.traverse_and_collect(expr);
125        }
126
127        if self.output.is_empty() {
128            return Ok(None);
129        }
130
131        let predicates = self
132            .output
133            .into_iter()
134            .map(|(column_id, predicates)| (column_id.to_string(), predicates))
135            .collect();
136        let applier = PredicatesIndexApplier::try_from(predicates);
137
138        Ok(Some(
139            InvertedIndexApplier::new(
140                self.region_dir,
141                self.metadata.region_id,
142                self.object_store,
143                Box::new(applier.context(BuildIndexApplierSnafu)?),
144                self.puffin_manager_factory,
145            )
146            .with_file_cache(self.file_cache)
147            .with_puffin_metadata_cache(self.puffin_metadata_cache)
148            .with_index_cache(self.inverted_index_cache),
149        ))
150    }
151
152    /// Recursively traverses expressions to collect predicates.
153    /// Results are stored in `self.output`.
154    fn traverse_and_collect(&mut self, expr: &Expr) {
155        let res = match expr {
156            Expr::Between(between) => self.collect_between(between),
157
158            Expr::InList(in_list) => self.collect_inlist(in_list),
159            Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
160                Operator::And => {
161                    self.traverse_and_collect(left);
162                    self.traverse_and_collect(right);
163                    Ok(())
164                }
165                Operator::Or => self.collect_or_eq_list(left, right),
166                Operator::Eq => self.collect_eq(left, right),
167                Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => {
168                    self.collect_comparison_expr(left, op, right)
169                }
170                Operator::RegexMatch => self.collect_regex_match(left, right),
171                _ => Ok(()),
172            },
173
174            // TODO(zhongzc): support more expressions, e.g. IsNull, IsNotNull, ...
175            _ => Ok(()),
176        };
177
178        if let Err(err) = res {
179            warn!(err; "Failed to collect predicates, ignore it. expr: {expr}");
180        }
181    }
182
183    /// Helper function to add a predicate to the output.
184    fn add_predicate(&mut self, column_id: ColumnId, predicate: Predicate) {
185        self.output.entry(column_id).or_default().push(predicate);
186    }
187
188    /// Helper function to get the column id and the column type of a column.
189    /// Returns `None` if the column is not a tag column or if the column is ignored.
190    fn column_id_and_type(
191        &self,
192        column_name: &str,
193    ) -> Result<Option<(ColumnId, ConcreteDataType)>> {
194        let column = self
195            .metadata
196            .column_by_name(column_name)
197            .context(ColumnNotFoundSnafu {
198                column: column_name,
199            })?;
200
201        if !self.indexed_column_ids.contains(&column.column_id) {
202            return Ok(None);
203        }
204
205        Ok(Some((
206            column.column_id,
207            column.column_schema.data_type.clone(),
208        )))
209    }
210
211    /// Helper function to get a non-null literal.
212    fn nonnull_lit(expr: &Expr) -> Option<&ScalarValue> {
213        match expr {
214            Expr::Literal(lit) if !lit.is_null() => Some(lit),
215            _ => None,
216        }
217    }
218
219    /// Helper function to get the column name of a column expression.
220    fn column_name(expr: &Expr) -> Option<&str> {
221        match expr {
222            Expr::Column(column) => Some(&column.name),
223            _ => None,
224        }
225    }
226
227    /// Helper function to encode a literal into bytes.
228    fn encode_lit(lit: &ScalarValue, data_type: ConcreteDataType) -> Result<Vec<u8>> {
229        let value = Value::try_from(lit.clone()).context(ConvertValueSnafu)?;
230        let mut bytes = vec![];
231        let field = SortField::new(data_type);
232        IndexValueCodec::encode_nonnull_value(value.as_value_ref(), &field, &mut bytes)?;
233        Ok(bytes)
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use api::v1::SemanticType;
240    use datafusion_common::Column;
241    use datafusion_expr::Between;
242    use datatypes::data_type::ConcreteDataType;
243    use datatypes::schema::ColumnSchema;
244    use index::inverted_index::search::predicate::{
245        Bound, Range, RangePredicate, RegexMatchPredicate,
246    };
247    use object_store::services::Memory;
248    use object_store::ObjectStore;
249    use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
250    use store_api::storage::RegionId;
251
252    use super::*;
253
254    pub(crate) fn test_region_metadata() -> RegionMetadata {
255        let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5678));
256        builder
257            .push_column_metadata(ColumnMetadata {
258                column_schema: ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
259                semantic_type: SemanticType::Tag,
260                column_id: 1,
261            })
262            .push_column_metadata(ColumnMetadata {
263                column_schema: ColumnSchema::new("b", ConcreteDataType::int64_datatype(), false),
264                semantic_type: SemanticType::Tag,
265                column_id: 2,
266            })
267            .push_column_metadata(ColumnMetadata {
268                column_schema: ColumnSchema::new("c", ConcreteDataType::string_datatype(), false),
269                semantic_type: SemanticType::Field,
270                column_id: 3,
271            })
272            .push_column_metadata(ColumnMetadata {
273                column_schema: ColumnSchema::new(
274                    "d",
275                    ConcreteDataType::timestamp_millisecond_datatype(),
276                    false,
277                ),
278                semantic_type: SemanticType::Timestamp,
279                column_id: 4,
280            })
281            .primary_key(vec![1, 2]);
282        builder.build().unwrap()
283    }
284
285    pub(crate) fn test_object_store() -> ObjectStore {
286        ObjectStore::new(Memory::default()).unwrap().finish()
287    }
288
289    pub(crate) fn tag_column() -> Expr {
290        Expr::Column(Column::from_name("a"))
291    }
292
293    pub(crate) fn tag_column2() -> Expr {
294        Expr::Column(Column::from_name("b"))
295    }
296
297    pub(crate) fn field_column() -> Expr {
298        Expr::Column(Column::from_name("c"))
299    }
300
301    pub(crate) fn nonexistent_column() -> Expr {
302        Expr::Column(Column::from_name("nonexistence"))
303    }
304
305    pub(crate) fn string_lit(s: impl Into<String>) -> Expr {
306        Expr::Literal(ScalarValue::Utf8(Some(s.into())))
307    }
308
309    pub(crate) fn int64_lit(i: impl Into<i64>) -> Expr {
310        Expr::Literal(ScalarValue::Int64(Some(i.into())))
311    }
312
313    pub(crate) fn encoded_string(s: impl Into<String>) -> Vec<u8> {
314        let mut bytes = vec![];
315        IndexValueCodec::encode_nonnull_value(
316            Value::from(s.into()).as_value_ref(),
317            &SortField::new(ConcreteDataType::string_datatype()),
318            &mut bytes,
319        )
320        .unwrap();
321        bytes
322    }
323
324    pub(crate) fn encoded_int64(s: impl Into<i64>) -> Vec<u8> {
325        let mut bytes = vec![];
326        IndexValueCodec::encode_nonnull_value(
327            Value::from(s.into()).as_value_ref(),
328            &SortField::new(ConcreteDataType::int64_datatype()),
329            &mut bytes,
330        )
331        .unwrap();
332        bytes
333    }
334
335    #[test]
336    fn test_collect_and_basic() {
337        let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_and_basic_");
338
339        let metadata = test_region_metadata();
340        let mut builder = InvertedIndexApplierBuilder::new(
341            "test".to_string(),
342            test_object_store(),
343            &metadata,
344            HashSet::from_iter([1, 2, 3]),
345            facotry,
346        );
347
348        let expr = Expr::BinaryExpr(BinaryExpr {
349            left: Box::new(Expr::BinaryExpr(BinaryExpr {
350                left: Box::new(tag_column()),
351                op: Operator::RegexMatch,
352                right: Box::new(string_lit("bar")),
353            })),
354            op: Operator::And,
355            right: Box::new(Expr::Between(Between {
356                expr: Box::new(tag_column2()),
357                negated: false,
358                low: Box::new(int64_lit(123)),
359                high: Box::new(int64_lit(456)),
360            })),
361        });
362
363        builder.traverse_and_collect(&expr);
364        let predicates = builder.output.get(&1).unwrap();
365        assert_eq!(predicates.len(), 1);
366        assert_eq!(
367            predicates[0],
368            Predicate::RegexMatch(RegexMatchPredicate {
369                pattern: "bar".to_string()
370            })
371        );
372        let predicates = builder.output.get(&2).unwrap();
373        assert_eq!(predicates.len(), 1);
374        assert_eq!(
375            predicates[0],
376            Predicate::Range(RangePredicate {
377                range: Range {
378                    lower: Some(Bound {
379                        inclusive: true,
380                        value: encoded_int64(123),
381                    }),
382                    upper: Some(Bound {
383                        inclusive: true,
384                        value: encoded_int64(456),
385                    }),
386                }
387            })
388        );
389    }
390}
mito2/sst/index/inverted_index/applier/builder.rs

mito2/sst/index/inverted_index/applier/
builder.rs
