query/optimizer/
constant_term.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::fmt;
16use std::fmt::Formatter;
17use std::hash::{Hash, Hasher};
18use std::sync::Arc;
19
20use arrow::array::BooleanArray;
21use common_base::term_token::extract_term_probes;
22use common_function::scalars::matches_term::MatchesTermFinder;
23use datafusion::config::ConfigOptions;
24use datafusion::error::Result as DfResult;
25use datafusion::physical_optimizer::PhysicalOptimizerRule;
26use datafusion::physical_plan::ExecutionPlan;
27use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder};
28use datafusion_common::ScalarValue;
29use datafusion_common::tree_node::{Transformed, TreeNode};
30use datafusion_expr::ColumnarValue;
31use datafusion_physical_expr::expressions::Literal;
32use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr};
33use datatypes::arrow_array::string_array_value_at_index;
34
35/// A physical expression that uses a pre-compiled term finder for the `matches_term` function.
36///
37/// This expression optimizes the `matches_term` function by pre-compiling the term
38/// when the term is a constant value. This avoids recompiling the term for each row
39/// during execution.
40#[derive(Debug)]
41pub struct PreCompiledMatchesTermExpr {
42    /// The text column expression to search in
43    text: Arc<dyn PhysicalExpr>,
44    /// The constant term to search for
45    term: String,
46    /// The pre-compiled term finder
47    finder: MatchesTermFinder,
48
49    /// No used but show how index tokenizes the term basically.
50    /// Not precise due to column options is unknown but for debugging purpose in most cases it's enough.
51    probes: Vec<String>,
52}
53
54impl fmt::Display for PreCompiledMatchesTermExpr {
55    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
56        write!(
57            f,
58            "MatchesConstTerm({}, term: \"{}\", probes: {:?})",
59            self.text, self.term, self.probes
60        )
61    }
62}
63
64impl Hash for PreCompiledMatchesTermExpr {
65    fn hash<H: Hasher>(&self, state: &mut H) {
66        self.text.hash(state);
67        self.term.hash(state);
68    }
69}
70
71impl PartialEq for PreCompiledMatchesTermExpr {
72    fn eq(&self, other: &Self) -> bool {
73        self.text.eq(&other.text) && self.term.eq(&other.term)
74    }
75}
76
77impl Eq for PreCompiledMatchesTermExpr {}
78
79impl PhysicalExpr for PreCompiledMatchesTermExpr {
80    fn data_type(
81        &self,
82        _input_schema: &arrow_schema::Schema,
83    ) -> datafusion::error::Result<arrow_schema::DataType> {
84        Ok(arrow_schema::DataType::Boolean)
85    }
86
87    fn nullable(&self, input_schema: &arrow_schema::Schema) -> datafusion::error::Result<bool> {
88        self.text.nullable(input_schema)
89    }
90
91    fn evaluate(
92        &self,
93        batch: &common_recordbatch::DfRecordBatch,
94    ) -> datafusion::error::Result<ColumnarValue> {
95        let num_rows = batch.num_rows();
96
97        let text_value = self.text.evaluate(batch)?;
98        let array = text_value.into_array(num_rows)?;
99
100        let mut result = BooleanArray::builder(num_rows);
101        for index in 0..array.len() {
102            match string_array_value_at_index(&array, index) {
103                Some(text) => {
104                    result.append_value(self.finder.find(text));
105                }
106                None => {
107                    result.append_null();
108                }
109            }
110        }
111
112        Ok(ColumnarValue::Array(Arc::new(result.finish())))
113    }
114
115    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
116        vec![&self.text]
117    }
118
119    fn with_new_children(
120        self: Arc<Self>,
121        children: Vec<Arc<dyn PhysicalExpr>>,
122    ) -> datafusion::error::Result<Arc<dyn PhysicalExpr>> {
123        Ok(Arc::new(PreCompiledMatchesTermExpr {
124            text: children[0].clone(),
125            term: self.term.clone(),
126            finder: self.finder.clone(),
127            probes: self.probes.clone(),
128        }))
129    }
130
131    fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
132        write!(f, "{}", self)
133    }
134}
135
136/// Optimizer rule that pre-compiles constant term in `matches_term` function.
137///
138/// This optimizer looks for `matches_term` function calls where the second argument
139/// (the term to match) is a constant value. When found, it replaces the function
140/// call with a specialized `PreCompiledMatchesTermExpr` that uses a pre-compiled
141/// term finder.
142///
143/// Example:
144/// ```sql
145/// -- Before optimization:
146/// matches_term(text_column, 'constant_term')
147///
148/// -- After optimization:
149/// PreCompiledMatchesTermExpr(text_column, 'constant_term')
150/// ```
151///
152/// This optimization improves performance by:
153/// 1. Pre-compiling the term once instead of for each row
154/// 2. Using a specialized expression that avoids function call overhead
155#[derive(Debug)]
156pub struct MatchesConstantTermOptimizer;
157
158impl PhysicalOptimizerRule for MatchesConstantTermOptimizer {
159    fn optimize(
160        &self,
161        plan: Arc<dyn ExecutionPlan>,
162        _config: &ConfigOptions,
163    ) -> DfResult<Arc<dyn ExecutionPlan>> {
164        let res = plan
165            .transform_down(&|plan: Arc<dyn ExecutionPlan>| {
166                if let Some(filter) = plan.downcast_ref::<FilterExec>() {
167                    let pred = filter.predicate().clone();
168                    let new_pred = pred.transform_down(&|expr: Arc<dyn PhysicalExpr>| {
169                        if let Some(func) = expr.downcast_ref::<ScalarFunctionExpr>() {
170                            if !func.name().eq_ignore_ascii_case("matches_term") {
171                                return Ok(Transformed::no(expr));
172                            }
173                            let args = func.args();
174                            if args.len() != 2 {
175                                return Ok(Transformed::no(expr));
176                            }
177
178                            if let Some(lit) = args[1].downcast_ref::<Literal>()
179                                && let ScalarValue::Utf8(Some(term)) = lit.value()
180                            {
181                                let finder = MatchesTermFinder::new(term);
182
183                                // Shown in plans to explain fulltext bloom pruning.
184                                let probes = extract_term_probes(term);
185
186                                let expr = PreCompiledMatchesTermExpr {
187                                    text: args[0].clone(),
188                                    term: term.clone(),
189                                    finder,
190                                    probes,
191                                };
192
193                                return Ok(Transformed::yes(Arc::new(expr)));
194                            }
195                        }
196
197                        Ok(Transformed::no(expr))
198                    })?;
199
200                    if new_pred.transformed {
201                        // `FilterExecBuilder::from(&FilterExec)` already carries over the
202                        // input, projection, default selectivity, batch size and fetch, so
203                        // only the predicate has to be replaced. Re-applying the projection
204                        // is not only pointless but wrong: `apply_projection` composes with
205                        // the projection copied from the builder.
206                        let exec = FilterExecBuilder::from(filter)
207                            .with_predicate(new_pred.data)
208                            .build()?;
209                        return Ok(Transformed::yes(Arc::new(exec) as _));
210                    }
211                }
212
213                Ok(Transformed::no(plan))
214            })?
215            .data;
216
217        Ok(res)
218    }
219
220    fn name(&self) -> &str {
221        "MatchesConstantTerm"
222    }
223
224    fn schema_check(&self) -> bool {
225        false
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use std::sync::Arc;
232
233    use arrow::array::{ArrayRef, StringArray, StringDictionaryBuilder};
234    use arrow::datatypes::{DataType, Field, Schema, UInt32Type};
235    use arrow::record_batch::RecordBatch;
236    use catalog::RegisterTableRequest;
237    use catalog::memory::MemoryCatalogManager;
238    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
239    use common_function::scalars::matches_term::MatchesTermFunction;
240    use common_function::scalars::udf::create_udf;
241    use datafusion::datasource::memory::MemorySourceConfig;
242    use datafusion::datasource::source::DataSourceExec;
243    use datafusion::physical_optimizer::PhysicalOptimizerRule;
244    use datafusion::physical_plan::filter::FilterExec;
245    use datafusion::physical_plan::get_plan_string;
246    use datafusion_common::{Column, DFSchema};
247    use datafusion_expr::expr::ScalarFunction;
248    use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
249    use datafusion_expr::{Expr, Literal, ScalarUDF};
250    use datafusion_physical_expr::{ScalarFunctionExpr, create_physical_expr};
251    use datatypes::prelude::ConcreteDataType;
252    use datatypes::schema::ColumnSchema;
253    use session::context::QueryContext;
254    use table::metadata::{TableInfoBuilder, TableMetaBuilder};
255    use table::test_util::EmptyTable;
256
257    use super::*;
258    use crate::parser::QueryLanguageParser;
259    use crate::{QueryEngineFactory, QueryEngineRef};
260
261    fn create_test_batch() -> RecordBatch {
262        let schema = Schema::new(vec![Field::new("text", DataType::Utf8, true)]);
263
264        let text_array = StringArray::from(vec![
265            Some("hello world"),
266            Some("greeting"),
267            Some("hello there"),
268            None,
269        ]);
270
271        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(text_array) as ArrayRef]).unwrap()
272    }
273
274    fn create_test_engine() -> QueryEngineRef {
275        let table_name = "test".to_string();
276        let columns = vec![
277            ColumnSchema::new(
278                "text".to_string(),
279                ConcreteDataType::string_datatype(),
280                false,
281            ),
282            ColumnSchema::new(
283                "timestamp".to_string(),
284                ConcreteDataType::timestamp_millisecond_datatype(),
285                false,
286            )
287            .with_time_index(true),
288        ];
289
290        let schema = Arc::new(datatypes::schema::Schema::new(columns));
291        let table_meta = TableMetaBuilder::empty()
292            .schema(schema)
293            .primary_key_indices(vec![])
294            .value_indices(vec![0])
295            .next_column_id(2)
296            .build()
297            .unwrap();
298        let table_info = TableInfoBuilder::default()
299            .name(&table_name)
300            .meta(table_meta)
301            .build()
302            .unwrap();
303        let table = EmptyTable::from_table_info(&table_info);
304        let catalog_list = MemoryCatalogManager::with_default_setup();
305        assert!(
306            catalog_list
307                .register_table_sync(RegisterTableRequest {
308                    catalog: DEFAULT_CATALOG_NAME.to_string(),
309                    schema: DEFAULT_SCHEMA_NAME.to_string(),
310                    table_name,
311                    table_id: 1024,
312                    table,
313                })
314                .is_ok()
315        );
316        QueryEngineFactory::new(
317            catalog_list,
318            None,
319            None,
320            None,
321            None,
322            false,
323            Default::default(),
324        )
325        .query_engine()
326    }
327
328    fn matches_term_udf() -> Arc<ScalarUDF> {
329        Arc::new(create_udf(Arc::new(MatchesTermFunction::default())))
330    }
331
332    #[test]
333    fn test_matches_term_optimization() {
334        let batch = create_test_batch();
335
336        // Create a predicate with a constant pattern
337        let predicate = create_physical_expr(
338            &Expr::ScalarFunction(ScalarFunction::new_udf(
339                matches_term_udf(),
340                vec![Expr::Column(Column::from_name("text")), "hello".lit()],
341            )),
342            &DFSchema::try_from(batch.schema().clone()).unwrap(),
343            &Default::default(),
344            &PhysicalPlanningContext::default(),
345        )
346        .unwrap();
347
348        let input = DataSourceExec::from_data_source(
349            MemorySourceConfig::try_new(&[vec![batch.clone()]], batch.schema(), None).unwrap(),
350        );
351        let filter = FilterExec::try_new(predicate, input).unwrap();
352
353        // Apply the optimizer
354        let optimizer = MatchesConstantTermOptimizer;
355        let optimized_plan = optimizer
356            .optimize(Arc::new(filter), &Default::default())
357            .unwrap();
358
359        let optimized_filter = optimized_plan.downcast_ref::<FilterExec>().unwrap();
360        let predicate = optimized_filter.predicate();
361
362        // The predicate should be a PreCompiledMatchesTermExpr
363        assert!(predicate.is::<PreCompiledMatchesTermExpr>());
364    }
365
366    #[test]
367    fn test_precompiled_matches_term_with_dictionary() {
368        let mut text = StringDictionaryBuilder::<UInt32Type>::new();
369        text.append_value("hello world");
370        text.append_value("greeting");
371        text.append_value("hello there");
372        text.append_null();
373        let schema = Arc::new(Schema::new(vec![Field::new(
374            "text",
375            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
376            true,
377        )]));
378        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(text.finish())]).unwrap();
379        let expr = PreCompiledMatchesTermExpr {
380            text: Arc::new(datafusion_physical_expr::expressions::Column::new(
381                "text", 0,
382            )),
383            term: "hello".to_string(),
384            finder: MatchesTermFinder::new("hello"),
385            probes: vec!["hello".to_string()],
386        };
387
388        let result = expr.evaluate(&batch).unwrap().into_array(4).unwrap();
389        assert_eq!(
390            result.as_any().downcast_ref::<BooleanArray>().unwrap(),
391            &BooleanArray::from(vec![Some(true), Some(false), Some(true), None])
392        );
393    }
394
395    #[test]
396    fn test_matches_term_no_optimization() {
397        let batch = create_test_batch();
398
399        // Create a predicate with a non-constant pattern
400        let predicate = create_physical_expr(
401            &Expr::ScalarFunction(ScalarFunction::new_udf(
402                matches_term_udf(),
403                vec![
404                    Expr::Column(Column::from_name("text")),
405                    Expr::Column(Column::from_name("text")),
406                ],
407            )),
408            &DFSchema::try_from(batch.schema().clone()).unwrap(),
409            &Default::default(),
410            &PhysicalPlanningContext::default(),
411        )
412        .unwrap();
413
414        let input = DataSourceExec::from_data_source(
415            MemorySourceConfig::try_new(&[vec![batch.clone()]], batch.schema(), None).unwrap(),
416        );
417        let filter = FilterExec::try_new(predicate, input).unwrap();
418
419        let optimizer = MatchesConstantTermOptimizer;
420        let optimized_plan = optimizer
421            .optimize(Arc::new(filter), &Default::default())
422            .unwrap();
423
424        let optimized_filter = optimized_plan.downcast_ref::<FilterExec>().unwrap();
425        let predicate = optimized_filter.predicate();
426
427        // The predicate should still be a ScalarFunctionExpr
428        assert!(predicate.is::<ScalarFunctionExpr>());
429    }
430
431    #[tokio::test]
432    async fn test_matches_term_optimization_from_sql() {
433        let sql = "WITH base AS (
434        SELECT text, timestamp FROM test 
435        WHERE MATCHES_TERM(text, 'hello wo_rld') 
436        AND timestamp > '2025-01-01 00:00:00'
437    ),
438    subquery1 AS (
439        SELECT * FROM base 
440        WHERE MATCHES_TERM(text, 'world')
441    ),
442    subquery2 AS (
443        SELECT * FROM test 
444        WHERE MATCHES_TERM(text, 'greeting') 
445        AND timestamp < '2025-01-02 00:00:00'
446    ),
447    union_result AS (
448        SELECT * FROM subquery1 
449        UNION ALL 
450        SELECT * FROM subquery2
451    ),
452    joined_data AS (
453        SELECT a.text, a.timestamp, b.text as other_text 
454        FROM union_result a 
455        JOIN test b ON a.timestamp = b.timestamp 
456        WHERE MATCHES_TERM(a.text, 'there')
457    )
458    SELECT text, other_text 
459    FROM joined_data 
460    WHERE MATCHES_TERM(text, '42') 
461    AND MATCHES_TERM(other_text, 'foo')";
462
463        let query_ctx = QueryContext::arc();
464
465        let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
466        let engine = create_test_engine();
467        let logical_plan = engine
468            .planner()
469            .plan(&stmt, query_ctx.clone())
470            .await
471            .unwrap();
472
473        let engine_ctx = engine.engine_context(query_ctx);
474        let state = engine_ctx.state();
475
476        let analyzed_plan = state
477            .analyzer()
478            .execute_and_check(logical_plan.clone(), state.config_options(), |_, _| {})
479            .unwrap();
480
481        let optimized_plan = state
482            .optimizer()
483            .optimize(analyzed_plan, state, |_, _| {})
484            .unwrap();
485
486        let physical_plan = state
487            .query_planner()
488            .create_physical_plan(&optimized_plan, state)
489            .await
490            .unwrap();
491
492        let plan_str = get_plan_string(&physical_plan).join("\n");
493        assert!(plan_str.contains("MatchesConstTerm(text@0, term: \"foo\", probes: [\"foo\"]"));
494        assert!(plan_str.contains(
495            "MatchesConstTerm(text@0, term: \"hello wo_rld\", probes: [\"hello\", \"wo\", \"rld\"]"
496        ));
497        assert!(plan_str.contains("MatchesConstTerm(text@0, term: \"world\", probes: [\"world\"]"));
498        assert!(
499            plan_str
500                .contains("MatchesConstTerm(text@0, term: \"greeting\", probes: [\"greeting\"]")
501        );
502        assert!(plan_str.contains("MatchesConstTerm(text@0, term: \"there\", probes: [\"there\"]"));
503        assert!(plan_str.contains("MatchesConstTerm(text@0, term: \"42\", probes: [\"42\"]"));
504        assert!(!plan_str.contains("matches_term"))
505    }
506}
query/optimizer/constant_term.rs

query/optimizer/
constant_term.rs
