From ebe4be0e61540f39e10448bcf77bb7fb142c429d Mon Sep 17 00:00:00 2001 From: Huaijin Date: Tue, 1 Jul 2025 14:22:30 +0800 Subject: [PATCH] feat: use StrinView to speed up query for main (#7400) Tow new env ``` ZO_SEARCH_INSPECTOR_ENABLED = false ZO_UTF8_VIEW_ENABLED = true ``` --- src/config/src/config.rs | 4 + .../datafusion/optimizer/rewrite_match.rs | 88 ++++++++++++++----- src/service/search/datafusion/udf/mod.rs | 10 +-- src/service/search/grpc/mod.rs | 1 + src/service/search/grpc/storage.rs | 7 +- src/service/search/grpc/utils.rs | 38 ++++++++ src/service/search/grpc/wal.rs | 26 +++++- src/service/search/index.rs | 19 +++- src/service/search/inspector.rs | 12 ++- src/service/search/sql.rs | 30 ++++++- 10 files changed, 191 insertions(+), 44 deletions(-) create mode 100644 src/service/search/grpc/utils.rs diff --git a/src/config/src/config.rs b/src/config/src/config.rs index a5c6ef334f..a968ad155d 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -1143,6 +1143,10 @@ pub struct Common { pub use_stream_settings_for_partitions_enabled: bool, #[env_config(name = "ZO_DASHBOARD_PLACEHOLDER", default = "_o2_all_")] pub dashboard_placeholder: String, + #[env_config(name = "ZO_SEARCH_INSPECTOR_ENABLED", default = false)] + pub search_inspector_enabled: bool, + #[env_config(name = "ZO_UTF8_VIEW_ENABLED", default = true)] + pub utf8_view_enabled: bool, } #[derive(EnvConfig)] diff --git a/src/service/search/datafusion/optimizer/rewrite_match.rs b/src/service/search/datafusion/optimizer/rewrite_match.rs index 2cf8ac6e36..bddd745a86 100644 --- a/src/service/search/datafusion/optimizer/rewrite_match.rs +++ b/src/service/search/datafusion/optimizer/rewrite_match.rs @@ -13,6 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . +use config::get_config; use datafusion::{ self, common::{ @@ -127,7 +128,11 @@ impl TreeNodeRewriter for MatchToFullTextMatch { .trim_start_matches('*') // contains .trim_end_matches('*') // prefix or contains .to_string(); // remove prefix and suffix * - let item = Expr::Literal(ScalarValue::Utf8(Some(format!("%{item}%")))); + let item = if get_config().common.utf8_view_enabled { + Expr::Literal(ScalarValue::Utf8View(Some(format!("%{item}%")))) + } else { + Expr::Literal(ScalarValue::Utf8(Some(format!("%{item}%")))) + }; for field in self.fields.iter() { let new_expr = Expr::Like(Like { negated: false, @@ -159,7 +164,11 @@ impl TreeNodeRewriter for MatchToFullTextMatch { ))); }; let mut expr_list = Vec::with_capacity(self.fields.len()); - let item = Expr::Literal(ScalarValue::Utf8(Some(item.to_string()))); + let item = if get_config().common.utf8_view_enabled { + Expr::Literal(ScalarValue::Utf8View(Some(item.to_string()))) + } else { + Expr::Literal(ScalarValue::Utf8(Some(item.to_string()))) + }; let distance = Expr::Literal(ScalarValue::Int64(Some(distance))); let fuzzy_expr = fuzzy_match_udf::FUZZY_MATCH_UDF.clone(); for field in self.fields.iter() { @@ -190,8 +199,9 @@ impl TreeNodeRewriter for MatchToFullTextMatch { mod tests { use std::sync::Arc; - use arrow::array::{Int64Array, StringArray}; + use arrow::array::{Array, Int64Array, StringArray, StringViewArray}; use arrow_schema::DataType; + use config::get_config; use datafusion::{ arrow::{ datatypes::{Field, Schema}, @@ -234,32 +244,64 @@ mod tests { ), ]; - // define a schema. - let schema = Arc::new(Schema::new(vec![ - Field::new("_timestamp", DataType::Int64, false), - Field::new("name", DataType::Utf8, false), - Field::new("log", DataType::Utf8, false), - ])); + let schema = if get_config().common.utf8_view_enabled { + // define a schema. + Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8View, false), + Field::new("log", DataType::Utf8View, false), + ])) + } else { + Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("log", DataType::Utf8, false), + ])) + }; + + let name_array: Arc = if get_config().common.utf8_view_enabled { + Arc::new(StringViewArray::from(vec![ + "open", + "observe", + "openobserve", + "OBserve", + "oo", + ])) + } else { + Arc::new(StringArray::from(vec![ + "open", + "observe", + "openobserve", + "OBserve", + "oo", + ])) + }; + + let log_array: Arc = if get_config().common.utf8_view_enabled { + Arc::new(StringViewArray::from(vec![ + "o2", + "obSERVE", + "openobserve", + "o2", + "oo", + ])) + } else { + Arc::new(StringArray::from(vec![ + "o2", + "obSERVE", + "openobserve", + "o2", + "oo", + ])) + }; // define data. let batch = RecordBatch::try_new( schema.clone(), vec![ Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), - Arc::new(StringArray::from(vec![ - "open", - "observe", - "openobserve", - "OBserve", - "oo", - ])), - Arc::new(StringArray::from(vec![ - "o2", - "obSERVE", - "openobserve", - "o2", - "oo", - ])), + name_array, + log_array, ], ) .unwrap(); diff --git a/src/service/search/datafusion/udf/mod.rs b/src/service/search/datafusion/udf/mod.rs index 4f52ee3115..99156a1592 100644 --- a/src/service/search/datafusion/udf/mod.rs +++ b/src/service/search/datafusion/udf/mod.rs @@ -55,15 +55,7 @@ pub(crate) const REGEX_NOT_MATCH_UDF_NAME: &str = "re_not_match"; /// The name of the regex_matches UDF given to DataFusion. pub(crate) const REGEX_MATCHES_UDF_NAME: &str = "re_matches"; -pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 13] = [ - ZoFunction { - name: "match_all_raw", - text: "match_all_raw('v')", - }, - ZoFunction { - name: "match_all_raw_ignore_case", - text: "match_all_raw_ignore_case('v')", - }, +pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 11] = [ ZoFunction { name: "match_all", text: "match_all('v')", diff --git a/src/service/search/grpc/mod.rs b/src/service/search/grpc/mod.rs index 9061c4ca9d..33cc50c395 100644 --- a/src/service/search/grpc/mod.rs +++ b/src/service/search/grpc/mod.rs @@ -21,6 +21,7 @@ use infra::errors::Result; pub mod flight; pub mod storage; +pub(crate) mod utils; pub mod wal; pub type SearchTable = Result<(Vec>, ScanStats)>; diff --git a/src/service/search/grpc/storage.rs b/src/service/search/grpc/storage.rs index 1f4e4f3fa9..6c00b90e33 100644 --- a/src/service/search/grpc/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -53,6 +53,7 @@ use crate::service::{ search::{ datafusion::exec, generate_search_schema_diff, + grpc::utils, index::IndexCondition, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, tantivy::puffin_directory::{ @@ -350,8 +351,10 @@ pub async fn search( if files.is_empty() { continue; } - let schema = schema_versions[ver].clone(); - let schema = schema.with_metadata(Default::default()); + let schema = schema_versions[ver] + .clone() + .with_metadata(Default::default()); + let schema = utils::change_schema_to_utf8_view(schema); let session = config::meta::search::Session { id: format!("{}-storage-{ver}", query.trace_id), diff --git a/src/service/search/grpc/utils.rs b/src/service/search/grpc/utils.rs new file mode 100644 index 0000000000..8e883ba91a --- /dev/null +++ b/src/service/search/grpc/utils.rs @@ -0,0 +1,38 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::sync::Arc; + +use arrow_schema::{DataType, Field, Schema}; +use config::get_config; + +pub fn change_schema_to_utf8_view(schema: Schema) -> Schema { + if !get_config().common.utf8_view_enabled { + return schema; + } + + let fields = schema + .fields() + .iter() + .map(|f| { + if f.data_type() == &DataType::Utf8 { + Arc::new(Field::new(f.name(), DataType::Utf8View, f.is_nullable())) + } else { + f.clone() + } + }) + .collect::>(); + Schema::new(fields) +} diff --git a/src/service/search/grpc/wal.rs b/src/service/search/grpc/wal.rs index a4cfab13e0..89f5dafb90 100644 --- a/src/service/search/grpc/wal.rs +++ b/src/service/search/grpc/wal.rs @@ -16,6 +16,7 @@ use std::{path::Path, sync::Arc}; use arrow::array::{ArrayRef, new_null_array}; +use arrow_schema::{DataType, Field}; use config::{ cluster::LOCAL_NODE, get_config, @@ -46,6 +47,7 @@ use crate::{ search::{ datafusion::{exec, table_provider::memtable::NewMemTable}, generate_filter_from_equal_items, generate_search_schema_diff, + grpc::utils, index::IndexCondition, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, match_source, @@ -275,6 +277,8 @@ pub async fn search_parquet( let schema = schema_versions[ver] .clone() .with_metadata(Default::default()); + let schema = utils::change_schema_to_utf8_view(schema); + let session = config::meta::search::Session { id: format!("{}-wal-{ver}", query.trace_id), storage_type: StorageType::Wal, @@ -438,10 +442,13 @@ pub async fn search_memtable( continue; } - let diff_fields = generate_search_schema_diff(&schema, &latest_schema_map); + // the field in latest_schema_map, but not in schema, + // so not in the diff_fields, result in Utf8View issue + // so we need to add it to the diff_fields + let mut diff_fields = generate_search_schema_diff(&schema, &latest_schema_map); let mut adapt_batches = Vec::with_capacity(record_batches.len()); for batch in record_batches { - adapt_batches.push(adapt_batch(latest_schema.clone(), batch)); + adapt_batches.push(adapt_batch(latest_schema.clone(), batch, &mut diff_fields)); } let record_batches = adapt_batches; @@ -644,7 +651,11 @@ async fn get_file_list( .await } -pub fn adapt_batch(latest_schema: Arc, batch: RecordBatch) -> RecordBatch { +fn adapt_batch( + latest_schema: Arc, + batch: RecordBatch, + diff_fields: &mut HashMap, +) -> RecordBatch { let batch_schema = &*batch.schema(); let batch_cols = batch.columns().to_vec(); @@ -654,6 +665,15 @@ pub fn adapt_batch(latest_schema: Arc, batch: RecordBatch) -> RecordBatc if let Some((idx, field)) = batch_schema.column_with_name(field_latest.name()) { cols.push(Arc::clone(&batch_cols[idx])); fields.push(field.clone()); + } else if *field_latest.data_type() == DataType::Utf8View { + // in memtable, the schema should be utf8 + cols.push(new_null_array(&DataType::Utf8, batch.num_rows())); + fields.push(Field::new( + field_latest.name(), + DataType::Utf8, + field_latest.is_nullable(), + )); + diff_fields.insert(field_latest.name().to_string(), DataType::Utf8View); } else { cols.push(new_null_array(field_latest.data_type(), batch.num_rows())); fields.push(field_latest.as_ref().clone()); diff --git a/src/service/search/index.rs b/src/service/search/index.rs index 6ce3c52c78..e24c514e1c 100644 --- a/src/service/search/index.rs +++ b/src/service/search/index.rs @@ -19,7 +19,7 @@ use std::{ }; use config::{ - INDEX_FIELD_NAME_FOR_ALL, + INDEX_FIELD_NAME_FOR_ALL, get_config, utils::tantivy::{query::contains_query::ContainsQuery, tokenizer::o2_collect_tokens}, }; use datafusion::{ @@ -563,7 +563,13 @@ impl Condition { .trim_start_matches('*') // contains .trim_end_matches('*') // prefix or contains .to_string(); - let term = Arc::new(Literal::new(ScalarValue::Utf8(Some(format!("%{value}%"))))); + let term = if get_config().common.utf8_view_enabled { + Arc::new(Literal::new(ScalarValue::Utf8View(Some(format!( + "%{value}%" + ))))) + } else { + Arc::new(Literal::new(ScalarValue::Utf8(Some(format!("%{value}%"))))) + }; let mut expr_list: Vec> = Vec::with_capacity(fst_fields.len()); for field in fst_fields.iter() { @@ -584,7 +590,11 @@ impl Condition { } Condition::FuzzyMatchAll(value, distance) => { let fuzzy_expr = Arc::new(fuzzy_match_udf::FUZZY_MATCH_UDF.clone()); - let term = Arc::new(Literal::new(ScalarValue::Utf8(Some(value.clone())))); + let term = if get_config().common.utf8_view_enabled { + Arc::new(Literal::new(ScalarValue::Utf8View(Some(value.to_string())))) + } else { + Arc::new(Literal::new(ScalarValue::Utf8(Some(value.to_string())))) + }; let distance = Arc::new(Literal::new(ScalarValue::Int64(Some(*distance as i64)))); let mut expr_list: Vec> = Vec::with_capacity(fst_fields.len()); @@ -779,6 +789,9 @@ fn get_scalar_value(value: &str, data_type: &DataType) -> Result, a DataType::Binary => Arc::new(Literal::new(ScalarValue::Binary(Some( value.as_bytes().to_vec(), )))), + DataType::Utf8View => { + Arc::new(Literal::new(ScalarValue::Utf8View(Some(value.to_string())))) + } _ => unimplemented!(), }) } diff --git a/src/service/search/inspector.rs b/src/service/search/inspector.rs index 54be846c9d..0533d228cf 100644 --- a/src/service/search/inspector.rs +++ b/src/service/search/inspector.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use config::{cluster::LOCAL_NODE, meta::cluster::NodeInfo}; +use config::{cluster::LOCAL_NODE, get_config, meta::cluster::NodeInfo}; use serde::{Deserialize, Serialize}; #[derive(Debug, Default, Serialize, Deserialize)] @@ -118,6 +118,14 @@ impl SearchInspectorFieldsBuilder { } pub fn search_inspector_fields(msg: String, kvs: SearchInspectorFields) -> String { + if !get_config().common.search_inspector_enabled { + return msg; + } + + search_inspector_fields_inner(msg, kvs) +} + +fn search_inspector_fields_inner(msg: String, kvs: SearchInspectorFields) -> String { if msg.is_empty() { return msg; } @@ -154,7 +162,7 @@ mod tests { .component("search".to_string()) .build(); - let result = search_inspector_fields(msg.clone(), fields); + let result = search_inspector_fields_inner(msg.clone(), fields); assert!(result.contains("test message")); assert!(result.contains("\"duration\":100")); assert!(result.contains("\"component\":\"search\"")); diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index 3035cebd65..e24ee851fb 100644 --- a/src/service/search/sql.rs +++ b/src/service/search/sql.rs @@ -15,7 +15,7 @@ use std::{ops::ControlFlow, sync::Arc}; -use arrow_schema::FieldRef; +use arrow_schema::{DataType, Field, FieldRef}; use chrono::Duration; use config::{ ALL_VALUES_COL_NAME, ID_COL_NAME, ORIGINAL_DATA_COL_NAME, TIMESTAMP_COL_NAME, get_config, @@ -327,7 +327,33 @@ impl Sql { } } + // 15. replace the Utf8 to Utf8View type + let final_schemas = if cfg.common.utf8_view_enabled { + let mut final_schemas = HashMap::with_capacity(used_schemas.len()); + for (stream, schema) in used_schemas.iter() { + let fields = schema + .schema() + .fields() + .iter() + .map(|f| { + if f.data_type() == &DataType::Utf8 { + Arc::new(Field::new(f.name(), DataType::Utf8View, f.is_nullable())) + } else { + f.clone() + } + }) + .collect::>(); + let new_schema = + Schema::new(fields).with_metadata(schema.schema().metadata().clone()); + final_schemas.insert(stream.clone(), Arc::new(SchemaCache::new(new_schema))); + } + final_schemas + } else { + used_schemas.clone() + }; + let is_complex = is_complex_query(&mut statement); + Ok(Sql { sql: statement.to_string(), is_complex, @@ -339,7 +365,7 @@ impl Sql { prefix_items: prefix_column_visitor.prefix_items, columns, aliases, - schemas: used_schemas, + schemas: final_schemas, limit, offset, time_range: Some((query.start_time, query.end_time)),