mito2/sst/index/fulltext_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
15use std::collections::BTreeMap;
16
17use common_base::term_token::extract_like_probes;
18use datafusion_common::ScalarValue;
19use datafusion_expr::expr::{Like, ScalarFunction};
20use datafusion_expr::{BinaryExpr, Expr, Operator};
21use datatypes::schema::FulltextBackend;
22use object_store::ObjectStore;
23use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
24use store_api::metadata::RegionMetadata;
25use store_api::region_request::PathType;
26use store_api::storage::{ColumnId, ConcreteDataType};
27
28use crate::cache::file_cache::FileCacheRef;
29use crate::cache::index::bloom_filter_index::BloomFilterIndexCacheRef;
30use crate::error::Result;
31use crate::sst::index::fulltext_index::applier::FulltextIndexApplier;
32use crate::sst::index::puffin_manager::PuffinManagerFactory;
33
34/// A request for fulltext index.
35///
36/// It contains all the queries and terms for a column.
37#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
38pub struct FulltextRequest {
39    pub queries: Vec<FulltextQuery>,
40    pub terms: Vec<FulltextTerm>,
41    pub likes: Vec<FulltextLike>,
42}
43
44/// A query to be matched in fulltext index.
45///
46/// `query` is the query to be matched, e.g. "+foo -bar" in `SELECT * FROM t WHERE matches(text, "+foo -bar")`.
47#[derive(Debug, Clone, PartialEq, Eq, Hash)]
48pub struct FulltextQuery(pub String);
49
50/// A term to be matched in fulltext index.
51///
52/// `term` is the term to be matched, e.g. "foo" in `SELECT * FROM t WHERE matches_term(text, "foo")`.
53/// `col_lowered` indicates whether the column is lowercased, e.g. `col_lowered = true` when `matches_term(lower(text), "foo")`.
54#[derive(Debug, Clone, PartialEq, Eq, Hash)]
55pub struct FulltextTerm {
56    pub col_lowered: bool,
57    pub term: String,
58}
59
60/// A case-sensitive `LIKE` pattern to be matched in fulltext index.
61///
62/// `pattern` is e.g. "%foo bar%" in `SELECT * FROM t WHERE text LIKE '%foo bar%'`.
63/// `col_lowered` indicates whether the column is lowercased, e.g. `lower(text) LIKE '%foo bar%'`.
64#[derive(Debug, Clone, PartialEq, Eq, Hash)]
65pub struct FulltextLike {
66    pub col_lowered: bool,
67    pub pattern: String,
68}
69
70/// `FulltextIndexApplierBuilder` is a builder for `FulltextIndexApplier`.
71pub struct FulltextIndexApplierBuilder<'a> {
72    table_dir: String,
73    path_type: PathType,
74    store: ObjectStore,
75    puffin_manager_factory: PuffinManagerFactory,
76    metadata: &'a RegionMetadata,
77    file_cache: Option<FileCacheRef>,
78    puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
79    bloom_filter_cache: Option<BloomFilterIndexCacheRef>,
80}
81
82impl<'a> FulltextIndexApplierBuilder<'a> {
83    /// Creates a new `FulltextIndexApplierBuilder`.
84    pub fn new(
85        table_dir: String,
86        path_type: PathType,
87        store: ObjectStore,
88        puffin_manager_factory: PuffinManagerFactory,
89        metadata: &'a RegionMetadata,
90    ) -> Self {
91        Self {
92            table_dir,
93            path_type,
94            store,
95            puffin_manager_factory,
96            metadata,
97            file_cache: None,
98            puffin_metadata_cache: None,
99            bloom_filter_cache: None,
100        }
101    }
102
103    /// Sets the file cache to be used by the `FulltextIndexApplier`.
104    pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
105        self.file_cache = file_cache;
106        self
107    }
108
109    /// Sets the puffin metadata cache to be used by the `FulltextIndexApplier`.
110    pub fn with_puffin_metadata_cache(
111        mut self,
112        puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
113    ) -> Self {
114        self.puffin_metadata_cache = puffin_metadata_cache;
115        self
116    }
117
118    /// Sets the bloom filter cache to be used by the `FulltextIndexApplier`.
119    pub fn with_bloom_filter_cache(
120        mut self,
121        bloom_filter_cache: Option<BloomFilterIndexCacheRef>,
122    ) -> Self {
123        self.bloom_filter_cache = bloom_filter_cache;
124        self
125    }
126
127    /// Builds `SstIndexApplier` from the given expressions.
128    pub fn build(self, exprs: &[Expr]) -> Result<Option<FulltextIndexApplier>> {
129        let mut requests = BTreeMap::new();
130        for expr in exprs {
131            Self::extract_requests(expr, self.metadata, &mut requests);
132        }
133
134        // Check if any requests have queries or terms
135        let has_requests = requests.iter().any(|(_, request)| {
136            !request.queries.is_empty() || !request.terms.is_empty() || !request.likes.is_empty()
137        });
138
139        Ok(has_requests.then(|| {
140            FulltextIndexApplier::new(
141                self.table_dir,
142                self.path_type,
143                self.store,
144                requests,
145                self.puffin_manager_factory,
146            )
147            .with_file_cache(self.file_cache)
148            .with_puffin_metadata_cache(self.puffin_metadata_cache)
149            .with_bloom_filter_cache(self.bloom_filter_cache)
150        }))
151    }
152
153    fn extract_requests(
154        expr: &Expr,
155        metadata: &'a RegionMetadata,
156        requests: &mut BTreeMap<ColumnId, FulltextRequest>,
157    ) {
158        match expr {
159            Expr::BinaryExpr(BinaryExpr {
160                left,
161                op: Operator::And,
162                right,
163            }) => {
164                Self::extract_requests(left, metadata, requests);
165                Self::extract_requests(right, metadata, requests);
166            }
167            Expr::Like(like) => {
168                if let Some((column_id, like)) = Self::expr_to_like(metadata, like) {
169                    requests.entry(column_id).or_default().likes.push(like);
170                }
171            }
172            Expr::ScalarFunction(func) => {
173                if let Some((column_id, query)) = Self::expr_to_query(metadata, func) {
174                    requests.entry(column_id).or_default().queries.push(query);
175                } else if let Some((column_id, term)) = Self::expr_to_term(metadata, func) {
176                    requests.entry(column_id).or_default().terms.push(term);
177                }
178            }
179            _ => {}
180        }
181    }
182
183    fn expr_to_query(
184        metadata: &RegionMetadata,
185        f: &ScalarFunction,
186    ) -> Option<(ColumnId, FulltextQuery)> {
187        if f.name() != "matches" {
188            return None;
189        }
190        if f.args.len() != 2 {
191            return None;
192        }
193
194        let Expr::Column(c) = &f.args[0] else {
195            return None;
196        };
197        let column = metadata.column_by_name(&c.name)?;
198
199        if column.column_schema.data_type != ConcreteDataType::string_datatype() {
200            return None;
201        }
202
203        let Expr::Literal(ScalarValue::Utf8(Some(query)), _) = &f.args[1] else {
204            return None;
205        };
206
207        Some((column.column_id, FulltextQuery(query.clone())))
208    }
209
210    fn expr_to_term(
211        metadata: &RegionMetadata,
212        f: &ScalarFunction,
213    ) -> Option<(ColumnId, FulltextTerm)> {
214        if f.name() != "matches_term" {
215            return None;
216        }
217        if f.args.len() != 2 {
218            return None;
219        }
220
221        let mut lowered = false;
222        let column;
223        match &f.args[0] {
224            Expr::Column(c) => {
225                column = c;
226            }
227            Expr::ScalarFunction(f) => {
228                let lower_arg = Self::extract_lower_arg(f)?;
229                lowered = true;
230                if let Expr::Column(c) = lower_arg {
231                    column = c;
232                } else {
233                    return None;
234                }
235            }
236            _ => return None,
237        }
238
239        let column = metadata.column_by_name(&column.name)?;
240        if column.column_schema.data_type != ConcreteDataType::string_datatype() {
241            return None;
242        }
243
244        let Expr::Literal(ScalarValue::Utf8(Some(term)), _) = &f.args[1] else {
245            return None;
246        };
247
248        Some((
249            column.column_id,
250            FulltextTerm {
251                col_lowered: lowered,
252                term: term.clone(),
253            },
254        ))
255    }
256
257    fn expr_to_like(metadata: &RegionMetadata, like: &Like) -> Option<(ColumnId, FulltextLike)> {
258        // `ILIKE` uses Unicode case folding, which disagrees with the index's
259        // `to_lowercase` on characters like 'ſ', so its probes could miss rows.
260        if like.negated || like.case_insensitive {
261            return None;
262        }
263        // Probes assume arrow's default `\` escape.
264        if !matches!(like.escape_char, None | Some('\\')) {
265            return None;
266        }
267
268        let (c, col_lowered) = match like.expr.as_ref() {
269            Expr::Column(c) => (c, false),
270            Expr::ScalarFunction(f) => match Self::extract_lower_arg(f)? {
271                Expr::Column(c) => (c, true),
272                _ => return None,
273            },
274            _ => return None,
275        };
276        let column = metadata.column_by_name(&c.name)?;
277        if column.column_schema.data_type != ConcreteDataType::string_datatype() {
278            return None;
279        }
280        // `LIKE` is common on unindexed columns and with patterns like '%word%' that have
281        // no probe; skip them instead of opening every SST's index for nothing.
282        let options = column.column_schema.fulltext_options().ok()??;
283        if !options.enable || options.backend != FulltextBackend::Bloom {
284            return None;
285        }
286
287        let Expr::Literal(
288            ScalarValue::Utf8(Some(pattern))
289            | ScalarValue::LargeUtf8(Some(pattern))
290            | ScalarValue::Utf8View(Some(pattern)),
291            _,
292        ) = like.pattern.as_ref()
293        else {
294            return None;
295        };
296        if extract_like_probes(pattern).is_empty() {
297            return None;
298        }
299
300        Some((
301            column.column_id,
302            FulltextLike {
303                col_lowered,
304                pattern: pattern.clone(),
305            },
306        ))
307    }
308
309    fn extract_lower_arg(lower_func: &ScalarFunction) -> Option<&Expr> {
310        if lower_func.args.len() != 1 {
311            return None;
312        }
313
314        if lower_func.name() != "lower" {
315            return None;
316        }
317
318        if lower_func.args.len() != 1 {
319            return None;
320        }
321
322        Some(&lower_func.args[0])
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use std::sync::Arc;
329
330    use api::v1::SemanticType;
331    use common_function::function::FunctionRef;
332    use common_function::function_factory::ScalarFunctionFactory;
333    use common_function::scalars::matches::MatchesFunction;
334    use common_function::scalars::matches_term::MatchesTermFunction;
335    use datafusion::functions::string::lower;
336    use datafusion_common::Column;
337    use datafusion_expr::expr::{Like, ScalarFunction};
338    use datafusion_expr::{Literal, ScalarUDF};
339    use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextOptions};
340    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
341    use store_api::storage::RegionId;
342
343    use super::*;
344
345    fn mock_metadata() -> RegionMetadata {
346        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
347        builder
348            .push_column_metadata(ColumnMetadata {
349                column_schema: ColumnSchema::new("text", ConcreteDataType::string_datatype(), true),
350                semantic_type: SemanticType::Field,
351                column_id: 1,
352            })
353            .push_column_metadata(ColumnMetadata {
354                column_schema: ColumnSchema::new(
355                    "ts",
356                    ConcreteDataType::timestamp_millisecond_datatype(),
357                    false,
358                ),
359                semantic_type: SemanticType::Timestamp,
360                column_id: 2,
361            });
362
363        builder.build().unwrap()
364    }
365
366    fn matches_func() -> Arc<ScalarUDF> {
367        Arc::new(
368            ScalarFunctionFactory::from(Arc::new(MatchesFunction::default()) as FunctionRef)
369                .provide(Default::default()),
370        )
371    }
372
373    fn matches_term_func() -> Arc<ScalarUDF> {
374        Arc::new(
375            ScalarFunctionFactory::from(Arc::new(MatchesTermFunction::default()) as FunctionRef)
376                .provide(Default::default()),
377        )
378    }
379
380    #[test]
381    fn test_expr_to_query_basic() {
382        let metadata = mock_metadata();
383
384        let func = ScalarFunction {
385            args: vec![Expr::Column(Column::from_name("text")), "foo".lit()],
386            func: matches_func(),
387        };
388
389        let (column_id, query) =
390            FulltextIndexApplierBuilder::expr_to_query(&metadata, &func).unwrap();
391        assert_eq!(column_id, 1);
392        assert_eq!(query, FulltextQuery("foo".to_string()));
393    }
394
395    #[test]
396    fn test_expr_to_query_wrong_num_args() {
397        let metadata = mock_metadata();
398
399        let func = ScalarFunction {
400            args: vec![Expr::Column(Column::from_name("text"))],
401            func: matches_func(),
402        };
403
404        assert!(FulltextIndexApplierBuilder::expr_to_query(&metadata, &func).is_none());
405    }
406
407    #[test]
408    fn test_expr_to_query_not_found_column() {
409        let metadata = mock_metadata();
410
411        let func = ScalarFunction {
412            args: vec![Expr::Column(Column::from_name("not_found")), "foo".lit()],
413            func: matches_func(),
414        };
415
416        assert!(FulltextIndexApplierBuilder::expr_to_query(&metadata, &func).is_none());
417    }
418
419    #[test]
420    fn test_expr_to_query_column_wrong_data_type() {
421        let metadata = mock_metadata();
422
423        let func = ScalarFunction {
424            args: vec![Expr::Column(Column::from_name("ts")), "foo".lit()],
425            func: matches_func(),
426        };
427
428        assert!(FulltextIndexApplierBuilder::expr_to_query(&metadata, &func).is_none());
429    }
430
431    #[test]
432    fn test_expr_to_query_pattern_not_string() {
433        let metadata = mock_metadata();
434
435        let func = ScalarFunction {
436            args: vec![Expr::Column(Column::from_name("text")), 42.lit()],
437            func: matches_func(),
438        };
439
440        assert!(FulltextIndexApplierBuilder::expr_to_query(&metadata, &func).is_none());
441    }
442
443    #[test]
444    fn test_expr_to_term_basic() {
445        let metadata = mock_metadata();
446
447        let func = ScalarFunction {
448            args: vec![Expr::Column(Column::from_name("text")), "foo".lit()],
449            func: matches_term_func(),
450        };
451
452        let (column_id, term) =
453            FulltextIndexApplierBuilder::expr_to_term(&metadata, &func).unwrap();
454        assert_eq!(column_id, 1);
455        assert_eq!(
456            term,
457            FulltextTerm {
458                col_lowered: false,
459                term: "foo".to_string(),
460            }
461        );
462    }
463
464    #[test]
465    fn test_expr_to_term_with_lower() {
466        let metadata = mock_metadata();
467
468        let lower_func_expr = ScalarFunction {
469            args: vec![Expr::Column(Column::from_name("text"))],
470            func: lower(),
471        };
472
473        let func = ScalarFunction {
474            args: vec![Expr::ScalarFunction(lower_func_expr), "foo".lit()],
475            func: matches_term_func(),
476        };
477
478        let (column_id, term) =
479            FulltextIndexApplierBuilder::expr_to_term(&metadata, &func).unwrap();
480        assert_eq!(column_id, 1);
481        assert_eq!(
482            term,
483            FulltextTerm {
484                col_lowered: true,
485                term: "foo".to_string(),
486            }
487        );
488    }
489
490    #[test]
491    fn test_expr_to_term_wrong_num_args() {
492        let metadata = mock_metadata();
493
494        let func = ScalarFunction {
495            args: vec![Expr::Column(Column::from_name("text"))],
496            func: matches_term_func(),
497        };
498
499        assert!(FulltextIndexApplierBuilder::expr_to_term(&metadata, &func).is_none());
500    }
501
502    #[test]
503    fn test_expr_to_term_wrong_function_name() {
504        let metadata = mock_metadata();
505
506        let func = ScalarFunction {
507            args: vec![Expr::Column(Column::from_name("text")), "foo".lit()],
508            func: matches_func(), // Using 'matches' instead of 'matches_term'
509        };
510
511        assert!(FulltextIndexApplierBuilder::expr_to_term(&metadata, &func).is_none());
512    }
513
514    #[test]
515    fn test_extract_lower_arg() {
516        let func = ScalarFunction {
517            args: vec![Expr::Column(Column::from_name("text"))],
518            func: lower(),
519        };
520
521        let arg = FulltextIndexApplierBuilder::extract_lower_arg(&func).unwrap();
522        match arg {
523            Expr::Column(c) => {
524                assert_eq!(c.name, "text");
525            }
526            _ => panic!("Expected Column expression"),
527        }
528    }
529
530    #[test]
531    fn test_extract_lower_arg_wrong_function() {
532        let func = ScalarFunction {
533            args: vec![Expr::Column(Column::from_name("text"))],
534            func: matches_func(), // Not 'lower'
535        };
536
537        assert!(FulltextIndexApplierBuilder::extract_lower_arg(&func).is_none());
538    }
539
540    #[test]
541    fn test_extract_requests() {
542        let metadata = mock_metadata();
543
544        // Create a matches expression
545        let matches_expr = Expr::ScalarFunction(ScalarFunction {
546            args: vec![Expr::Column(Column::from_name("text")), "foo".lit()],
547            func: matches_func(),
548        });
549
550        let mut requests = BTreeMap::new();
551        FulltextIndexApplierBuilder::extract_requests(&matches_expr, &metadata, &mut requests);
552
553        assert_eq!(requests.len(), 1);
554        let request = requests.get(&1).unwrap();
555        assert_eq!(request.queries.len(), 1);
556        assert_eq!(request.terms.len(), 0);
557        assert_eq!(request.queries[0], FulltextQuery("foo".to_string()));
558    }
559
560    fn mock_metadata_with_fulltext(backend: FulltextBackend) -> RegionMetadata {
561        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
562        builder
563            .push_column_metadata(ColumnMetadata {
564                column_schema: ColumnSchema::new("text", ConcreteDataType::string_datatype(), true)
565                    .with_fulltext_options(FulltextOptions::new_unchecked(
566                        true,
567                        FulltextAnalyzer::English,
568                        false,
569                        backend,
570                        10240,
571                        0.01,
572                    ))
573                    .unwrap(),
574                semantic_type: SemanticType::Field,
575                column_id: 1,
576            })
577            .push_column_metadata(ColumnMetadata {
578                column_schema: ColumnSchema::new(
579                    "ts",
580                    ConcreteDataType::timestamp_millisecond_datatype(),
581                    false,
582                ),
583                semantic_type: SemanticType::Timestamp,
584                column_id: 2,
585            });
586
587        builder.build().unwrap()
588    }
589
590    #[test]
591    fn test_expr_to_like() {
592        let like = |negated, case_insensitive, escape_char, pattern: &str| Like {
593            negated,
594            expr: Box::new(Expr::Column(Column::from_name("text"))),
595            pattern: Box::new(pattern.lit()),
596            escape_char,
597            case_insensitive,
598        };
599        let bloom = mock_metadata_with_fulltext(FulltextBackend::Bloom);
600
601        for escape_char in [None, Some('\\')] {
602            assert_eq!(
603                FulltextIndexApplierBuilder::expr_to_like(
604                    &bloom,
605                    &like(false, false, escape_char, "%foo bar baz%")
606                ),
607                Some((
608                    1,
609                    FulltextLike {
610                        col_lowered: false,
611                        pattern: "%foo bar baz%".to_string(),
612                    }
613                ))
614            );
615        }
616        let lowered = Like {
617            expr: Box::new(Expr::ScalarFunction(ScalarFunction {
618                args: vec![Expr::Column(Column::from_name("text"))],
619                func: lower(),
620            })),
621            ..like(false, false, None, "%foo bar baz%")
622        };
623        assert_eq!(
624            FulltextIndexApplierBuilder::expr_to_like(&bloom, &lowered),
625            Some((
626                1,
627                FulltextLike {
628                    col_lowered: true,
629                    pattern: "%foo bar baz%".to_string(),
630                }
631            ))
632        );
633        // NOT LIKE, ILIKE and other escape characters can't use the probes, and
634        // '%foo%' has none.
635        for like in [
636            like(true, false, None, "%foo bar baz%"),
637            like(false, true, None, "%foo bar baz%"),
638            like(false, false, Some('!'), "%foo bar baz%"),
639            like(false, false, None, "%foo%"),
640        ] {
641            assert_eq!(
642                FulltextIndexApplierBuilder::expr_to_like(&bloom, &like),
643                None
644            );
645        }
646        // Only columns with a bloom fulltext index.
647        for metadata in [
648            mock_metadata(),
649            mock_metadata_with_fulltext(FulltextBackend::Tantivy),
650        ] {
651            assert_eq!(
652                FulltextIndexApplierBuilder::expr_to_like(
653                    &metadata,
654                    &like(false, false, None, "%foo bar baz%")
655                ),
656                None
657            );
658        }
659    }
660
661    #[test]
662    fn test_extract_multiple_requests() {
663        let metadata = mock_metadata();
664
665        // Create a matches expression
666        let matches_expr = Expr::ScalarFunction(ScalarFunction {
667            args: vec![Expr::Column(Column::from_name("text")), "foo".lit()],
668            func: matches_func(),
669        });
670
671        // Create a matches_term expression
672        let matches_term_expr = Expr::ScalarFunction(ScalarFunction {
673            args: vec![Expr::Column(Column::from_name("text")), "bar".lit()],
674            func: matches_term_func(),
675        });
676
677        // Create a binary expression combining both
678        let binary_expr = Expr::BinaryExpr(BinaryExpr {
679            left: Box::new(matches_expr),
680            op: Operator::And,
681            right: Box::new(matches_term_expr),
682        });
683
684        let mut requests = BTreeMap::new();
685        FulltextIndexApplierBuilder::extract_requests(&binary_expr, &metadata, &mut requests);
686
687        assert_eq!(requests.len(), 1);
688        let request = requests.get(&1).unwrap();
689        assert_eq!(request.queries.len(), 1);
690        assert_eq!(request.terms.len(), 1);
691        assert_eq!(request.queries[0], FulltextQuery("foo".to_string()));
692        assert_eq!(
693            request.terms[0],
694            FulltextTerm {
695                col_lowered: false,
696                term: "bar".to_string(),
697            }
698        );
699    }
700}
mito2/sst/index/fulltext_index/applier/builder.rs

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