From f680073307bde4414937c1173752edea2515538f Mon Sep 17 00:00:00 2001 From: Taiming Liu Date: Fri, 23 May 2025 20:15:46 -0700 Subject: [PATCH] feat: add _timestamp to tantivy index to optimize histogram queries (#6801) --- src/cli/basic/cli.rs | 2 +- src/common/infra/mod.rs | 6 +- src/common/utils/websocket.rs | 2 +- src/config/src/config.rs | 5 +- src/config/src/meta/inverted_index.rs | 25 ++ src/config/src/meta/sql.rs | 3 +- src/config/src/utils/file.rs | 9 +- src/config/src/utils/sql.rs | 4 +- src/config/src/utils/time.rs | 16 ++ src/job/files/parquet.rs | 31 ++ src/job/mod.rs | 4 +- src/migration/mod.rs | 4 +- src/proto/proto/cluster/plan.proto | 11 +- src/proto/src/generated/cluster.rs | 14 +- src/service/db/instance.rs | 29 -- src/service/db/metas.rs | 97 +++++++ src/service/db/mod.rs | 3 +- src/service/db/version.rs | 33 --- .../datafusion/distributed_plan/rewrite.rs | 4 +- .../datafusion/plan/tantivy_count_exec.rs | 174 +++++++++++- src/service/search/grpc/flight.rs | 41 ++- src/service/search/grpc/storage.rs | 152 ++++++++-- src/service/search/grpc/wal.rs | 3 - src/service/search/sql.rs | 268 +++++++++++++----- .../search/tantivy/puffin_directory/mod.rs | 11 +- .../search/tantivy/puffin_directory/reader.rs | 35 +++ 26 files changed, 771 insertions(+), 215 deletions(-) delete mode 100644 src/service/db/instance.rs create mode 100644 src/service/db/metas.rs delete mode 100644 src/service/db/version.rs diff --git a/src/cli/basic/cli.rs b/src/cli/basic/cli.rs index 42e2e5122b..0b3355f360 100644 --- a/src/cli/basic/cli.rs +++ b/src/cli/basic/cli.rs @@ -332,7 +332,7 @@ pub async fn cli() -> Result { let component = command.get_one::("component").unwrap(); match component.as_str() { "version" => { - println!("version: {}", db::version::get().await?); + println!("version: {}", db::metas::version::get().await?); } "user" => { db::user::cache().await?; diff --git a/src/common/infra/mod.rs b/src/common/infra/mod.rs index 2cac1a3abe..a872d9115d 100644 --- a/src/common/infra/mod.rs +++ b/src/common/infra/mod.rs @@ -15,7 +15,7 @@ use ::config::{cache_instance_id, ider}; -use crate::service::db::instance; +use crate::service::db::metas; pub mod cluster; pub mod config; @@ -25,12 +25,12 @@ pub mod wal; pub async fn init() -> Result<(), anyhow::Error> { // set instance id - let instance_id = match instance::get().await { + let instance_id = match metas::instance::get().await { Ok(Some(instance)) => instance, Ok(None) | Err(_) => { log::info!("Generating new instance id"); let id = ider::generate(); - let _ = instance::set(&id).await; + let _ = metas::instance::set(&id).await; id } }; diff --git a/src/common/utils/websocket.rs b/src/common/utils/websocket.rs index a7cf7768bc..85656e7917 100644 --- a/src/common/utils/websocket.rs +++ b/src/common/utils/websocket.rs @@ -81,7 +81,7 @@ pub fn update_histogram_interval_in_query( .pop() .unwrap(); - visit_statements_mut(&mut statement, |stmt| { + let _ = visit_statements_mut(&mut statement, |stmt| { if let Statement::Query(query) = stmt { if let sqlparser::ast::SetExpr::Select(select) = query.body.as_mut() { for projection in &mut select.projection { diff --git a/src/config/src/config.rs b/src/config/src/config.rs index b50a9a96c5..5807e84ad6 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -384,13 +384,12 @@ pub async fn get_sns_client() -> &'static aws_sdk_sns::Client { } pub static BLOCKED_STREAMS: Lazy> = Lazy::new(|| { - let blocked_streams = get_config() + get_config() .common .blocked_streams .split(',') .map(|x| x.to_string()) - .collect(); - blocked_streams + .collect() }); #[derive(EnvConfig)] diff --git a/src/config/src/meta/inverted_index.rs b/src/config/src/meta/inverted_index.rs index 84d07ffb65..8c647d6d64 100644 --- a/src/config/src/meta/inverted_index.rs +++ b/src/config/src/meta/inverted_index.rs @@ -65,6 +65,7 @@ impl std::fmt::Display for InvertedIndexTantivyMode { pub enum InvertedIndexOptimizeMode { SimpleSelect(usize, bool), SimpleCount, + SimpleHistogram(i64, u64, usize), } impl std::fmt::Display for InvertedIndexOptimizeMode { @@ -74,6 +75,12 @@ impl std::fmt::Display for InvertedIndexOptimizeMode { write!(f, "simple_select(limit: {}, ascend: {})", limit, ascend) } InvertedIndexOptimizeMode::SimpleCount => write!(f, "simple_count"), + InvertedIndexOptimizeMode::SimpleHistogram(min_value, bucket_width, num_buckets) => { + write!( + f, + "simple_histogram(min_value: {min_value}, bucket_width: {bucket_width}, num_buckets: {num_buckets})" + ) + } } } } @@ -87,6 +94,13 @@ impl From for InvertedIndexOptimizeMode { Some(cluster_rpc::idx_optimize_mode::Mode::SimpleCount(_)) => { InvertedIndexOptimizeMode::SimpleCount } + Some(cluster_rpc::idx_optimize_mode::Mode::SimpleHistogram(select)) => { + InvertedIndexOptimizeMode::SimpleHistogram( + select.min_value, + select.bucket_width, + select.num_buckets as usize, + ) + } None => panic!("Invalid InvertedIndexOptimizeMode"), } } @@ -108,6 +122,17 @@ impl From for cluster_rpc::IdxOptimizeMode { cluster_rpc::SimpleCount {}, )), }, + InvertedIndexOptimizeMode::SimpleHistogram(min_value, bucket_width, num_buckets) => { + cluster_rpc::IdxOptimizeMode { + mode: Some(cluster_rpc::idx_optimize_mode::Mode::SimpleHistogram( + cluster_rpc::SimpleHistogram { + min_value, + bucket_width, + num_buckets: num_buckets as u32, + }, + )), + } + } } } } diff --git a/src/config/src/meta/sql.rs b/src/config/src/meta/sql.rs index 6ca6e89812..7e7f432829 100644 --- a/src/config/src/meta/sql.rs +++ b/src/config/src/meta/sql.rs @@ -61,9 +61,10 @@ pub fn resolve_stream_names(sql: &str) -> Result, anyhow::Error> { pub fn resolve_stream_names_with_type(sql: &str) -> Result, anyhow::Error> { let dialect = &PostgreSqlDialect {}; - let statement = dbg!(DFParser::parse_sql_with_dialect(sql, dialect))? + let statement = DFParser::parse_sql_with_dialect(sql, dialect)? .pop_back() .ok_or(anyhow::anyhow!("Failed to parse sql"))?; + log::debug!("statement: {:?}", statement); let (table_refs, _) = resolve_table_references(&statement, true)?; let mut tables = Vec::new(); for table in table_refs { diff --git a/src/config/src/utils/file.rs b/src/config/src/utils/file.rs index 20481a967c..9c010190ee 100644 --- a/src/config/src/utils/file.rs +++ b/src/config/src/utils/file.rs @@ -128,10 +128,11 @@ pub async fn scan_files_with_channel( let limit = limit.unwrap_or_default(); let mut files = Vec::with_capacity(std::cmp::max(16, limit)); let dir = std::fs::read_dir(root).map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::Other, - format!("Error reading directory: {}, err: {}", root.display(), e), - ) + std::io::Error::other(format!( + "Error reading directory: {}, err: {}", + root.display(), + e + )) })?; for entry in dir { let path = entry?.path(); diff --git a/src/config/src/utils/sql.rs b/src/config/src/utils/sql.rs index c2085ea2a4..aba386bc58 100644 --- a/src/config/src/utils/sql.rs +++ b/src/config/src/utils/sql.rs @@ -207,7 +207,7 @@ fn has_union(query: &Query) -> bool { fn has_subquery(stat: &Statement) -> bool { let mut visitor = SubqueryVisitor::new(); - stat.visit(&mut visitor); + let _ = stat.visit(&mut visitor); visitor.is_subquery } @@ -238,7 +238,7 @@ impl Visitor for SubqueryVisitor { fn has_timestamp(stat: &Statement) -> bool { let mut visitor = TimestampVisitor::new(); - stat.visit(&mut visitor); + let _ = stat.visit(&mut visitor); visitor.timestamp_selected } diff --git a/src/config/src/utils/time.rs b/src/config/src/utils/time.rs index a15664a0f9..e9070d217a 100644 --- a/src/config/src/utils/time.rs +++ b/src/config/src/utils/time.rs @@ -122,7 +122,23 @@ pub fn parse_str_to_time(s: &str) -> Result, anyhow::Error> { if s.len() == 19 { let fmt = "%Y-%m-%dT%H:%M:%S"; NaiveDateTime::parse_from_str(s, fmt)?.and_utc() + } else if s.contains('.') { + // Handle formats with decimal seconds: "2025-05-14T01:15:25.047" + // First check if it has milliseconds (3 decimal places) + if s.split('.').next_back().unwrap_or("").len() == 3 { + let fmt = "%Y-%m-%dT%H:%M:%S%.3f"; + NaiveDateTime::parse_from_str(s, fmt)?.and_utc() + } else if s.split('.').next_back().unwrap_or("").len() == 6 { + // Handle microseconds (6 decimal places) + let fmt = "%Y-%m-%dT%H:%M:%S%.6f"; + NaiveDateTime::parse_from_str(s, fmt)?.and_utc() + } else { + // Fall back to RFC3339 parsing for other decimal formats + let t = DateTime::parse_from_rfc3339(s)?; + t.into() + } } else { + // Other formats with 'T' but no spaces or decimal points let t = DateTime::parse_from_rfc3339(s)?; t.into() } diff --git a/src/job/files/parquet.rs b/src/job/files/parquet.rs index cd84ff966f..1a5bdaf4fd 100644 --- a/src/job/files/parquet.rs +++ b/src/job/files/parquet.rs @@ -1576,6 +1576,9 @@ pub(crate) async fn generate_tantivy_index( tantivy_schema_builder.add_text_field(INDEX_FIELD_NAME_FOR_ALL, fts_opts); } for field in index_fields.iter() { + if field == TIMESTAMP_COL_NAME { + continue; + } let index_opts = tantivy::schema::TextOptions::default().set_indexing_options( tantivy::schema::TextFieldIndexing::default() .set_index_option(tantivy::schema::IndexRecordOption::Basic) @@ -1584,6 +1587,8 @@ pub(crate) async fn generate_tantivy_index( ); tantivy_schema_builder.add_text_field(field, index_opts); } + // add _timestamp field to tantivy schema + tantivy_schema_builder.add_i64_field(TIMESTAMP_COL_NAME, tantivy::schema::FAST); let tantivy_schema = tantivy_schema_builder.build(); let fts_field = tantivy_schema.get_field(INDEX_FIELD_NAME_FOR_ALL).ok(); @@ -1641,6 +1646,32 @@ pub(crate) async fn generate_tantivy_index( } } + // process _timestamp field + let column_data = match inverted_idx_batch.column_by_name(TIMESTAMP_COL_NAME) { + Some(column_data) => match column_data.as_any().downcast_ref::() { + Some(column_data) => column_data, + None => { + // generate empty array to ensure the tantivy and parquet have same rows + &Int64Array::from(vec![0; num_rows]) + } + }, + None => { + // generate empty array to ensure the tantivy and parquet have same rows + &Int64Array::from(vec![0; num_rows]) + } + }; + let ts_field = tantivy_schema.get_field(TIMESTAMP_COL_NAME).unwrap(); // unwrap directly since added above + const YIELD_THRESHOLD: usize = 100; + let mut batch_size = 0; + for (i, doc) in docs.iter_mut().enumerate() { + doc.add_i64(ts_field, column_data.value(i)); + batch_size += 1; + if batch_size >= YIELD_THRESHOLD { + tokio::task::coop::consume_budget().await; + batch_size = 0; + } + } + tx.send(docs).await?; } Ok(total_num_rows) diff --git a/src/job/mod.rs b/src/job/mod.rs index f58c63f764..de03110ce4 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -96,7 +96,9 @@ pub async fn init() -> Result<(), anyhow::Error> { .expect("organization cache sync failed"); // check version - db::version::set().await.expect("db version set failed"); + db::metas::version::set() + .await + .expect("db version set failed"); // Auth auditing should be done by router also #[cfg(feature = "enterprise")] diff --git a/src/migration/mod.rs b/src/migration/mod.rs index d0b4a238ec..8a43404433 100644 --- a/src/migration/mod.rs +++ b/src/migration/mod.rs @@ -18,7 +18,7 @@ use infra::db::{ORM_CLIENT, ORM_CLIENT_DDL, connect_to_orm, connect_to_orm_ddl}; use schema::migrate_resource_names; use version_compare::Version; -use crate::service::db::version; +use crate::service::db::metas; pub mod dashboards; pub mod file_list; @@ -105,7 +105,7 @@ pub async fn init_db() -> std::result::Result<(), anyhow::Error> { // we initialize both clients here to avoid potential deadlock afterwards ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; // check version upgrade - let old_version = version::get().await.unwrap_or("v0.0.0".to_string()); + let old_version = metas::version::get().await.unwrap_or("v0.0.0".to_string()); check_upgrade(&old_version, config::VERSION).await?; #[allow(deprecated)] diff --git a/src/proto/proto/cluster/plan.proto b/src/proto/proto/cluster/plan.proto index 638ab46bc5..069da4845c 100644 --- a/src/proto/proto/cluster/plan.proto +++ b/src/proto/proto/cluster/plan.proto @@ -63,8 +63,9 @@ message SuperClusterInfo { message IdxOptimizeMode { oneof mode { - SimpleSelect simple_select = 1; - SimpleCount simple_count = 2; + SimpleSelect simple_select = 1; + SimpleCount simple_count = 2; + SimpleHistogram simple_histogram = 3; } } @@ -75,6 +76,12 @@ message SimpleSelect { message SimpleCount {} +message SimpleHistogram { + int64 min_value = 1; + uint64 bucket_width = 2; + uint32 num_buckets = 3; +} + message KvItem { string key = 1; string value = 2; diff --git a/src/proto/src/generated/cluster.rs b/src/proto/src/generated/cluster.rs index 3ffaf5d6b7..8d4c42a2f6 100644 --- a/src/proto/src/generated/cluster.rs +++ b/src/proto/src/generated/cluster.rs @@ -2968,7 +2968,7 @@ pub struct SuperClusterInfo { #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct IdxOptimizeMode { - #[prost(oneof = "idx_optimize_mode::Mode", tags = "1, 2")] + #[prost(oneof = "idx_optimize_mode::Mode", tags = "1, 2, 3")] pub mode: ::core::option::Option, } /// Nested message and enum types in `IdxOptimizeMode`. @@ -2980,6 +2980,8 @@ pub mod idx_optimize_mode { SimpleSelect(super::SimpleSelect), #[prost(message, tag = "2")] SimpleCount(super::SimpleCount), + #[prost(message, tag = "3")] + SimpleHistogram(super::SimpleHistogram), } } #[allow(clippy::derive_partial_eq_without_eq)] @@ -2995,6 +2997,16 @@ pub struct SimpleSelect { pub struct SimpleCount {} #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct SimpleHistogram { + #[prost(int64, tag = "1")] + pub min_value: i64, + #[prost(uint64, tag = "2")] + pub bucket_width: u64, + #[prost(uint32, tag = "3")] + pub num_buckets: u32, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct KvItem { #[prost(string, tag = "1")] pub key: ::prost::alloc::string::String, diff --git a/src/service/db/instance.rs b/src/service/db/instance.rs deleted file mode 100644 index f7c5c8ab9c..0000000000 --- a/src/service/db/instance.rs +++ /dev/null @@ -1,29 +0,0 @@ -// 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 infra::errors::Result; - -pub async fn get() -> Result> { - let ret = super::get("/instance/").await?; - let loc_value = String::from_utf8_lossy(&ret).to_string(); - let loc_value = loc_value.trim().trim_matches('"').to_string(); - let value = Some(loc_value); - Ok(value) -} - -pub async fn set(id: &str) -> Result<()> { - let data = bytes::Bytes::from(id.to_string()); - super::put("/instance/", data, super::NO_NEED_WATCH, None).await -} diff --git a/src/service/db/metas.rs b/src/service/db/metas.rs new file mode 100644 index 0000000000..e5f77f44ea --- /dev/null +++ b/src/service/db/metas.rs @@ -0,0 +1,97 @@ +// 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 crate::service::db; + +// pub static TANTIVY_INDEX_UPDATED_AT: Lazy = Lazy::new(get_or_create_idx_updated_at); + +pub mod version { + use super::db; + + pub async fn get() -> Result { + let ret = db::get("/meta/kv/version").await?; + let version = std::str::from_utf8(&ret).unwrap(); + Ok(version.to_string()) + } + + pub async fn set() -> Result<(), anyhow::Error> { + db::put( + "/meta/kv/version", + bytes::Bytes::from(config::VERSION), + db::NO_NEED_WATCH, + None, + ) + .await?; + Ok(()) + } +} + +pub mod instance { + use infra::errors::Result; + + use super::db; + + pub async fn get() -> Result> { + let ret = db::get("/instance/").await?; + let loc_value = String::from_utf8_lossy(&ret).to_string(); + let loc_value = loc_value.trim().trim_matches('"').to_string(); + let value = Some(loc_value); + Ok(value) + } + + pub async fn set(id: &str) -> Result<()> { + let data = bytes::Bytes::from(id.to_string()); + db::put("/instance/", data, db::NO_NEED_WATCH, None).await + } +} + +pub mod tantivy_index { + use tokio::sync::OnceCell; + + use super::db; + + static TIMESTAMP_UPDATED_AT: OnceCell = OnceCell::const_new(); + + pub async fn get_ttv_timestamp_updated_at() -> i64 { + TIMESTAMP_UPDATED_AT + .get_or_init(get_or_create_idx_updated_at) + .await + .to_owned() + } + + async fn get_or_create_idx_updated_at() -> i64 { + let key = "/tantivy/_timestamp/updated_at"; + match db::get(key).await { + Ok(ret) if !ret.is_empty() => { + let timestamp = String::from_utf8_lossy(&ret) + .to_string() + .parse::() + .unwrap(); + timestamp + } + _ => { + let timestamp = config::utils::time::now_micros(); + let data = bytes::Bytes::from(timestamp.to_string()); + if let Err(e) = db::put(key, data, db::NO_NEED_WATCH, None).await { + log::warn!( + "[db::metas] Error storing tantivy _timestamp index updated_at: {}", + e + ); + } + timestamp + } + } + } +} diff --git a/src/service/db/mod.rs b/src/service/db/mod.rs index 34803d249e..02f850a75c 100644 --- a/src/service/db/mod.rs +++ b/src/service/db/mod.rs @@ -29,10 +29,10 @@ pub mod distinct_values; pub mod enrichment_table; pub mod file_list; pub mod functions; -pub mod instance; #[cfg(feature = "enterprise")] pub mod keys; pub mod kv; +pub mod metas; pub mod metrics; #[cfg(feature = "enterprise")] pub mod ofga; @@ -46,7 +46,6 @@ pub mod session; pub mod short_url; pub mod syslog; pub mod user; -pub mod version; pub(crate) use infra_db::{Event, NEED_WATCH, NO_NEED_WATCH, get_coordinator}; diff --git a/src/service/db/version.rs b/src/service/db/version.rs deleted file mode 100644 index 7c0ebf5c06..0000000000 --- a/src/service/db/version.rs +++ /dev/null @@ -1,33 +0,0 @@ -// 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 crate::service::db; - -pub async fn get() -> Result { - let ret = db::get("/meta/kv/version").await?; - let version = std::str::from_utf8(&ret).unwrap(); - Ok(version.to_string()) -} - -pub async fn set() -> Result<(), anyhow::Error> { - db::put( - "/meta/kv/version", - bytes::Bytes::from(config::VERSION), - db::NO_NEED_WATCH, - None, - ) - .await?; - Ok(()) -} diff --git a/src/service/search/datafusion/distributed_plan/rewrite.rs b/src/service/search/datafusion/distributed_plan/rewrite.rs index 98e5e05768..b2fd31fa24 100644 --- a/src/service/search/datafusion/distributed_plan/rewrite.rs +++ b/src/service/search/datafusion/distributed_plan/rewrite.rs @@ -53,7 +53,7 @@ impl RemoteScanRewriter { equal_keys: HashMap>, match_all_keys: Vec, index_condition: Option, - index_optimizer_mode: Option, + index_optimize_mode: Option, is_leader: bool, opentelemetry_context: opentelemetry::Context, ) -> Self { @@ -66,7 +66,7 @@ impl RemoteScanRewriter { equal_keys, match_all_keys, index_condition, - index_optimizer_mode, + index_optimize_mode, is_leader, opentelemetry_context, ), diff --git a/src/service/search/datafusion/plan/tantivy_count_exec.rs b/src/service/search/datafusion/plan/tantivy_count_exec.rs index 336fd7bd8b..7769340b9a 100644 --- a/src/service/search/datafusion/plan/tantivy_count_exec.rs +++ b/src/service/search/datafusion/plan/tantivy_count_exec.rs @@ -15,7 +15,9 @@ use std::{any::Any, sync::Arc}; -use arrow::array::{Array, Int64Array}; +use arrow::array::{ + Array, Int64Array, TimestampMicrosecondArray, TimestampNanosecondArray, UInt64Array, +}; use config::meta::{inverted_index::InvertedIndexOptimizeMode, stream::FileKey}; use datafusion::{ arrow::{array::RecordBatch, datatypes::SchemaRef}, @@ -36,29 +38,32 @@ use crate::service::search::{ }; #[derive(Debug)] -pub struct TantivyCountExec { +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 + index_optimize_mode: InvertedIndexOptimizeMode, // Type of query the ttv index optimizes } -impl TantivyCountExec { - /// Create a new TantivyCountExec +impl TantivyOptimizeExec { + /// Create a new TantivyOptimizeExec pub fn new( query: Arc, schema: SchemaRef, file_list: Vec, index_condition: IndexCondition, + index_optimize_mode: InvertedIndexOptimizeMode, ) -> Self { let cache = Self::compute_properties(Arc::clone(&schema)); - TantivyCountExec { + TantivyOptimizeExec { query, schema, file_list, index_condition, cache, + index_optimize_mode, } } @@ -74,7 +79,7 @@ impl TantivyCountExec { } } -impl DisplayAs for TantivyCountExec { +impl DisplayAs for TantivyOptimizeExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { // display up to 5 file keys, let file_keys = self @@ -86,15 +91,15 @@ impl DisplayAs for TantivyCountExec { .join(", "); write!( f, - "TantivyCountExec: files: {}, file_list: [{file_keys}]", + "TantivyOptimizeExec: files: {}, file_list: [{file_keys}]", self.file_list.len() ) } } -impl ExecutionPlan for TantivyCountExec { +impl ExecutionPlan for TantivyOptimizeExec { fn name(&self) -> &'static str { - "TantivyCountExec" + "TantivyOptimizeExec" } fn as_any(&self) -> &dyn Any { @@ -123,7 +128,7 @@ impl ExecutionPlan for TantivyCountExec { ) -> Result { if partition >= 1 { return internal_err!( - "TantivyCountExec invalid partition {partition} (expected partition: 0)" + "TantivyOptimizeExec invalid partition {partition} (expected partition: 0)" ); } @@ -132,6 +137,7 @@ impl ExecutionPlan for TantivyCountExec { self.file_list.clone(), Some(self.index_condition.clone()), self.schema.clone(), + self.index_optimize_mode.clone(), ); let stream = futures::stream::once(fut); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -150,12 +156,13 @@ async fn adapt_tantivy_result( mut file_list: Vec, index_condition: Option, schema: SchemaRef, + idx_optimize_mode: InvertedIndexOptimizeMode, ) -> Result { - let (idx_took, error, total_hits) = filter_file_list_by_tantivy_index( + let (idx_took, error, total_hits, histogram_counts) = filter_file_list_by_tantivy_index( query.clone(), &mut file_list, index_condition, - Some(InvertedIndexOptimizeMode::SimpleCount), + Some(idx_optimize_mode.clone()), ) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; @@ -165,7 +172,7 @@ async fn adapt_tantivy_result( } log::info!( - "[trace_id {}] search->storage: stream {}/{}/{}, TantivyCountExec execution time {} ms", + "[trace_id {}] search->storage: stream {}/{}/{}, TantivyOptimizeExec execution time {} ms", query.trace_id, query.org_id, query.stream_type, @@ -173,9 +180,146 @@ async fn adapt_tantivy_result( idx_took ); - let array = vec![Arc::new(Int64Array::from(vec![total_hits as i64])) as Arc]; + let array = match idx_optimize_mode { + InvertedIndexOptimizeMode::SimpleCount => { + vec![Arc::new(Int64Array::from(vec![total_hits as i64])) as Arc] + } + InvertedIndexOptimizeMode::SimpleHistogram(min_value, bucket_width, num_buckets) => { + create_histogram_arrow_array( + &schema, + histogram_counts, + min_value, + bucket_width, + num_buckets, + )? + } + _ => { + return internal_err!( + "Only count and histogram optimize modes are supported by TantivyOptimizeExec" + ); + } + }; RecordBatch::try_new(schema, array).map_err(|e| { - DataFusionError::Internal(format!("TantivyCountExec create record batch error: {e}")) + DataFusionError::Internal(format!( + "TantivyOptimizeExec create record batch error: {e}", + )) }) } + +/// Creates a RecordBatch containing histogram data with timestamps and counts +/// +/// Parameters: +/// - schema: The expected schema for the result +/// - histogram_counts: Vector of counts for each histogram bucket +/// - min_value: The minimum timestamp value (start of first bucket) +/// - bucket_width: Width of each bucket in microseconds +/// - num_buckets: Expected number of buckets +fn create_histogram_arrow_array( + schema: &SchemaRef, + histogram_counts: Vec, + min_value: i64, + bucket_width: u64, + num_buckets: usize, +) -> Result>, DataFusionError> { + // Validate inputs + if bucket_width == 0 { + return Err(DataFusionError::Internal( + "Histogram bucket width cannot be zero".to_string(), + )); + } + + // Verify schema has expected structure + if schema.fields().len() != 2 { + return Err(DataFusionError::Internal(format!( + "Expected schema with 2 fields for histogram, got {}", + schema.fields().len() + ))); + } + + // Ensure we have the right number of buckets (pad or truncate if necessary) + let normalized_counts = if histogram_counts.len() != num_buckets { + log::warn!( + "Histogram counts length ({}) doesn't match expected buckets ({}), normalizing data", + histogram_counts.len(), + num_buckets + ); + let mut normalized = histogram_counts; + normalized.resize(num_buckets, 0); + normalized + } else { + histogram_counts + }; + + // Pre-allocate vectors with exact capacity + let mut timestamp_values = Vec::with_capacity(num_buckets); + let mut count_values = Vec::with_capacity(num_buckets); + + for (i, &count) in normalized_counts.iter().enumerate() { + if count > 0 { + let bucket_offset = i64::try_from(i) + .and_then(|i| i64::try_from(bucket_width).map(|w| i * w)) + .map_err(|_| { + DataFusionError::Internal( + "Bucket timestamp calculation would overflow".to_string(), + ) + })?; + let bucket_timestamp = min_value.checked_add(bucket_offset).ok_or_else(|| { + DataFusionError::Internal("Bucket timestamp calculation would overflow".to_string()) + })?; + timestamp_values.push(bucket_timestamp); + count_values.push(count as i64); + } + } + + // Get field data types from schema to ensure we create the right array types + let timestamp_field = &schema.fields()[0]; + let count_field = &schema.fields()[1]; + + // Create arrays with proper types based on schema + let timestamp_array = match timestamp_field.data_type() { + arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _) => { + Arc::new(TimestampMicrosecondArray::from(timestamp_values)) as Arc + } + arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => { + // Convert microseconds to nanoseconds + let nano_values = timestamp_values + .iter() + .map(|&ts| { + ts.checked_mul(1000).ok_or_else(|| { + DataFusionError::Internal( + "Timestamp conversion to nanoseconds would overflow".to_string(), + ) + }) + }) + .collect::, DataFusionError>>()?; + Arc::new(TimestampNanosecondArray::from(nano_values)) as Arc + } + // Handle other timestamp types as needed + _ => { + return Err(DataFusionError::Internal(format!( + "Unexpected timestamp type in histogram schema: {:?}", + timestamp_field.data_type() + ))); + } + }; + + // Create count array + let count_array = match count_field.data_type() { + arrow_schema::DataType::Int64 => Arc::new(Int64Array::from(count_values)) as Arc, + arrow_schema::DataType::UInt64 => { + // Convert to unsigned if schema requires it + let u64_values = count_values.iter().map(|&c| c as u64).collect::>(); + Arc::new(UInt64Array::from(u64_values)) as Arc + } + // Add other numeric types as needed + _ => { + return Err(DataFusionError::Internal(format!( + "Unexpected count type in histogram schema: {:?}", + count_field.data_type() + ))); + } + }; + + Ok(vec![timestamp_array, count_array]) +} diff --git a/src/service/search/grpc/flight.rs b/src/service/search/grpc/flight.rs index 6cb6edde6d..00a503a180 100644 --- a/src/service/search/grpc/flight.rs +++ b/src/service/search/grpc/flight.rs @@ -55,7 +55,7 @@ use crate::service::{ empty_exec::NewEmptyExec, }, exec::{prepare_datafusion_context, register_udf}, - plan::tantivy_count_exec::TantivyCountExec, + plan::tantivy_count_exec::TantivyOptimizeExec, table_provider::uniontable::NewUnionTable, }, index::IndexCondition, @@ -161,7 +161,8 @@ pub async fn search( } }) .collect_vec(); - let index_updated_at = get_stream_setting_index_updated_at(&stream_settings, stream_created_at); + let mut index_updated_at = + get_stream_setting_index_updated_at(&stream_settings, stream_created_at); // construct partition filters let search_partition_keys: Vec<(String, String)> = req @@ -187,7 +188,7 @@ pub async fn search( use_inverted_index: req.index_info.use_inverted_index, }); - let mut idx_optimize_rule: Option = + let idx_optimize_rule: Option = req.index_info.index_optimize_mode.clone().map(|x| x.into()); // get all tables @@ -231,13 +232,23 @@ pub async fn search( ) ); - if physical_plan.name() == "AggregateExec" - && physical_plan.schema().fields().len() == 1 + let mut storage_search_idx_optimize_rule = idx_optimize_rule.clone(); + let is_aggregate_exec = physical_plan.name() == "AggregateExec"; + let is_simple_count = physical_plan.schema().fields().len() == 1 && matches!( idx_optimize_rule, Some(InvertedIndexOptimizeMode::SimpleCount) - ) - { + ); + let is_simple_histogram = matches!( + idx_optimize_rule, + Some(InvertedIndexOptimizeMode::SimpleHistogram(..)) + ); + if is_simple_histogram { + let ttv_timestamp_updated_at = + db::metas::tantivy_index::get_ttv_timestamp_updated_at().await; + index_updated_at = index_updated_at.max(ttv_timestamp_updated_at); + } + if is_aggregate_exec && (is_simple_count || is_simple_histogram) { let (tantivy_files, datafusion_files) = split_file_list_by_time_range( file_list, req.search_info.start_time, @@ -246,7 +257,14 @@ pub async fn search( ); tantivy_file_list = tantivy_files; file_list = datafusion_files; - idx_optimize_rule = None; + storage_search_idx_optimize_rule = None; + log::debug!( + "[trace_id {}] flight->search: after_split_file idx: {}, datafusion_files: {}, optimize_rule: {:?}", + trace_id, + tantivy_file_list.len(), + file_list.len(), + storage_search_idx_optimize_rule + ); } // sort by max_ts, the latest file should be at the top @@ -262,7 +280,7 @@ pub async fn search( file_stats_cache.clone(), index_condition.clone(), fst_fields.clone(), - idx_optimize_rule, + storage_search_idx_optimize_rule, ) .await { @@ -354,11 +372,12 @@ pub async fn search( if !tantivy_file_list.is_empty() { scan_stats.add(&collect_stats(&tantivy_file_list)); - let tantivy_exec = Arc::new(TantivyCountExec::new( + let tantivy_exec = Arc::new(TantivyOptimizeExec::new( query_params, physical_plan.schema(), tantivy_file_list, index_condition.unwrap(), + 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 _])); } @@ -476,7 +495,7 @@ fn split_file_list_by_time_range( file_list.into_iter().partition(|file| { file.meta.min_ts >= start_time && file.meta.max_ts <= end_time - && file.meta.min_ts > index_updated_at + && file.meta.min_ts >= index_updated_at && file.meta.index_size > 0 }) } diff --git a/src/service/search/grpc/storage.rs b/src/service/search/grpc/storage.rs index 2bec8166e6..d8c274cd14 100644 --- a/src/service/search/grpc/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -18,7 +18,7 @@ use std::{collections::HashSet, sync::Arc}; use anyhow::Context; use arrow_schema::Schema; use config::{ - FILE_EXT_TANTIVY, FILE_EXT_TANTIVY_FOLDER, INDEX_FIELD_NAME_FOR_ALL, + FILE_EXT_TANTIVY, FILE_EXT_TANTIVY_FOLDER, INDEX_FIELD_NAME_FOR_ALL, TIMESTAMP_COL_NAME, cluster::LOCAL_NODE, get_config, is_local_disk_storage, meta::{ @@ -142,7 +142,7 @@ pub async fn search( let mut idx_took = 0; let mut is_add_filter_back = false; if use_inverted_index { - (idx_took, is_add_filter_back, _) = filter_file_list_by_tantivy_index( + (idx_took, is_add_filter_back, ..) = filter_file_list_by_tantivy_index( query.clone(), &mut files, index_condition.clone(), @@ -344,9 +344,6 @@ pub async fn search( let mut tables = Vec::new(); let start = std::time::Instant::now(); for (ver, files) in files_group { - if files.is_empty() { - continue; - } if files.is_empty() { continue; } @@ -506,8 +503,8 @@ pub async fn filter_file_list_by_tantivy_index( query: Arc, file_list: &mut Vec, index_condition: Option, - idx_optimize_rule: Option, -) -> Result<(usize, bool, usize), Error> { + idx_optimize_mode: Option, +) -> Result<(usize, bool, usize, Vec), Error> { let start = std::time::Instant::now(); let cfg = get_config(); @@ -606,7 +603,7 @@ pub async fn filter_file_list_by_tantivy_index( let time_range = query.time_range.unwrap_or((0, 0)); let index_parquet_files = index_file_names.into_iter().map(|(_, f)| f).collect_vec(); let (mut index_parquet_files, query_limit) = - if let Some(InvertedIndexOptimizeMode::SimpleSelect(limit, _ascend)) = idx_optimize_rule { + if let Some(InvertedIndexOptimizeMode::SimpleSelect(limit, _ascend)) = idx_optimize_mode { if limit > 0 { ( group_files_by_time_range(index_parquet_files, target_partitions), @@ -627,6 +624,7 @@ pub async fn filter_file_list_by_tantivy_index( let mut no_more_files = false; let mut total_hits = 0; + let mut total_histogram_hits = vec![]; let group_num = index_parquet_files.len(); let max_group_len = index_parquet_files .iter() @@ -668,7 +666,7 @@ pub async fn filter_file_list_by_tantivy_index( // Spawn a task for each file, wherein full text search and // secondary index search queries are executed let index_condition_clone = index_condition.clone(); - let idx_optimize_rule_clone = idx_optimize_rule.clone(); + let idx_optimize_rule_clone = idx_optimize_mode.clone(); let permit = semaphore.clone().acquire_owned().await.unwrap(); let task = tokio::task::spawn(async move { // spawn a new task for catching the panic error @@ -726,13 +724,21 @@ pub async fn filter_file_list_by_tantivy_index( } }; for result in tasks { - let result: anyhow::Result<(String, Option, usize)> = result; // Each result corresponds to a file in the file list match result { - Ok((file_name, bitvec, hits_in_file)) => { + Ok((file_name, bitvec, hits_in_file, histogram_hits)) => { + let histogram_hits_count = histogram_hits.iter().sum::(); total_hits += hits_in_file; + if !histogram_hits.is_empty() { + total_histogram_hits.push(histogram_hits); + } if file_name.is_empty() && bitvec.is_none() { // no need inverted index for this file, need add filter back + log::warn!( + "[trace_id {}] search->tantivy: no hits for index_condition: {:?}. Adding the parquet file back for Datafusion search", + query.trace_id, + index_condition, + ); is_add_filter_back = true; continue; } @@ -762,6 +768,14 @@ pub async fn filter_file_list_by_tantivy_index( hits_in_file, file_name ); + } else if histogram_hits_count > 0 { + log::debug!( + "[trace_id {}] search->tantivy: histogram hits for index_condition {:?} found {} in {}", + query.trace_id, + index_condition, + hits_in_file, + file_name + ); } else { log::debug!( "[trace_id {}] search->tantivy: no match found in index for file {}", @@ -789,14 +803,31 @@ pub async fn filter_file_list_by_tantivy_index( } } + let final_histogram_hits = if total_histogram_hits.is_empty() { + Vec::new() + } else { + // note: all histogram_hits should have the same length + let len = total_histogram_hits[0].len(); + (0..len) + .map(|i| { + total_histogram_hits + .iter() + .map(|v| v.get(i).unwrap_or(&0)) + .sum::() + }) + .collect() + }; + let histogram_hits_sum = final_histogram_hits.iter().sum::(); + log::info!( "{}", search_inspector_fields( format!( - "[trace_id {}] search->tantivy: total hits for index_condition: {:?} found {} rows, is_add_filter_back: {}, file_num: {}, took: {} ms", + "[trace_id {}] search->tantivy: total hits for index_condition: {:?} found {} rows, {} histogram_hits, is_add_filter_back: {}, file_num: {}, took: {} ms", query.trace_id, index_condition, total_hits, + histogram_hits_sum, is_add_filter_back, file_list_map.len(), search_start.elapsed().as_millis() @@ -807,8 +838,9 @@ pub async fn filter_file_list_by_tantivy_index( .search_role("follower".to_string()) .duration(search_start.elapsed().as_millis() as usize) .desc(format!( - "found {} rows, is_add_filter_back: {}, file_num: {}", + "found {} rows, {} histogram_hits, is_add_filter_back: {}, file_num: {}", total_hits, + histogram_hits_sum, is_add_filter_back, file_list_map.len(), )) @@ -821,6 +853,7 @@ pub async fn filter_file_list_by_tantivy_index( start.elapsed().as_millis() as usize, is_add_filter_back, total_hits, + final_histogram_hits, )) } @@ -847,7 +880,7 @@ async fn search_tantivy_index( index_condition: Option, idx_optimize_rule: Option, parquet_file: &FileKey, -) -> anyhow::Result<(String, Option, usize)> { +) -> anyhow::Result<(String, Option, usize, Vec)> { let file_account = parquet_file.account.clone(); let Some(ttv_file_name) = convert_parquet_idx_file_name_to_tantivy_file(&parquet_file.key) else { @@ -951,18 +984,22 @@ async fn search_tantivy_index( warm_terms.insert(field, HashMap::new()); } } - warm_up_terms(&tantivy_searcher, &warm_terms).await?; + let need_fast_field = idx_optimize_rule + .as_ref() + .is_some_and(|rule| matches!(rule, InvertedIndexOptimizeMode::SimpleHistogram(..))); + warm_up_terms(&tantivy_searcher, &warm_terms, need_fast_field).await?; } // search the index let file_in_range = parquet_file.meta.min_ts <= time_range.1 && parquet_file.meta.max_ts >= time_range.0; let idx_optimize_rule_clone = idx_optimize_rule.clone(); + // TODO(taiming): refactor the return type throughout the tantivy index search let matched_docs = tokio::task::spawn_blocking(move || match (file_in_range, idx_optimize_rule_clone) { (false, _) | (true, None) => tantivy_searcher .search(&query, &tantivy::collector::DocSetCollector) - .map(|ret| (ret, 0)), + .map(|ret| (ret, 0, vec![])), (true, Some(InvertedIndexOptimizeMode::SimpleSelect(limit, ascend))) => { tantivy_searcher .search( @@ -983,22 +1020,53 @@ async fn search_tantivy_index( ( ret.into_iter().map(|(_, doc)| doc).collect::>(), 0, + vec![], ) }) } (true, Some(InvertedIndexOptimizeMode::SimpleCount)) => tantivy_searcher .search(&query, &tantivy::collector::Count) - .map(|ret| (HashSet::new(), ret)), + .map(|ret| (HashSet::new(), ret, vec![])), + ( + true, + Some(InvertedIndexOptimizeMode::SimpleHistogram( + min_value, + bucket_width, + num_buckets, + )), + ) => { + // fail the function if field not in tantivy schema + if tantivy_schema.get_field(TIMESTAMP_COL_NAME).is_err() { + log::warn!("_timestamp not index in tantivy file: {}", ttv_file_name); + return Ok((HashSet::new(), 0, vec![])); + } + tantivy_searcher + .search( + &query, + &tantivy::collector::HistogramCollector::new::( + TIMESTAMP_COL_NAME.to_string(), + min_value, + bucket_width, + num_buckets, + ), + ) + .map(|ret| (HashSet::new(), 0, ret)) + } }) .await??; // return early if no matches in tantivy - let (matched_docs, total_hits) = matched_docs; - if total_hits > 0 { - return Ok((parquet_file.key.to_string(), None, total_hits)); + let (matched_docs, total_hits, histogram_hits) = matched_docs; + if total_hits > 0 || !histogram_hits.is_empty() { + return Ok(( + parquet_file.key.to_string(), + None, + total_hits, + histogram_hits, + )); } if matched_docs.is_empty() { - return Ok((parquet_file.key.to_string(), None, 0)); + return Ok((parquet_file.key.to_string(), None, 0, vec![])); } // return early if the number of matched docs is too large if cfg.limit.inverted_index_skip_threshold > 0 @@ -1014,7 +1082,7 @@ async fn search_tantivy_index( cfg.limit.inverted_index_skip_threshold, parquet_file.key ); - return Ok(("".to_string(), None, 0)); + return Ok(("".to_string(), None, 0, vec![])); } // Prepare a vec of segment offsets @@ -1032,7 +1100,7 @@ async fn search_tantivy_index( for doc in matched_docs { res.set(doc.doc_id as usize, true); } - Ok((parquet_file.key.to_string(), Some(res), matched_num)) + Ok((parquet_file.key.to_string(), Some(res), matched_num, vec![])) } // Group files by time range @@ -1214,4 +1282,42 @@ mod tests { let max_index = find_max_group_index(&groups); assert_eq!(max_index, 1); } + + #[test] + fn test_histogram_i64() { + const MARGIN_IN_BYTES: usize = 1_000_000; + const MEMORY_BUDGET_NUM_BYTES_MIN: usize = ((MARGIN_IN_BYTES as u32) * 15u32) as usize; + + let mut schema_builder = tantivy::schema::SchemaBuilder::new(); + let val_field = schema_builder.add_i64_field(TIMESTAMP_COL_NAME, tantivy::schema::FAST); + let schema = schema_builder.build(); + let index = tantivy::index::Index::create_in_ram(schema); + let mut writer = index + .writer_with_num_threads(1, MEMORY_BUDGET_NUM_BYTES_MIN) + .unwrap(); + writer + .add_document(tantivy::doc!(val_field=>12i64)) + .unwrap(); + writer + .add_document(tantivy::doc!(val_field=>-30i64)) + .unwrap(); + writer + .add_document(tantivy::doc!(val_field=>-12i64)) + .unwrap(); + writer + .add_document(tantivy::doc!(val_field=>-10i64)) + .unwrap(); + writer.commit().unwrap(); + let reader = index.reader().unwrap(); + let searcher = reader.searcher(); + let all_query = tantivy::query::AllQuery; + let histogram_collector = tantivy::collector::HistogramCollector::new( + TIMESTAMP_COL_NAME.to_string(), + -20i64, + 10u64, + 4, + ); + let histogram = searcher.search(&all_query, &histogram_collector).unwrap(); + assert_eq!(histogram, vec![1, 1, 0, 1]); + } } diff --git a/src/service/search/grpc/wal.rs b/src/service/search/grpc/wal.rs index 00e391c674..3846de3842 100644 --- a/src/service/search/grpc/wal.rs +++ b/src/service/search/grpc/wal.rs @@ -275,9 +275,6 @@ pub async fn search_parquet( let mut tables = Vec::new(); let start = std::time::Instant::now(); for (ver, files) in files_group { - if files.is_empty() { - continue; - } if files.is_empty() { continue; } diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index 48594e1277..e5982a9408 100644 --- a/src/service/search/sql.rs +++ b/src/service/search/sql.rs @@ -140,12 +140,12 @@ impl Sql { // 2. rewrite track_total_hits if query.track_total_hits { let mut trace_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut trace_total_hits_visitor); + let _ = statement.visit(&mut trace_total_hits_visitor); } // 3. get column name, alias, group by, order by let mut column_visitor = ColumnVisitor::new(&total_schemas); - statement.visit(&mut column_visitor); + let _ = statement.visit(&mut column_visitor); let columns = column_visitor.columns.clone(); let aliases = column_visitor @@ -180,7 +180,7 @@ impl Sql { // 4. get match_all() value let mut match_visitor = MatchVisitor::new(); - statement.visit(&mut match_visitor); + let _ = statement.visit(&mut match_visitor); let need_fst_fields = match_visitor.match_items.is_some(); // 5. check if have full text search filed in stream @@ -225,25 +225,25 @@ impl Sql { // 7. get partition column value let mut partition_column_visitor = PartitionColumnVisitor::new(&used_schemas); - statement.visit(&mut partition_column_visitor); + let _ = statement.visit(&mut partition_column_visitor); // 8. get prefix column value let mut prefix_column_visitor = PrefixColumnVisitor::new(&used_schemas); - statement.visit(&mut prefix_column_visitor); + let _ = statement.visit(&mut prefix_column_visitor); // 9. pick up histogram interval let mut histogram_interval_visitor = - HistogramIntervalVistor::new(Some((query.start_time, query.end_time))); - statement.visit(&mut histogram_interval_visitor); + HistogramIntervalVisitor::new(Some((query.start_time, query.end_time))); + let _ = statement.visit(&mut histogram_interval_visitor); // NOTE: only this place modify the sql // 10. add _timestamp and _o2_id if need if !is_complex_query(&mut statement) { let mut add_timestamp_visitor = AddTimestampVisitor::new(); - statement.visit(&mut add_timestamp_visitor); + let _ = statement.visit(&mut add_timestamp_visitor); if o2_id_is_needed(&used_schemas, &search_event_type) { let mut add_o2_id_visitor = AddO2IdVisitor::new(); - statement.visit(&mut add_o2_id_visitor); + let _ = statement.visit(&mut add_o2_id_visitor); } } @@ -262,7 +262,7 @@ impl Sql { cfg.common.feature_query_remove_filter_with_index, cfg.common.inverted_index_count_optimizer_enabled, ); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); index_condition = index_visitor.index_condition; can_optimize = index_visitor.can_optimize; } @@ -280,9 +280,28 @@ impl Sql { )); } - // 13. check `select count(*) from table where match_all` optimizer - if can_optimize && is_simple_count_query(&mut statement) { - index_optimize_mode = Some(InvertedIndexOptimizeMode::SimpleCount); + // 13. check other inverted index optimize modes + // `select count(*) from table where match_all` -> SimpleCount + // or `select histogram(..), count(*) from table where match_all` -> SimpleHistogram + if can_optimize && index_optimize_mode.is_none() { + let mut visitor = OtherIndexOptimizeModeVisitor::new(); + let _ = statement.visit(&mut visitor); + if visitor.is_simple_count { + index_optimize_mode = Some(InvertedIndexOptimizeMode::SimpleCount); + } else if visitor.is_simple_histogram && histogram_interval_visitor.interval.is_some() { + let bucket_width = histogram_interval_visitor.interval.unwrap() as u64 * 1_000_000; + // round the bucket edges to even start + let rounding_by = bucket_width as i64; + let min_value = (query.start_time / rounding_by) * rounding_by; + let max_value = (query.end_time / rounding_by) * rounding_by; + let num_buckets = + ((max_value - min_value) as f64 / bucket_width as f64).ceil() as usize; + index_optimize_mode = Some(InvertedIndexOptimizeMode::SimpleHistogram( + min_value, + bucket_width, + num_buckets, + )); + } } Ok(Sql { @@ -667,7 +686,7 @@ impl VisitorMut for ColumnVisitor<'_> { if let Some(order_by) = query.order_by.as_mut() { for order in order_by.exprs.iter_mut() { let mut name_visitor = FieldNameVisitor::new(); - order.expr.visit(&mut name_visitor); + let _ = order.expr.visit(&mut name_visitor); if name_visitor.field_names.len() == 1 { let expr_name = name_visitor.field_names.iter().next().unwrap().to_string(); self.order_by.push(( @@ -697,7 +716,7 @@ impl VisitorMut for ColumnVisitor<'_> { if let GroupByExpr::Expressions(exprs, _) = &mut select.group_by { for expr in exprs.iter_mut() { let mut name_visitor = FieldNameVisitor::new(); - expr.visit(&mut name_visitor); + let _ = expr.visit(&mut name_visitor); if name_visitor.field_names.len() == 1 { let expr_name = name_visitor.field_names.iter().next().unwrap().to_string(); self.group_by.push(expr_name); @@ -1098,7 +1117,7 @@ impl VisitorMut for AddTimestampVisitor { match item { SelectItem::UnnamedExpr(expr) => { let mut visitor = FieldNameVisitor::new(); - expr.visit(&mut visitor); + let _ = expr.visit(&mut visitor); if visitor.field_names.contains(TIMESTAMP_COL_NAME) { has_timestamp = true; break; @@ -1106,7 +1125,7 @@ impl VisitorMut for AddTimestampVisitor { } SelectItem::ExprWithAlias { expr, alias: _ } => { let mut visitor = FieldNameVisitor::new(); - expr.visit(&mut visitor); + let _ = expr.visit(&mut visitor); if visitor.field_names.contains(TIMESTAMP_COL_NAME) { has_timestamp = true; break; @@ -1151,7 +1170,7 @@ impl VisitorMut for AddO2IdVisitor { match item { SelectItem::UnnamedExpr(expr) => { let mut visitor = FieldNameVisitor::new(); - expr.visit(&mut visitor); + let _ = expr.visit(&mut visitor); if visitor.field_names.contains(ID_COL_NAME) { has_o2_id = true; break; @@ -1159,7 +1178,7 @@ impl VisitorMut for AddO2IdVisitor { } SelectItem::ExprWithAlias { expr, alias: _ } => { let mut visitor = FieldNameVisitor::new(); - expr.visit(&mut visitor); + let _ = expr.visit(&mut visitor); if visitor.field_names.contains(ID_COL_NAME) { has_o2_id = true; break; @@ -1183,66 +1202,92 @@ impl VisitorMut for AddO2IdVisitor { } } -fn is_simple_count_query(statement: &mut Statement) -> bool { - let mut visitor = SimpleCountVisitor::new(); - statement.visit(&mut visitor); - visitor.is_simple_count -} - // check if the query is simple count query -// 1. don't has subquery -// 2. don't has join -// 3. don't has group by -// 4. only has count(*) -// 5. don't has SetOperation(UNION/EXCEPT/INTERSECT of two queries) -// 6. don't has distinct -struct SimpleCountVisitor { +// 1. doesn't have subquery +// 2. doesn't have join +// 3. doesn't have group by +// 4. doesn't have SetOperation(UNION/EXCEPT/INTERSECT of two queries) +// 5. doesn't have distinct +// +// either only has count(*) -> SimpleCount +// or has histogram(...) and count(*) -> SimpleHistogram +struct OtherIndexOptimizeModeVisitor { pub is_simple_count: bool, + pub is_simple_histogram: bool, } -impl SimpleCountVisitor { +impl OtherIndexOptimizeModeVisitor { fn new() -> Self { Self { is_simple_count: false, + is_simple_histogram: false, } } } -impl VisitorMut for SimpleCountVisitor { +impl VisitorMut for OtherIndexOptimizeModeVisitor { type Break = (); fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow { - // Check if the query is select * if let sqlparser::ast::SetExpr::Select(select) = query.body.as_ref() { - if select.projection.len() != 1 + if select.projection.len() > 2 || select.distinct.is_some() || select.from.len() > 1 - || select.from.iter().any(|from| !from.joins.is_empty()) - || !matches!(select.group_by, GroupByExpr::Expressions(ref expr, _) if expr.is_empty()) + || select.from.iter().any(|from| { + !from.joins.is_empty() + || matches!( + from.relation, + TableFactor::Derived { .. } | TableFactor::Function { .. } + ) + }) { return ControlFlow::Break(()); } - self.is_simple_count = match &select.projection[0] { - SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. } => { - if let Expr::Function(func) = expr { - let name = trim_quotes(&func.name.to_string().to_lowercase()); - name == "count" - && func.filter.is_none() - && func.over.is_none() - && func.within_group.is_empty() - && matches!( - &func.args, - FunctionArguments::List(list) if list.args.len() == 1 && trim_quotes(&list.args[0].to_string()) == "*" - ) - } else { - false + // 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, } - _ => 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_histogram = true; + } } - if !self.is_simple_count { + if self.is_simple_count || self.is_simple_histogram { return ControlFlow::Break(()); } ControlFlow::Continue(()) @@ -1254,6 +1299,7 @@ impl VisitorMut for SimpleCountVisitor { Expr::Subquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. } ) { self.is_simple_count = false; + self.is_simple_histogram = false; return ControlFlow::Break(()); } ControlFlow::Continue(()) @@ -1262,7 +1308,7 @@ impl VisitorMut for SimpleCountVisitor { fn is_complex_query(statement: &mut Statement) -> bool { let mut visitor = ComplexQueryVisitor::new(); - statement.visit(&mut visitor); + let _ = statement.visit(&mut visitor); visitor.is_complex } @@ -1344,12 +1390,13 @@ impl VisitorMut for ComplexQueryVisitor { } } -struct HistogramIntervalVistor { +#[derive(Debug)] +struct HistogramIntervalVisitor { pub interval: Option, time_range: Option<(i64, i64)>, } -impl HistogramIntervalVistor { +impl HistogramIntervalVisitor { fn new(time_range: Option<(i64, i64)>) -> Self { Self { interval: None, @@ -1358,7 +1405,7 @@ impl HistogramIntervalVistor { } } -impl VisitorMut for HistogramIntervalVistor { +impl VisitorMut for HistogramIntervalVisitor { type Break = (); fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { @@ -1764,7 +1811,7 @@ pub fn check_or_add_order_by_timestamp(sql: &str, is_asc: bool) -> infra::errors return Ok(sql.to_string()); } let mut visitor = AddOrderingTermVisitor::new(TIMESTAMP_COL_NAME.to_string(), is_asc); - statement.visit(&mut visitor); + let _ = statement.visit(&mut visitor); Ok(statement.to_string()) } @@ -1814,7 +1861,7 @@ pub fn add_new_filters_with_and_operator( return Ok(sql.to_string()); } let mut visitor = AddNewFiltersWithAndOperatorVisitor::new(filters); - statement.visit(&mut visitor); + let _ = statement.visit(&mut visitor); Ok(statement.to_string()) } @@ -1903,7 +1950,7 @@ mod tests { let mut index_fields = HashSet::new(); index_fields.insert("name".to_string()); let mut index_visitor = IndexVisitor::new_from_index_fields(index_fields, true, true); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); let expected = "name=a AND (name=b OR (_all:good AND _all:bar))"; let expected_sql = "SELECT * FROM t WHERE age = 1 AND (match_all('foo') OR age = 2)"; assert_eq!( @@ -1923,7 +1970,7 @@ mod tests { let mut index_fields = HashSet::new(); index_fields.insert("name".to_string()); let mut index_visitor = IndexVisitor::new_from_index_fields(index_fields, true, true); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); let expected = ""; let expected_sql = "SELECT * FROM t WHERE name IS NOT NULL AND age > 1 AND (match_all('foo') OR abs(age) = 2)"; assert_eq!( @@ -1947,7 +1994,7 @@ mod tests { let mut index_fields = HashSet::new(); index_fields.insert("name".to_string()); let mut index_visitor = IndexVisitor::new_from_index_fields(index_fields, true, true); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); let expected = ""; let expected_sql = "SELECT * FROM t WHERE (name = 'b' OR (match_all('good') AND match_all('bar'))) OR (match_all('foo') OR age = 2)"; assert_eq!( @@ -1971,7 +2018,7 @@ mod tests { let mut index_fields = HashSet::new(); index_fields.insert("name".to_string()); let mut index_visitor = IndexVisitor::new_from_index_fields(index_fields, true, true); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); let expected = "((name=b OR (_all:good AND _all:bar)) OR (_all:foo AND name=c))"; let expected_sql = "SELECT * FROM t"; assert_eq!( @@ -1991,7 +2038,7 @@ mod tests { let mut index_fields = HashSet::new(); index_fields.insert("name".to_string()); let mut index_visitor = IndexVisitor::new_from_index_fields(index_fields, true, true); - statement.visit(&mut index_visitor); + let _ = statement.visit(&mut index_visitor); let expected = "((_all:good AND _all:bar) OR (_all:foo AND name=c))"; let expected_sql = "SELECT * FROM t WHERE (foo = 'b' OR foo = 'c') AND foo = 'd'"; assert_eq!( @@ -2009,7 +2056,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM t WHERE name = 'a'"; assert_eq!(statement.to_string(), expected_sql); } @@ -2022,7 +2069,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM t WHERE name = 'a'"; assert_eq!(statement.to_string(), expected_sql); } @@ -2035,7 +2082,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM t1 JOIN t2 ON t1.name = t2.name WHERE t1.name = 'openobserve'"; assert_eq!(statement.to_string(), expected_sql); } @@ -2048,7 +2095,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM t1 WHERE name NOT IN (SELECT name FROM t2)"; assert_eq!(statement.to_string(), expected_sql); @@ -2062,7 +2109,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM (SELECT name FROM t1 UNION SELECT name FROM t2)"; assert_eq!(statement.to_string(), expected_sql); @@ -2076,7 +2123,7 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM ((SELECT name FROM t1) UNION (SELECT name FROM t2))"; assert_eq!(statement.to_string(), expected_sql); } @@ -2089,11 +2136,17 @@ mod tests { .pop() .unwrap(); let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); - statement.visit(&mut track_total_hits_visitor); + let _ = statement.visit(&mut track_total_hits_visitor); let expected_sql = "SELECT count(*) AS zo_sql_num FROM (SELECT name FROM t1 UNION SELECT name FROM t2 UNION SELECT name FROM t3)"; assert_eq!(statement.to_string(), expected_sql); } + fn is_simple_count_query(statement: &mut Statement) -> bool { + let mut visitor = OtherIndexOptimizeModeVisitor::new(); + let _ = statement.visit(&mut visitor); + visitor.is_simple_count + } + #[test] fn test_is_simple_count_visit1() { let sql = "SELECT count(*) from t"; @@ -2154,6 +2207,77 @@ mod tests { assert_eq!(is_simple_count_query(&mut statement), false); } + fn is_simple_histogram_query(statement: &mut Statement) -> bool { + let mut visitor = OtherIndexOptimizeModeVisitor::new(); + let _ = statement.visit(&mut visitor); + visitor.is_simple_histogram + } + + #[test] + fn test_is_simple_histogram_visit1() { + let sql = "select histogram(_timestamp, '10 second') AS zo_sql_key, count(*) AS zo_sql_num from \"default\" GROUP BY zo_sql_key ORDER BY zo_sql_key"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), true); + } + + #[test] + fn test_is_simple_histogram_visit2() { + // Test with additional where clause + let sql = "select histogram(_timestamp, '1m') as h, count(*) from t where name = 'test'"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), true); + } + + #[test] + fn test_is_simple_histogram_visit3() { + // Test with wrong order of projections (count before histogram) + let sql = "select count(*), histogram(_timestamp, '1m') from t"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), false); + } + + #[test] + fn test_is_simple_histogram_visit4() { + // Test with subquery - should fail + let sql = "select histogram(_timestamp, '1m'), count(*) from (select * from t)"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), false); + } + + #[test] + fn test_is_simple_histogram_visit5() { + // Test with additional projection - should fail + let sql = "select histogram(_timestamp, '1m'), count(*), name from t"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), false); + } + + #[test] + fn test_is_simple_histogram_visit6() { + // Test with join - should fail + let sql = "select histogram(_timestamp, '1m'), count(*) from t1 join t2"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + assert_eq!(is_simple_histogram_query(&mut statement), false); + } + #[test] fn test_check_or_add_order_by_timestamp_no_order_asc() { let sql = "SELECT * FROM logs"; diff --git a/src/service/search/tantivy/puffin_directory/mod.rs b/src/service/search/tantivy/puffin_directory/mod.rs index 92dfae9265..d2d1084c78 100644 --- a/src/service/search/tantivy/puffin_directory/mod.rs +++ b/src/service/search/tantivy/puffin_directory/mod.rs @@ -30,11 +30,14 @@ pub mod reader; pub mod reader_cache; pub mod writer; -// We do not need all of the tantivy files, only the .term and .idx files -// for getting doc IDs and also the meta.json file +// We do not need all of the tantivy files, only specific ones: +// - .term and .idx files for getting doc IDs +// - .pos files for position information +// - .fast files for fast fields +// - meta.json file for index metadata // This might change in the future when we add more features to the index -const ALLOWED_FILE_EXT: &[&str] = &["term", "idx", "pos"]; -const EMPTY_FILE_EXT: &[&str] = &["fast", "fieldnorm", "store"]; +const ALLOWED_FILE_EXT: &[&str] = &["term", "idx", "pos", "fast"]; +const EMPTY_FILE_EXT: &[&str] = &["fieldnorm", "store"]; const META_JSON: &str = "meta.json"; const FOOTER_CACHE: &str = "footer_cache"; diff --git a/src/service/search/tantivy/puffin_directory/reader.rs b/src/service/search/tantivy/puffin_directory/reader.rs index 040f10190c..e3a8a25ed4 100644 --- a/src/service/search/tantivy/puffin_directory/reader.rs +++ b/src/service/search/tantivy/puffin_directory/reader.rs @@ -14,12 +14,14 @@ // along with this program. If not, see . use std::{ + collections::HashSet, io, ops::Range, path::{Path, PathBuf}, sync::Arc, }; +use config::TIMESTAMP_COL_NAME; use futures::future::try_join_all; use hashbrown::HashMap; use tantivy::{ @@ -211,10 +213,13 @@ impl Directory for PuffinDirReader { pub async fn warm_up_terms( searcher: &tantivy::Searcher, terms_grouped_by_field: &HashMap>, + need_fast_field: bool, ) -> anyhow::Result<()> { let mut warm_up_fields_futures = Vec::new(); let mut warm_up_fields_term_futures = Vec::new(); let mut warm_up_terms_futures = Vec::new(); + let mut warm_up_fast_fields_futures = Vec::new(); + let mut warmed_segments = HashSet::new(); for (field, terms) in terms_grouped_by_field { for segment_reader in searcher.segment_readers() { let inv_idx = segment_reader.inverted_index(*field)?; @@ -231,6 +236,18 @@ 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); + } + } } } if !warm_up_fields_futures.is_empty() { @@ -242,5 +259,23 @@ pub async fn warm_up_terms( if !warm_up_terms_futures.is_empty() { try_join_all(warm_up_terms_futures).await?; } + if !warm_up_fast_fields_futures.is_empty() { + try_join_all(warm_up_fast_fields_futures).await?; + } + Ok(()) +} + +async fn warm_up_fastfield( + fast_field_reader: &tantivy::fastfield::FastFieldReaders, +) -> anyhow::Result<()> { + let columns = fast_field_reader + .list_dynamic_column_handles(TIMESTAMP_COL_NAME) + .await?; + futures::future::try_join_all( + columns + .into_iter() + .map(|col| async move { col.file_slice().read_bytes_async().await }), + ) + .await?; Ok(()) }