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

query/optimizer/
constant_term.rs
