From d8349974f500070ff47da4da19ac9f259dcf9cde Mon Sep 17 00:00:00 2001 From: Hengfei Yang Date: Tue, 27 May 2025 19:46:32 +0800 Subject: [PATCH] perf: use tantivy to optimize histogram & count without filter (#6930) --- src/service/db/schema.rs | 6 +- .../datafusion/plan/tantivy_count_exec.rs | 12 +- src/service/search/grpc/flight.rs | 2 +- src/service/search/index.rs | 12 +- src/service/search/sql.rs | 108 ++++++++++++------ .../search/tantivy/puffin_directory/reader.rs | 24 ++-- 6 files changed, 105 insertions(+), 59 deletions(-) diff --git a/src/service/db/schema.rs b/src/service/db/schema.rs index cd04b3cd38..ef13d9ca21 100644 --- a/src/service/db/schema.rs +++ b/src/service/db/schema.rs @@ -508,10 +508,10 @@ pub async fn cache() -> Result<(), anyhow::Error> { let entry = schemas.entry(item_key).or_insert(Vec::new()); entry.push((start_dt, val)); if i % 1000 == 0 { - log::info!("Cache schema group progress: {} / {}", i, items_num); + log::info!("Cache schema progress: {}/{}", i, items_num); } } - log::info!("Cache schema group done"); + log::info!("Stream schemas Cached {} schemas", items_num); let keys_num = schemas.keys().len(); let keys = schemas.keys().map(|k| k.to_string()).collect::>(); for (i, item_key) in keys.iter().enumerate() { @@ -564,7 +564,7 @@ pub async fn cache() -> Result<(), anyhow::Error> { w.insert(item_key.to_string(), schema_versions); drop(w); if i % 1000 == 0 { - log::info!("Stream schemas Cached progress: {} / {}", i, keys.len()); + log::info!("Stream schemas Cached progress: {}/{}", i, keys.len()); } } log::info!("Stream schemas Cached {} streams", keys_num); diff --git a/src/service/search/datafusion/plan/tantivy_count_exec.rs b/src/service/search/datafusion/plan/tantivy_count_exec.rs index 7769340b9a..bd48a655d2 100644 --- a/src/service/search/datafusion/plan/tantivy_count_exec.rs +++ b/src/service/search/datafusion/plan/tantivy_count_exec.rs @@ -40,10 +40,10 @@ use crate::service::search::{ #[derive(Debug)] pub struct TantivyOptimizeExec { query: Arc, - schema: SchemaRef, // The schema for the produced row - file_list: Vec, // The list of files to read - index_condition: IndexCondition, // The condition to filter the rows - cache: PlanProperties, // Cached properties of this plan + schema: SchemaRef, // The schema for the produced row + file_list: Vec, // The list of files to read + index_condition: Option, // The condition to filter the rows + cache: PlanProperties, // Cached properties of this plan index_optimize_mode: InvertedIndexOptimizeMode, // Type of query the ttv index optimizes } @@ -53,7 +53,7 @@ impl TantivyOptimizeExec { query: Arc, schema: SchemaRef, file_list: Vec, - index_condition: IndexCondition, + index_condition: Option, index_optimize_mode: InvertedIndexOptimizeMode, ) -> Self { let cache = Self::compute_properties(Arc::clone(&schema)); @@ -135,7 +135,7 @@ impl ExecutionPlan for TantivyOptimizeExec { let fut = adapt_tantivy_result( self.query.clone(), self.file_list.clone(), - Some(self.index_condition.clone()), + self.index_condition.clone(), self.schema.clone(), self.index_optimize_mode.clone(), ); diff --git a/src/service/search/grpc/flight.rs b/src/service/search/grpc/flight.rs index 00a503a180..c45181b55e 100644 --- a/src/service/search/grpc/flight.rs +++ b/src/service/search/grpc/flight.rs @@ -376,7 +376,7 @@ pub async fn search( query_params, physical_plan.schema(), tantivy_file_list, - index_condition.unwrap(), + index_condition, idx_optimize_rule.unwrap(), // guaranteed Some, if tantivy_file_list is not empty )); physical_plan = Arc::new(UnionExec::new(vec![physical_plan, tantivy_exec as _])); diff --git a/src/service/search/index.rs b/src/service/search/index.rs index d8d907441d..c52864fb96 100644 --- a/src/service/search/index.rs +++ b/src/service/search/index.rs @@ -34,7 +34,10 @@ use serde::{Deserialize, Serialize}; use sqlparser::ast::{BinaryOperator, Expr, FunctionArguments}; use tantivy::{ Term, - query::{BooleanQuery, FuzzyTermQuery, Occur, PhrasePrefixQuery, Query, RegexQuery, TermQuery}, + query::{ + AllQuery, BooleanQuery, FuzzyTermQuery, Occur, PhrasePrefixQuery, Query, RegexQuery, + TermQuery, + }, schema::{Field, IndexRecordOption, Schema}, }; @@ -101,6 +104,7 @@ pub enum Condition { Regex(String, String), MatchAll(String), FuzzyMatchAll(String, u8), + All(), Or(Box, Box), And(Box, Box), } @@ -193,6 +197,7 @@ impl Condition { "{}:fuzzy({}, {})", INDEX_FIELD_NAME_FOR_ALL, value, distance ), + Condition::All() => "ALL".to_string(), Condition::Or(left, right) => format!("({} OR {})", left.to_query(), right.to_query()), Condition::And(left, right) => { format!("({} AND {})", left.to_query(), right.to_query()) @@ -350,6 +355,7 @@ impl Condition { let term = Term::from_field_text(default_field, value); Box::new(FuzzyTermQuery::new(term, *distance, false)) } + Condition::All() => Box::new(AllQuery {}), Condition::Or(left, right) => { let left_query = left.to_tantivy_query(schema, default_field)?; let right_query = right.to_tantivy_query(schema, default_field)?; @@ -381,6 +387,7 @@ impl Condition { Condition::FuzzyMatchAll(..) => { fields.insert(INDEX_FIELD_NAME_FOR_ALL.to_string()); } + Condition::All() => {} Condition::Or(left, right) | Condition::And(left, right) => { fields.extend(left.get_tantivy_fields()); fields.extend(right.get_tantivy_fields()); @@ -407,6 +414,7 @@ impl Condition { Condition::FuzzyMatchAll(..) => { fields.extend(fst_fields.iter().cloned()); } + Condition::All() => {} Condition::Or(left, right) | Condition::And(left, right) => { fields.extend(left.get_schema_fields(fst_fields)); fields.extend(right.get_schema_fields(fst_fields)); @@ -492,6 +500,7 @@ impl Condition { } Ok(disjunction(expr_list)) } + Condition::All() => Ok(Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))), Condition::Or(left, right) => { let left = left.to_physical_expr(schema, fst_fields)?; let right = right.to_physical_expr(schema, fst_fields)?; @@ -512,6 +521,7 @@ impl Condition { Condition::Regex(..) => false, Condition::MatchAll(v) => is_blank_or_alphanumeric(v), Condition::FuzzyMatchAll(..) => false, + Condition::All() => true, Condition::Or(left, right) => left.can_remove_filter() && right.can_remove_filter(), Condition::And(left, right) => left.can_remove_filter() && right.can_remove_filter(), } diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index e5982a9408..95e0ee7141 100644 --- a/src/service/search/sql.rs +++ b/src/service/search/sql.rs @@ -57,7 +57,7 @@ use super::{ FUZZY_MATCH_ALL_UDF_NAME, MATCH_ALL_RAW_IGNORE_CASE_UDF_NAME, MATCH_ALL_RAW_UDF_NAME, MATCH_ALL_UDF_NAME, }, - index::{IndexCondition, get_index_condition_from_expr}, + index::{Condition, IndexCondition, get_index_condition_from_expr}, request::Request, utils::{conjunction, is_field, is_value, split_conjunction, trim_quotes}, }; @@ -266,6 +266,12 @@ impl Sql { index_condition = index_visitor.index_condition; can_optimize = index_visitor.can_optimize; } + // use all condition for histogram without filter + if use_inverted_index && can_optimize && index_condition.is_none() { + index_condition = Some(IndexCondition { + conditions: vec![Condition::All()], + }); + } // 12. check `select * from table where match_all()` optimizer let mut index_optimize_mode = None; @@ -739,6 +745,21 @@ impl VisitorMut for ColumnVisitor<'_> { checking_inverted_index_inner(&index_fields, expr); } } + } else if is_simple_count_query(select) || is_simple_histogram_query(select) { + // if there is no selection, but have histogram and fst_fields, also can use + // inverted index + if self.schemas.len() == 1 { + for (_, schema) in self.schemas.iter() { + let stream_settings = unwrap_stream_settings(schema.schema()); + let fts_fields = get_stream_setting_fts_fields(&stream_settings); + let index_fields = get_stream_setting_index_fields(&stream_settings); + let index_fields = itertools::chain(fts_fields.iter(), index_fields.iter()) + .collect::>(); + if !index_fields.is_empty() { + self.use_inverted_index = true; + } + } + } } } if let Some(Expr::Value(Value::Number(n, _))) = query.limit.as_ref() { @@ -827,6 +848,9 @@ impl VisitorMut for IndexVisitor { if self.is_remove_filter || can_remove_filter { select.selection = other_expr; } + } else if is_simple_count_query(select) || is_simple_histogram_query(select) { + // if there is no selection, but have histogram, also can use inverted index + self.can_optimize = true; } } ControlFlow::Continue(()) @@ -1244,46 +1268,11 @@ impl VisitorMut for OtherIndexOptimizeModeVisitor { return ControlFlow::Break(()); } - // Check if the query is a simple `select sql_func(*)` without modifiers - let is_sql_func = |select: &SelectItem, fn_name: &str, with_star: bool| -> bool { - match select { - SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. } => { - if let Expr::Function(func) = expr { - let name = trim_quotes(&func.name.to_string().to_lowercase()); - // Check function name matches and has no special modifiers - let has_no_modifiers = func.filter.is_none() - && func.over.is_none() - && func.within_group.is_empty(); - - // If with_start is true, check for single "*" argument - let has_valid_args = if with_star { - matches!( - &func.args, - FunctionArguments::List(list) - if list.args.len() == 1 - && trim_quotes(&list.args[0].to_string()) == "*" - ) - } else { - true - }; - - name == fn_name && has_no_modifiers && has_valid_args - } else { - false - } - } - _ => false, - } - }; - if select.projection.len() == 1 && matches!(select.group_by, GroupByExpr::Expressions(ref expr, _) if expr.is_empty()) { - self.is_simple_count = is_sql_func(&select.projection[0], "count", true); - } else if select.projection.len() == 2 - && is_sql_func(&select.projection[0], "histogram", false) - && is_sql_func(&select.projection[1], "count", true) - { + self.is_simple_count = is_simple_count_query(select); + } else if is_simple_histogram_query(select) { self.is_simple_histogram = true; } } @@ -1312,6 +1301,49 @@ fn is_complex_query(statement: &mut Statement) -> bool { visitor.is_complex } +// check if the query is only count(*) query +fn is_simple_count_query(select: &Select) -> bool { + select.projection.len() == 1 && is_sql_func(&select.projection[0], "count", true) +} + +// check if the query is only histogram & count query +fn is_simple_histogram_query(select: &Select) -> bool { + select.projection.len() == 2 + && is_sql_func(&select.projection[0], "histogram", false) + && is_sql_func(&select.projection[1], "count", true) +} + +// Check if the query is a simple `select sql_func(*)` without modifiers +fn is_sql_func(select: &SelectItem, fn_name: &str, with_star: bool) -> bool { + match select { + SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. } => { + if let Expr::Function(func) = expr { + let name = trim_quotes(&func.name.to_string().to_lowercase()); + // Check function name matches and has no special modifiers + let has_no_modifiers = + func.filter.is_none() && func.over.is_none() && func.within_group.is_empty(); + + // If with_start is true, check for single "*" argument + let has_valid_args = if with_star { + matches!( + &func.args, + FunctionArguments::List(list) + if list.args.len() == 1 + && trim_quotes(&list.args[0].to_string()) == "*" + ) + } else { + true + }; + + name == fn_name && has_no_modifiers && has_valid_args + } else { + false + } + } + _ => false, + } +} + // check if the query is complex query // 1. has subquery // 2. has join diff --git a/src/service/search/tantivy/puffin_directory/reader.rs b/src/service/search/tantivy/puffin_directory/reader.rs index e3a8a25ed4..aace21aa25 100644 --- a/src/service/search/tantivy/puffin_directory/reader.rs +++ b/src/service/search/tantivy/puffin_directory/reader.rs @@ -236,20 +236,23 @@ pub async fn warm_up_terms( warm_up_terms_futures .push(async move { inv_idx_clone.warm_postings(term, *position_needed).await }); } + } + } - // only warm up fast fields if needed - if need_fast_field { - // only warm up fast fields once per segment - let segment_id = segment_reader.segment_id(); - if !warmed_segments.contains(&segment_id) { - let fast_field_reader = segment_reader.fast_fields(); - warm_up_fast_fields_futures - .push(async move { warm_up_fastfield(fast_field_reader).await }); - warmed_segments.insert(segment_id); - } + // only warm up fast fields if needed + if need_fast_field { + for segment_reader in searcher.segment_readers() { + // only warm up fast fields once per segment + let segment_id = segment_reader.segment_id(); + if !warmed_segments.contains(&segment_id) { + let fast_field_reader = segment_reader.fast_fields(); + warm_up_fast_fields_futures + .push(async move { warm_up_fastfield(fast_field_reader).await }); + warmed_segments.insert(segment_id); } } } + if !warm_up_fields_futures.is_empty() { try_join_all(warm_up_fields_futures).await?; } @@ -265,6 +268,7 @@ pub async fn warm_up_terms( Ok(()) } +// warm up the fast field, only support _timestamp field async fn warm_up_fastfield( fast_field_reader: &tantivy::fastfield::FastFieldReaders, ) -> anyhow::Result<()> {