Skip to main content

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::{BTreeMap, HashSet};
22
23#[cfg(test)]
24use common_function::scalars::json::json_get::JsonGetWithType;
25use common_telemetry::warn;
26use datafusion_common::ScalarValue;
27use datafusion_expr::{BinaryExpr, Expr, Operator};
28use datatypes::data_type::ConcreteDataType;
29#[cfg(test)]
30use datatypes::extension::json::Json2ExtensionType;
31use datatypes::value::Value;
32use index::inverted_index::search::predicate::Predicate;
33use index::target::IndexTarget;
34use mito_codec::index::IndexValueCodec;
35use mito_codec::row_converter::SortField;
36use object_store::ObjectStore;
37use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
38use snafu::{OptionExt, ResultExt};
39use store_api::metadata::RegionMetadata;
40use store_api::region_request::PathType;
41use store_api::storage::ColumnId;
42
43use crate::cache::file_cache::FileCacheRef;
44use crate::cache::index::inverted_index::InvertedIndexCacheRef;
45use crate::error::{ColumnNotFoundSnafu, ConvertValueSnafu, EncodeSnafu, Result};
46use crate::sst::index::inverted_index::applier::{InvertedIndexApplier, index_target_data_type};
47use crate::sst::index::puffin_manager::PuffinManagerFactory;
48
49/// Constructs an [`InvertedIndexApplier`] which applies predicates to SST files during scan.
50pub(crate) struct InvertedIndexApplierBuilder<'a> {
51    /// Directory of the table, required argument for constructing [`InvertedIndexApplier`].
52    table_dir: String,
53
54    /// Path type for generating file paths.
55    path_type: PathType,
56
57    /// Object store, required argument for constructing [`InvertedIndexApplier`].
58    object_store: ObjectStore,
59
60    /// File cache, required argument for constructing [`InvertedIndexApplier`].
61    file_cache: Option<FileCacheRef>,
62
63    /// Metadata of the region, used to get metadata like column type.
64    metadata: &'a RegionMetadata,
65
66    /// Column ids with column-level inverted indexes, excluding JSON path indexes.
67    indexed_column_ids: HashSet<ColumnId>,
68
69    /// Stores predicates during traversal on the Expr tree.
70    output: BTreeMap<IndexTarget, Vec<Predicate>>,
71
72    /// The puffin manager factory.
73    puffin_manager_factory: PuffinManagerFactory,
74
75    /// Cache for inverted index.
76    inverted_index_cache: Option<InvertedIndexCacheRef>,
77
78    /// Cache for puffin metadata.
79    puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
80}
81
82impl<'a> InvertedIndexApplierBuilder<'a> {
83    /// Creates a new [`InvertedIndexApplierBuilder`].
84    pub fn new(
85        table_dir: String,
86        path_type: PathType,
87        object_store: ObjectStore,
88        metadata: &'a RegionMetadata,
89        indexed_column_ids: HashSet<ColumnId>,
90        puffin_manager_factory: PuffinManagerFactory,
91    ) -> Self {
92        Self {
93            table_dir,
94            path_type,
95            object_store,
96            metadata,
97            indexed_column_ids,
98            output: BTreeMap::default(),
99            puffin_manager_factory,
100            file_cache: None,
101            inverted_index_cache: None,
102            puffin_metadata_cache: None,
103        }
104    }
105
106    /// Sets the file cache.
107    pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
108        self.file_cache = file_cache;
109        self
110    }
111
112    /// Sets the puffin metadata cache.
113    pub fn with_puffin_metadata_cache(
114        mut self,
115        puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
116    ) -> Self {
117        self.puffin_metadata_cache = puffin_metadata_cache;
118        self
119    }
120
121    /// Sets the inverted index cache.
122    pub fn with_inverted_index_cache(
123        mut self,
124        inverted_index_cache: Option<InvertedIndexCacheRef>,
125    ) -> Self {
126        self.inverted_index_cache = inverted_index_cache;
127        self
128    }
129
130    /// Consumes the builder to construct an [`InvertedIndexApplier`], optionally returned based on
131    /// the expressions provided. If no predicates match, returns `None`.
132    pub fn build(mut self, exprs: &[Expr]) -> Result<Option<InvertedIndexApplier>> {
133        for expr in exprs {
134            self.traverse_and_collect(expr);
135        }
136
137        if self.output.is_empty() {
138            return Ok(None);
139        }
140
141        let expected_predicate_target_types = self.expected_predicate_target_types();
142
143        Ok(Some(
144            InvertedIndexApplier::new(
145                self.table_dir,
146                self.path_type,
147                self.object_store,
148                self.puffin_manager_factory,
149                self.output,
150                expected_predicate_target_types,
151            )?
152            .with_file_cache(self.file_cache)
153            .with_puffin_metadata_cache(self.puffin_metadata_cache)
154            .with_index_cache(self.inverted_index_cache),
155        ))
156    }
157
158    /// Returns `(target, data_type)` pairs for predicate targets
159    /// collected in `self.output`.
160    ///
161    /// The data types are resolved from the latest region manifest. Columns
162    /// that no longer exist in the latest metadata are skipped.
163    fn expected_predicate_target_types(&self) -> BTreeMap<IndexTarget, ConcreteDataType> {
164        self.output
165            .keys()
166            .filter_map(|target| {
167                let data_type = index_target_data_type(self.metadata, target)?;
168                Some((target.clone(), data_type))
169            })
170            .collect()
171    }
172
173    /// Recursively traverses expressions to collect predicates.
174    /// Results are stored in `self.output`.
175    fn traverse_and_collect(&mut self, expr: &Expr) {
176        let res = match expr {
177            Expr::Between(between) => self.collect_between(between),
178
179            Expr::InList(in_list) => self.collect_inlist(in_list),
180            Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
181                Operator::And => {
182                    self.traverse_and_collect(left);
183                    self.traverse_and_collect(right);
184                    Ok(())
185                }
186                Operator::Or => self.collect_or_eq_list(left, right),
187                Operator::Eq => self.collect_eq(left, right),
188                Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => {
189                    self.collect_comparison_expr(left, op, right)
190                }
191                Operator::RegexMatch => self.collect_regex_match(left, right),
192                _ => Ok(()),
193            },
194
195            // TODO(zhongzc): support more expressions, e.g. IsNull, IsNotNull, ...
196            _ => Ok(()),
197        };
198
199        if let Err(err) = res {
200            warn!(err; "Failed to collect predicates, ignore it. expr: {expr}");
201        }
202    }
203
204    /// Helper function to add a predicate to the output.
205    fn add_predicate(&mut self, target: IndexTarget, predicate: Predicate) {
206        self.output.entry(target).or_default().push(predicate);
207    }
208
209    /// Helper function to get the column id and the column type of a column.
210    /// Returns `None` if the column is not a tag column or if the column is ignored.
211    fn column_id_and_type(
212        &self,
213        column_name: &str,
214    ) -> Result<Option<(ColumnId, ConcreteDataType)>> {
215        let column = self
216            .metadata
217            .column_by_name(column_name)
218            .context(ColumnNotFoundSnafu {
219                column: column_name,
220            })?;
221
222        if !self.indexed_column_ids.contains(&column.column_id) {
223            return Ok(None);
224        }
225
226        Ok(Some((
227            column.column_id,
228            column.column_schema.data_type.clone(),
229        )))
230    }
231
232    /// Resolves an expression to an inverted-index target and its value type.
233    ///
234    /// Plain column references resolve only if the column is indexed.
235    /// JSON accesses require a matching path type hint with inverted indexing enabled.
236    /// Two-argument `json_get` uses the hint's type. With three arguments, the
237    /// third argument must be a typed NULL matching the hint's type; different
238    /// string representations are considered compatible.
239    fn expr_to_index_target(&self, expr: &Expr) -> Result<Option<(IndexTarget, ConcreteDataType)>> {
240        if let Some(name) = Self::column_name(expr) {
241            return Ok(self
242                .column_id_and_type(name)?
243                .map(|(id, ty)| (IndexTarget::ColumnId(id), ty)));
244        }
245
246        crate::sst::index::column::json_index_target(self.metadata, expr, |hint| {
247            hint.inverted_index
248        })
249    }
250
251    /// Helper function to get a non-null literal.
252    fn nonnull_lit(expr: &Expr) -> Option<&ScalarValue> {
253        match expr {
254            Expr::Literal(lit, _) if !lit.is_null() => Some(lit),
255            _ => None,
256        }
257    }
258
259    /// Helper function to get the column name of a column expression.
260    fn column_name(expr: &Expr) -> Option<&str> {
261        match expr {
262            Expr::Column(column) => Some(&column.name),
263            _ => None,
264        }
265    }
266
267    /// Helper function to encode a literal into bytes.
268    fn encode_lit(lit: &ScalarValue, data_type: ConcreteDataType) -> Result<Vec<u8>> {
269        let value = Value::try_from(lit.clone()).context(ConvertValueSnafu)?;
270        let mut bytes = vec![];
271        let field = SortField::new(data_type);
272        IndexValueCodec::encode_nonnull_value(value.as_value_ref(), &field, &mut bytes)
273            .context(EncodeSnafu)?;
274        Ok(bytes)
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use std::sync::Arc;
281
282    use api::v1::SemanticType;
283    use datafusion_common::Column;
284    use datafusion_expr::{Between, Literal};
285    use datatypes::data_type::ConcreteDataType;
286    use datatypes::schema::ColumnSchema;
287    use index::inverted_index::search::predicate::{
288        Bound, Range, RangePredicate, RegexMatchPredicate,
289    };
290    use object_store::ObjectStore;
291    use object_store::services::Memory;
292    use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
293    use store_api::storage::RegionId;
294
295    use super::*;
296
297    #[test]
298    fn test_json_target_resolution() {
299        use datafusion_expr::col;
300        let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_target_resolution");
301        let metadata = json_metadata();
302        let builder = InvertedIndexApplierBuilder::new(
303            "test".into(),
304            PathType::Bare,
305            test_object_store(),
306            &metadata,
307            HashSet::new(),
308            factory,
309        );
310        for (path, expected) in [
311            ("$.a", vec!["a"]),
312            ("$.\"a.b\"", vec!["a.b"]),
313            ("$.nested.value", vec!["nested", "value"]),
314        ] {
315            let (target, _) = builder
316                .expr_to_index_target(&json_expr(path, None))
317                .unwrap()
318                .unwrap();
319            assert_eq!(
320                target,
321                IndexTarget::new_json_path(1, expected.into_iter().map(String::from).collect())
322                    .unwrap()
323            );
324        }
325        for path in ["$.missing", "$.unindexed", "$.a[0]", "$.*", "$"] {
326            assert!(
327                builder
328                    .expr_to_index_target(&json_expr(path, None))
329                    .unwrap()
330                    .is_none(),
331                "{path}"
332            );
333        }
334        let data_type = ConcreteDataType::int64_datatype();
335        let expected = Some((
336            IndexTarget::new_json_path(1, vec!["n".into()]).unwrap(),
337            data_type,
338        ));
339        for ty in [None, Some(ScalarValue::Int64(None))] {
340            let resolved = builder.expr_to_index_target(&json_expr("$.n", ty)).unwrap();
341            assert_eq!(resolved, expected);
342        }
343        assert!(
344            builder
345                .expr_to_index_target(&json_expr("$.n", Some(ScalarValue::Float64(None))))
346                .unwrap()
347                .is_none()
348        );
349        assert!(
350            builder
351                .expr_to_index_target(&col("attrs"))
352                .unwrap()
353                .is_none()
354        );
355    }
356
357    #[test]
358    fn test_json_plan_checks_sst_type_hints() {
359        use datafusion_expr::{col, lit};
360        use datatypes::extension::json::JsonMetadata;
361        use datatypes::json::{JsonSettings, JsonTypeHint};
362
363        let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_plan_type_hints");
364        let metadata = json_metadata();
365        for include_column in [false, true] {
366            let builder = InvertedIndexApplierBuilder::new(
367                "test".into(),
368                PathType::Bare,
369                test_object_store(),
370                &metadata,
371                HashSet::from([2]),
372                factory.clone(),
373            );
374            let mut exprs = vec![json_expr("$.n", None).eq(lit(42_i64))];
375            if include_column {
376                exprs.push(col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(0), None))));
377            }
378            let applier = builder.build(&exprs).unwrap().unwrap();
379            let plan = applier
380                .plan_for_sst(&Arc::new(metadata.clone()))
381                .unwrap()
382                .unwrap();
383            assert_eq!(plan.predicate_key, applier.default_plan.predicate_key);
384
385            // Changed types, missing paths, and disabled hints cannot use the JSON index.
386            // Compatible ordinary predicates must still be applied.
387            for hints in [
388                vec![JsonTypeHint {
389                    path: vec!["n".into()],
390                    data_type: ConcreteDataType::float64_datatype(),
391                    inverted_index: true,
392                    skipping_index: None,
393                    fulltext_index: None,
394                }],
395                vec![],
396                vec![JsonTypeHint {
397                    path: vec!["n".into()],
398                    data_type: ConcreteDataType::int64_datatype(),
399                    inverted_index: false,
400                    skipping_index: None,
401                    fulltext_index: None,
402                }],
403            ] {
404                let mut sst = metadata.clone();
405                sst.column_metadatas[0].column_schema.with_extension_type(
406                    &Json2ExtensionType::new(Arc::new(JsonMetadata::new(
407                        JsonSettings::try_new(hints, Some(0)).unwrap(),
408                    ))),
409                );
410                let plan = applier.plan_for_sst(&Arc::new(sst)).unwrap();
411                if !include_column {
412                    assert!(plan.is_none());
413                    continue;
414                }
415                let expected = BTreeMap::from([(
416                    IndexTarget::ColumnId(2),
417                    applier.predicates[&IndexTarget::ColumnId(2)].clone(),
418                )]);
419                assert_eq!(
420                    plan.unwrap().predicate_key,
421                    InvertedIndexApplier::build_apply_plan(&expected)
422                        .unwrap()
423                        .predicate_key
424                );
425            }
426        }
427    }
428
429    #[test]
430    fn test_json_predicate_collection() {
431        use datafusion_expr::lit;
432
433        let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_predicate_collection");
434        let metadata = json_metadata();
435        let mut builder = InvertedIndexApplierBuilder::new(
436            "test".into(),
437            PathType::Bare,
438            test_object_store(),
439            &metadata,
440            HashSet::new(),
441            factory,
442        );
443        let numeric_target = IndexTarget::new_json_path(1, vec!["n".into()]).unwrap();
444        let n = json_expr("$.n", Some(ScalarValue::Int64(None)));
445        for expr in [
446            n.clone().eq(lit(42_i64)),
447            lit(40_i64).lt(n.clone()),
448            n.clone().between(lit(1_i64), lit(99_i64)),
449            n.clone().in_list(vec![lit(1_i64), lit(2_i64)], false),
450            n.clone().eq(lit(1_i64)).or(n.eq(lit(2_i64))),
451        ] {
452            builder.output.clear();
453            builder.traverse_and_collect(&expr);
454            assert_eq!(builder.output.len(), 1, "{expr}");
455            assert_eq!(
456                builder.output.get(&numeric_target).unwrap().len(),
457                1,
458                "{expr}"
459            );
460        }
461        builder.output.clear();
462        let expr = json_expr("$.a", None)
463            .eq(lit("x"))
464            .or(json_expr("$.b", None).eq(lit("y")));
465        builder.traverse_and_collect(&expr);
466        assert!(builder.output.is_empty());
467        let string_target = IndexTarget::new_json_path(1, vec!["a".into()]).unwrap();
468        for pattern in [
469            ScalarValue::Utf8(Some("a.*[ip]".into())),
470            ScalarValue::Utf8View(Some("a.*[ip]".into())),
471            ScalarValue::LargeUtf8(Some("a.*[ip]".into())),
472        ] {
473            builder.output.clear();
474            builder
475                .collect_regex_match(&json_expr("$.a", None), &Expr::Literal(pattern, None))
476                .unwrap();
477            assert_eq!(builder.output.len(), 1);
478            assert_eq!(builder.output.get(&string_target).unwrap().len(), 1);
479        }
480    }
481
482    #[test]
483    fn test_collect_and_basic() {
484        let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_collect_and_basic_");
485
486        let metadata = test_region_metadata();
487        let mut builder = InvertedIndexApplierBuilder::new(
488            "test".to_string(),
489            PathType::Bare,
490            test_object_store(),
491            &metadata,
492            HashSet::from_iter([1, 2, 3]),
493            factory,
494        );
495
496        let expr = Expr::BinaryExpr(BinaryExpr {
497            left: Box::new(Expr::BinaryExpr(BinaryExpr {
498                left: Box::new(tag_column()),
499                op: Operator::RegexMatch,
500                right: Box::new(string_lit("bar")),
501            })),
502            op: Operator::And,
503            right: Box::new(Expr::Between(Between {
504                expr: Box::new(tag_column2()),
505                negated: false,
506                low: Box::new(int64_lit(123)),
507                high: Box::new(int64_lit(456)),
508            })),
509        });
510
511        builder.traverse_and_collect(&expr);
512        let predicates = builder
513            .output
514            .get(&index::target::IndexTarget::ColumnId(1))
515            .unwrap();
516        assert_eq!(predicates.len(), 1);
517        assert_eq!(
518            predicates[0],
519            Predicate::RegexMatch(RegexMatchPredicate {
520                pattern: "bar".to_string()
521            })
522        );
523        let predicates = builder
524            .output
525            .get(&index::target::IndexTarget::ColumnId(2))
526            .unwrap();
527        assert_eq!(predicates.len(), 1);
528        assert_eq!(
529            predicates[0],
530            Predicate::Range(RangePredicate {
531                range: Range {
532                    lower: Some(Bound {
533                        inclusive: true,
534                        value: encoded_int64(123),
535                    }),
536                    upper: Some(Bound {
537                        inclusive: true,
538                        value: encoded_int64(456),
539                    }),
540                }
541            })
542        );
543    }
544
545    pub(crate) fn test_region_metadata() -> RegionMetadata {
546        let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5678));
547        builder
548            .push_column_metadata(ColumnMetadata {
549                column_schema: ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
550                semantic_type: SemanticType::Tag,
551                column_id: 1,
552            })
553            .push_column_metadata(ColumnMetadata {
554                column_schema: ColumnSchema::new("b", ConcreteDataType::int64_datatype(), false),
555                semantic_type: SemanticType::Tag,
556                column_id: 2,
557            })
558            .push_column_metadata(ColumnMetadata {
559                column_schema: ColumnSchema::new("c", ConcreteDataType::string_datatype(), false),
560                semantic_type: SemanticType::Field,
561                column_id: 3,
562            })
563            .push_column_metadata(ColumnMetadata {
564                column_schema: ColumnSchema::new(
565                    "d",
566                    ConcreteDataType::timestamp_millisecond_datatype(),
567                    false,
568                ),
569                semantic_type: SemanticType::Timestamp,
570                column_id: 4,
571            })
572            .primary_key(vec![1, 2]);
573        builder.build().unwrap()
574    }
575
576    pub(crate) fn test_object_store() -> ObjectStore {
577        ObjectStore::new(Memory::default()).unwrap()
578    }
579
580    pub(crate) fn tag_column() -> Expr {
581        Expr::Column(Column::from_name("a"))
582    }
583
584    pub(crate) fn tag_column2() -> Expr {
585        Expr::Column(Column::from_name("b"))
586    }
587
588    pub(crate) fn field_column() -> Expr {
589        Expr::Column(Column::from_name("c"))
590    }
591
592    pub(crate) fn nonexistent_column() -> Expr {
593        Expr::Column(Column::from_name("nonexistence"))
594    }
595
596    pub(crate) fn string_lit(s: impl Into<String>) -> Expr {
597        s.into().lit()
598    }
599
600    pub(crate) fn int64_lit(i: impl Into<i64>) -> Expr {
601        i.into().lit()
602    }
603
604    pub(crate) fn encoded_string(s: impl Into<String>) -> Vec<u8> {
605        let mut bytes = vec![];
606        IndexValueCodec::encode_nonnull_value(
607            Value::from(s.into()).as_value_ref(),
608            &SortField::new(ConcreteDataType::string_datatype()),
609            &mut bytes,
610        )
611        .unwrap();
612        bytes
613    }
614
615    pub(crate) fn encoded_int64(s: impl Into<i64>) -> Vec<u8> {
616        let mut bytes = vec![];
617        IndexValueCodec::encode_nonnull_value(
618            Value::from(s.into()).as_value_ref(),
619            &SortField::new(ConcreteDataType::int64_datatype()),
620            &mut bytes,
621        )
622        .unwrap();
623        bytes
624    }
625
626    fn json_expr(path: &str, ty: Option<ScalarValue>) -> Expr {
627        use common_function::scalars::udf::create_udf;
628        use datafusion_expr::expr::ScalarFunction;
629        let mut args = vec![datafusion_expr::col("attrs"), string_lit(path)];
630        if let Some(ty) = ty {
631            args.push(Expr::Literal(ty, None));
632        }
633        Expr::ScalarFunction(ScalarFunction::new_udf(
634            Arc::new(create_udf(Arc::new(JsonGetWithType::default()))),
635            args,
636        ))
637    }
638
639    fn json_metadata() -> RegionMetadata {
640        use datatypes::extension::json::JsonMetadata;
641        use datatypes::json::{JsonSettings, JsonTypeHint};
642        use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
643        let mut column = ColumnSchema::new(
644            "attrs",
645            ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
646            true,
647        );
648        let hints = [
649            (vec!["a"], true),
650            (vec!["b"], true),
651            (vec!["a.b"], true),
652            (vec!["nested", "value"], true),
653            (vec!["unindexed"], false),
654        ]
655        .into_iter()
656        .map(|(path, inverted_index)| JsonTypeHint {
657            path: path.into_iter().map(String::from).collect(),
658            data_type: ConcreteDataType::string_datatype(),
659            inverted_index,
660            skipping_index: None,
661            fulltext_index: None,
662        })
663        .chain(std::iter::once(JsonTypeHint {
664            path: vec!["n".into()],
665            data_type: ConcreteDataType::int64_datatype(),
666            inverted_index: true,
667            skipping_index: None,
668            fulltext_index: None,
669        }))
670        .collect();
671        column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
672            JsonSettings::try_new(hints, Some(0)).unwrap(),
673        ))));
674        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
675        builder
676            .push_column_metadata(ColumnMetadata {
677                column_schema: column,
678                column_id: 1,
679                semantic_type: SemanticType::Field,
680            })
681            .push_column_metadata(ColumnMetadata {
682                column_schema: ColumnSchema::new(
683                    "ts",
684                    ConcreteDataType::timestamp_millisecond_datatype(),
685                    false,
686                ),
687                column_id: 2,
688                semantic_type: SemanticType::Timestamp,
689            });
690        builder.build().unwrap()
691    }
692}