From c2b7b2c728e9ad985cf5dc2f974ac2b02cafda32 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Tue, 23 Sep 2025 18:22:57 +0800 Subject: [PATCH] feat: support index optimizer for join/subuqery in multi stream (#8575) --- .../distributed_plan/remote_scan.rs | 9 + .../search/datafusion/optimizer/mod.rs | 20 +- .../index_optimizer/count.rs | 16 +- .../index_optimizer/distinct.rs | 19 +- .../index_optimizer/histogram.rs | 18 +- .../physical_optimizer/index_optimizer/mod.rs | 362 +++++++++++++++++- .../index_optimizer/select.rs | 22 +- .../index_optimizer/topn.rs | 19 +- .../index_optimizer/utils.rs | 65 ++++ .../physical_optimizer/remote_scan.rs | 35 +- src/service/search/sql/rewriter/index.rs | 5 +- 11 files changed, 486 insertions(+), 104 deletions(-) diff --git a/src/service/search/datafusion/distributed_plan/remote_scan.rs b/src/service/search/datafusion/distributed_plan/remote_scan.rs index c5793f2ece..85fe5df6ea 100644 --- a/src/service/search/datafusion/distributed_plan/remote_scan.rs +++ b/src/service/search/datafusion/distributed_plan/remote_scan.rs @@ -162,6 +162,15 @@ impl RemoteScanExec { pub fn analyze(&self) -> bool { self.remote_scan_node.search_infos.is_analyze } + + #[cfg(test)] + pub fn index_optimize_mode(&self) -> Option { + self.remote_scan_node + .index_info + .index_optimize_mode + .clone() + .map(|x| x.into()) + } } impl DisplayAs for RemoteScanExec { diff --git a/src/service/search/datafusion/optimizer/mod.rs b/src/service/search/datafusion/optimizer/mod.rs index f86efedfcf..c17a647391 100644 --- a/src/service/search/datafusion/optimizer/mod.rs +++ b/src/service/search/datafusion/optimizer/mod.rs @@ -13,7 +13,10 @@ // 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 std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use config::{ALL_VALUES_COL_NAME, ORIGINAL_DATA_COL_NAME, datafusion::request::Request}; use datafusion::{ @@ -36,6 +39,7 @@ use datafusion::{ physical_optimizer::{PhysicalOptimizerRule, limit_pushdown::LimitPushdown}, physical_plan::ExecutionPlan, prelude::SessionContext, + sql::TableReference, }; use infra::schema::get_stream_setting_index_fields; #[cfg(feature = "enterprise")] @@ -194,12 +198,11 @@ pub fn generate_physical_optimizer_rules( } // should after remote scan - if sql.stream_names.len() == 1 { - let stream_name = &sql.stream_names[0]; - let schema = sql.schemas.get(stream_name).unwrap(); + let mut index_fields: HashMap> = HashMap::new(); + for (stream_name, schema) in sql.schemas.iter() { let stream_settings = infra::schema::unwrap_stream_settings(schema.schema()); - let index_fields = get_stream_setting_index_fields(&stream_settings); - let index_fields = index_fields + let idx_fields = get_stream_setting_index_fields(&stream_settings); + let idx_fields = idx_fields .into_iter() .filter_map(|index_field| { if schema.contains_field(&index_field) { @@ -208,9 +211,10 @@ pub fn generate_physical_optimizer_rules( None } }) - .collect(); - rules.push(Arc::new(LeaderIndexOptimizerRule::new(index_fields)) as _); + .collect::>(); + index_fields.insert(stream_name.clone(), idx_fields); } + rules.push(Arc::new(LeaderIndexOptimizerRule::new(index_fields)) as _); rules.push(Arc::new(LimitPushdown::new()) as _); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/count.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/count.rs index 5764bd4ffa..4a6fc72877 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/count.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/count.rs @@ -24,6 +24,8 @@ use datafusion::{ physical_plan::{ExecutionPlan, aggregates::AggregateExec}, }; +use crate::service::search::datafusion::optimizer::physical_optimizer::index_optimizer::utils::is_complex_plan; + #[rustfmt::skip] /// check if the plan is like: /// select count(*) from stream @@ -75,19 +77,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleCountVisitor { self.is_simple_count = false; return Ok(TreeNodeRecursion::Stop); } - } else if node.name() == "HashJoinExec" - || node.name() == "RecursiveQueryExec" - || node.name() == "UnionExec" - || node.name() == "InterleaveExec" - || node.name() == "UnnestExec" - || node.name() == "CrossJoinExec" - || node.name() == "NestedLoopJoinExec" - || node.name() == "SymmetricHashJoinExec" - || node.name() == "SortMergeJoinExec" - || node.name() == "PartialSortExec" - || node.name() == "BoundedWindowAggExec" - || node.name() == "WindowAggExec" - { + } else if is_complex_plan(node) { // if encounter complex plan, stop visiting self.is_simple_count = false; return Ok(TreeNodeRecursion::Stop); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/distinct.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/distinct.rs index 1be46ab85f..c1276c4da9 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/distinct.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/distinct.rs @@ -28,8 +28,9 @@ use datafusion::{ }, }; -use crate::service::search::datafusion::optimizer::physical_optimizer::utils::{ - get_column_name, is_column, is_only_timestamp_filter, +use crate::service::search::datafusion::optimizer::physical_optimizer::{ + index_optimizer::utils::is_complex_plan, + utils::{get_column_name, is_column, is_only_timestamp_filter}, }; #[rustfmt::skip] @@ -144,19 +145,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleDistinctVisitor { // If projection doesn't have exactly 2 expressions, stop visiting self.simple_distinct = None; return Ok(TreeNodeRecursion::Stop); - } else if node.name() == "HashJoinExec" - || node.name() == "RecursiveQueryExec" - || node.name() == "UnionExec" - || node.name() == "InterleaveExec" - || node.name() == "UnnestExec" - || node.name() == "CrossJoinExec" - || node.name() == "NestedLoopJoinExec" - || node.name() == "SymmetricHashJoinExec" - || node.name() == "SortMergeJoinExec" - || node.name() == "PartialSortExec" - || node.name() == "BoundedWindowAggExec" - || node.name() == "WindowAggExec" - { + } else if is_complex_plan(node) { // if encounter complex plan, stop visiting self.simple_distinct = None; return Ok(TreeNodeRecursion::Stop); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/histogram.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/histogram.rs index 559630d100..2968d33dd2 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/histogram.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/histogram.rs @@ -29,7 +29,9 @@ use datafusion::{ scalar::ScalarValue, }; -use crate::service::search::datafusion::optimizer::physical_optimizer::utils::get_column_name; +use crate::service::search::datafusion::optimizer::physical_optimizer::{ + index_optimizer::utils::is_complex_plan, utils::get_column_name, +}; #[rustfmt::skip] /// SimpleHistogram(i64, u64, usize): select histogram(_timestamp, '1m') as ts, count(*) as cnt from table where match_all() group by ts; @@ -118,19 +120,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleHistogramVisitor { // If projection doesn't have exactly 2 expressions, stop visiting self.simple_histogram = None; return Ok(TreeNodeRecursion::Stop); - } else if node.name() == "HashJoinExec" - || node.name() == "RecursiveQueryExec" - || node.name() == "UnionExec" - || node.name() == "InterleaveExec" - || node.name() == "UnnestExec" - || node.name() == "CrossJoinExec" - || node.name() == "NestedLoopJoinExec" - || node.name() == "SymmetricHashJoinExec" - || node.name() == "SortMergeJoinExec" - || node.name() == "PartialSortExec" - || node.name() == "BoundedWindowAggExec" - || node.name() == "WindowAggExec" - { + } else if is_complex_plan(node) { // If encounter complex plan, stop visiting self.simple_histogram = None; return Ok(TreeNodeRecursion::Stop); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs index 0e7bce9788..5a4bb78b4b 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs @@ -13,13 +13,16 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{collections::HashSet, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use config::meta::inverted_index::IndexOptimizeMode; use datafusion::{ common::{ Result, - tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, + tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor}, }, config::ConfigOptions, physical_optimizer::PhysicalOptimizerRule, @@ -27,6 +30,7 @@ use datafusion::{ ExecutionPlan, aggregates::AggregateExec, sorts::sort_preserving_merge::SortPreservingMergeExec, }, + sql::TableReference, }; use parking_lot::Mutex; @@ -38,10 +42,10 @@ mod topn; mod utils; use crate::service::search::datafusion::{ - distributed_plan::remote_scan::RemoteScanExec, + distributed_plan::{empty_exec::NewEmptyExec, remote_scan::RemoteScanExec}, optimizer::physical_optimizer::index_optimizer::{ count::is_simple_count, distinct::is_simple_distinct, histogram::is_simple_histogram, - select::is_simple_select, topn::is_simple_topn, + select::is_simple_select, topn::is_simple_topn, utils::is_complex_plan, }, }; @@ -113,6 +117,10 @@ impl TreeNodeRewriter for FollowerIndexOptimizer { type Node = Arc; fn f_up(&mut self, plan: Self::Node) -> Result> { + if is_complex_plan(&plan) { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + } + if plan .as_any() .downcast_ref::() @@ -148,11 +156,11 @@ impl TreeNodeRewriter for FollowerIndexOptimizer { /// like order and limit #[derive(Default, Debug)] pub struct LeaderIndexOptimizerRule { - index_fields: HashSet, + index_fields: HashMap>, } impl LeaderIndexOptimizerRule { - pub fn new(index_fields: HashSet) -> Self { + pub fn new(index_fields: HashMap>) -> Self { Self { index_fields } } } @@ -182,11 +190,11 @@ impl PhysicalOptimizerRule for LeaderIndexOptimizerRule { } struct LeaderIndexOptimizer { - index_fields: HashSet, + index_fields: HashMap>, } impl LeaderIndexOptimizer { - pub fn new(index_fields: HashSet) -> Self { + pub fn new(index_fields: HashMap>) -> Self { Self { index_fields } } } @@ -195,28 +203,45 @@ impl TreeNodeRewriter for LeaderIndexOptimizer { type Node = Arc; fn f_up(&mut self, plan: Self::Node) -> Result> { + if is_complex_plan(&plan) { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + } + if plan .as_any() .downcast_ref::() .is_some() { - // Check for SimpleTopN first + // Get the index fields of the underlying table + let mut visitor = TableNameVisitor::new(); + plan.visit(&mut visitor)?; + let Some(table_name) = visitor.table_name else { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + }; + let index_fields = self + .index_fields + .get(&table_name) + .cloned() + .unwrap_or(HashSet::new()); + + // check if the query is simple topn or simple distinct if let Some(index_optimize_mode) = - is_simple_topn(Arc::clone(&plan), self.index_fields.clone()) + is_simple_topn(Arc::clone(&plan), index_fields.clone()) { + // Check for SimpleTopN let mut rewriter = IndexOptimizerRewrite::new(index_optimize_mode); let plan = plan.rewrite(&mut rewriter)?.data; return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop)); - } - // Check for SimpleDistinct - if let Some(index_optimize_mode) = - is_simple_distinct(Arc::clone(&plan), self.index_fields.clone()) + } else if let Some(index_optimize_mode) = + is_simple_distinct(Arc::clone(&plan), index_fields.clone()) { + // Check for SimpleDistinct let mut rewriter = IndexOptimizerRewrite::new(index_optimize_mode); let plan = plan.rewrite(&mut rewriter)?.data; return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop)); + } else { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Continue)); } - return Ok(Transformed::new(plan, false, TreeNodeRecursion::Continue)); } Ok(Transformed::no(plan)) } @@ -250,3 +275,310 @@ impl TreeNodeRewriter for IndexOptimizerRewrite { Ok(Transformed::no(node)) } } + +// visit physical plan to get underlying table name +struct TableNameVisitor { + table_name: Option, +} + +impl TableNameVisitor { + pub fn new() -> Self { + Self { table_name: None } + } +} + +impl<'n> TreeNodeVisitor<'n> for TableNameVisitor { + type Node = Arc; + + fn f_up(&mut self, node: &'n Self::Node) -> Result { + let name = node.name(); + if name == "NewEmptyExec" { + let table = node.as_any().downcast_ref::().unwrap(); + self.table_name = Some(TableReference::from(table.name())); + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + }; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::{ + execution::{SessionStateBuilder, runtime_env::RuntimeEnvBuilder}, + physical_plan::ExecutionPlan, + prelude::{SessionConfig, SessionContext}, + }; + + use super::*; + use crate::service::search::datafusion::{ + optimizer::physical_optimizer::{ + index_optimizer::utils::tests::get_remote_scan, remote_scan::RemoteScanRule, + }, + table_provider::empty_table::NewEmptyTable, + }; + + #[test] + fn test_table_name_visitor_extracts_table_name_from_new_empty_exec() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let plan: Arc = Arc::new(NewEmptyExec::new( + "my_table", + schema.clone().into(), + None, + &[], + None, + false, + schema.clone().into(), + )); + + let mut visitor = TableNameVisitor::new(); + plan.visit(&mut visitor).unwrap(); + let table = visitor.table_name.expect("table name should be found"); + assert_eq!( + table.to_string(), + TableReference::from("my_table").to_string() + ); + } + + #[test] + fn test_index_optimizer_rewrite_transforms_remote_scan_exec() { + // Build a minimal child plan and a default RemoteScanExec + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let child: Arc = Arc::new(NewEmptyExec::new( + "child", + schema.clone().into(), + None, + &[], + None, + false, + schema.clone().into(), + )); + + let remote_node = + crate::service::search::datafusion::distributed_plan::node::RemoteScanNode::default(); + let remote = + crate::service::search::datafusion::distributed_plan::remote_scan::RemoteScanExec::new( + Arc::clone(&child), + remote_node, + ) + .expect("construct remote scan exec"); + let plan: Arc = Arc::new(remote); + + // Apply rewrite with a concrete mode and assert it reports transformed=true + let mode = IndexOptimizeMode::SimpleTopN("field".to_string(), 10, true); + let mut rewriter = IndexOptimizerRewrite::new(mode.clone()); + let result = plan.rewrite(&mut rewriter).unwrap(); + assert!(result.transformed, "plan should be marked as transformed"); + assert_eq!(result.data.name(), "RemoteScanExec"); + let remote_scan = result + .data + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(remote_scan.index_optimize_mode(), Some(mode)); + } + + #[tokio::test] + async fn test_follower_rule_sets_mode_for_simple_select() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + + let index_optimizer_mode = Arc::new(Mutex::new(None)); + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(12)) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_physical_optimizer_rule(Arc::new(FollowerIndexOptimizerRule::new( + (0, 0), + index_optimizer_mode.clone(), + ))) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema.clone()).with_partitions(12); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let sql = "SELECT * FROM t ORDER BY _timestamp DESC LIMIT 10"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let _plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + assert_eq!( + index_optimizer_mode.lock().clone(), + Some(IndexOptimizeMode::SimpleSelect(10, false)) + ); + } + + #[tokio::test] + async fn test_follower_rule_sets_mode_for_simple_count() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + + let index_optimizer_mode = Arc::new(Mutex::new(None)); + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(12)) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_physical_optimizer_rule(Arc::new(FollowerIndexOptimizerRule::new( + (0, 0), + index_optimizer_mode.clone(), + ))) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema.clone()).with_partitions(12); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let sql = "SELECT count(*) FROM t"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let _plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + assert_eq!( + index_optimizer_mode.lock().clone(), + Some(IndexOptimizeMode::SimpleCount) + ); + } + + #[tokio::test] + async fn test_leader_rule_topn() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("id", DataType::Utf8, false), + ])); + let mut index_fields: HashMap> = HashMap::new(); + index_fields.insert( + TableReference::from("t"), + HashSet::from(["name".to_string()]), + ); + + let mut file_id_lists = hashbrown::HashMap::new(); + file_id_lists.insert(TableReference::from("t"), vec![]); + let remote_scan_rule = RemoteScanRule::new_test(file_id_lists, false); + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(12)) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_physical_optimizer_rule(Arc::new(remote_scan_rule)) + .with_physical_optimizer_rule(Arc::new(LeaderIndexOptimizerRule::new(index_fields))) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema.clone()).with_partitions(12); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let sql = "select name, count(*) as cnt from t group by name order by cnt desc limit 10"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + let remote_scan = get_remote_scan(plan); + assert_eq!( + remote_scan[0].index_optimize_mode(), + Some(IndexOptimizeMode::SimpleTopN("name".to_string(), 10, false)) + ) + } + + #[tokio::test] + async fn test_leader_rule_distinct() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("id", DataType::Utf8, false), + ])); + let mut index_fields: HashMap> = HashMap::new(); + index_fields.insert( + TableReference::from("t"), + HashSet::from(["name".to_string()]), + ); + + let mut file_id_lists = hashbrown::HashMap::new(); + file_id_lists.insert(TableReference::from("t"), vec![]); + let remote_scan_rule = RemoteScanRule::new_test(file_id_lists, false); + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(12)) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_physical_optimizer_rule(Arc::new(remote_scan_rule)) + .with_physical_optimizer_rule(Arc::new(LeaderIndexOptimizerRule::new(index_fields))) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema.clone()).with_partitions(12); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let sql = "select distinct name from t group by name order by name limit 10"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + let remote_scan = get_remote_scan(plan); + assert_eq!( + remote_scan[0].index_optimize_mode(), + Some(IndexOptimizeMode::SimpleDistinct( + "name".to_string(), + 10, + true + )) + ) + } + + #[tokio::test] + async fn test_leader_rule_subquery() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("kubernetes_namespace_name", DataType::Utf8, false), + Field::new("kubernetes_container_name", DataType::Utf8, false), + Field::new("log", DataType::Utf8, false), + ])); + let mut index_fields: HashMap> = HashMap::new(); + index_fields.insert( + TableReference::from("t"), + HashSet::from(["kubernetes_namespace_name".to_string()]), + ); + + let mut file_id_lists = hashbrown::HashMap::new(); + file_id_lists.insert(TableReference::from("t"), vec![]); + let remote_scan_rule = RemoteScanRule::new_test(file_id_lists, false); + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_target_partitions(12)) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_physical_optimizer_rule(Arc::new(remote_scan_rule)) + .with_physical_optimizer_rule(Arc::new(LeaderIndexOptimizerRule::new(index_fields))) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema.clone()).with_partitions(12); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let sql = "select kubernetes_namespace_name, + array_agg(distinct kubernetes_container_name) as container_name + from t + where log like '%zinc%' + and kubernetes_namespace_name in ( + select distinct kubernetes_namespace_name + from t + order by kubernetes_namespace_name limit 10000) + group by kubernetes_namespace_name + order by kubernetes_namespace_name + limit 10"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + let remote_scan = get_remote_scan(plan); + assert_eq!( + remote_scan[0].index_optimize_mode(), + Some(IndexOptimizeMode::SimpleDistinct( + "kubernetes_namespace_name".to_string(), + 10000, + true + )) + ) + } +} diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/select.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/select.rs index 4a51a6a1ab..a08a7625ea 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/select.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/select.rs @@ -24,8 +24,9 @@ use datafusion::{ physical_plan::{ExecutionPlan, sorts::sort_preserving_merge::SortPreservingMergeExec}, }; -use crate::service::search::datafusion::optimizer::physical_optimizer::utils::{ - get_column_name, is_column, +use crate::service::search::datafusion::optimizer::physical_optimizer::{ + index_optimizer::utils::is_complex_plan, + utils::{get_column_name, is_column}, }; #[rustfmt::skip] @@ -77,23 +78,10 @@ impl<'n> TreeNodeVisitor<'n> for SimpleSelectVisitor { return Ok(TreeNodeRecursion::Continue); } } - // if not simple distinct, stop visiting + // if not simple select, stop visiting self.is_simple_select = None; return Ok(TreeNodeRecursion::Stop); - } else if node.name() == "HashJoinExec" - || node.name() == "RecursiveQueryExec" - || node.name() == "UnionExec" - || node.name() == "InterleaveExec" - || node.name() == "UnnestExec" - || node.name() == "CrossJoinExec" - || node.name() == "NestedLoopJoinExec" - || node.name() == "SymmetricHashJoinExec" - || node.name() == "SortMergeJoinExec" - || node.name() == "PartialSortExec" - || node.name() == "BoundedWindowAggExec" - || node.name() == "WindowAggExec" - || node.name() == "AggregateExec" - { + } else if is_complex_plan(node) || node.name() == "AggregateExec" { // if encounter complex plan, stop visiting self.is_simple_select = None; return Ok(TreeNodeRecursion::Stop); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/topn.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/topn.rs index 7443f8e4e6..ba808df181 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/topn.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/topn.rs @@ -27,8 +27,9 @@ use datafusion::{ }, }; -use crate::service::search::datafusion::optimizer::physical_optimizer::utils::{ - get_column_name, is_column, +use crate::service::search::datafusion::optimizer::physical_optimizer::{ + index_optimizer::utils::is_complex_plan, + utils::{get_column_name, is_column}, }; #[rustfmt::skip] @@ -148,19 +149,7 @@ impl<'n> TreeNodeVisitor<'n> for SimpleTopnVisitor { // If AggregateExec doesn't match SimpleTopN pattern, stop visiting self.simple_topn = None; return Ok(TreeNodeRecursion::Stop); - } else if node.name() == "HashJoinExec" - || node.name() == "RecursiveQueryExec" - || node.name() == "UnionExec" - || node.name() == "InterleaveExec" - || node.name() == "UnnestExec" - || node.name() == "CrossJoinExec" - || node.name() == "NestedLoopJoinExec" - || node.name() == "SymmetricHashJoinExec" - || node.name() == "SortMergeJoinExec" - || node.name() == "PartialSortExec" - || node.name() == "BoundedWindowAggExec" - || node.name() == "WindowAggExec" - { + } else if is_complex_plan(node) { // If encounter complex plan, stop visiting self.simple_topn = None; return Ok(TreeNodeRecursion::Stop); diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/utils.rs b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/utils.rs index 6c91de3b40..9bf03dcc41 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/utils.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/index_optimizer/utils.rs @@ -13,6 +13,31 @@ // 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 datafusion::{common::tree_node::TreeNode, physical_plan::ExecutionPlan}; + +// check if the plan contains join, union, interleave, unnest, partial sort(use for streaming +// table), bounded window agg, window agg +pub fn is_complex_plan(node: &Arc) -> bool { + node.exists(|plan| { + Ok(plan.name() == "HashJoinExec" + || plan.name() == "RecursiveQueryExec" + || plan.name() == "UnionExec" + || plan.name() == "InterleaveExec" + || plan.name() == "UnnestExec" + || plan.name() == "CrossJoinExec" + || plan.name() == "NestedLoopJoinExec" + || plan.name() == "SymmetricHashJoinExec" + || plan.name() == "SortMergeJoinExec" + || plan.name() == "PartialSortExec" + || plan.name() == "BoundedWindowAggExec" + || plan.name() == "WindowAggExec" + || plan.children().len() > 1) + }) + .unwrap_or(true) +} + #[cfg(test)] pub(crate) mod tests { use std::sync::Arc; @@ -28,6 +53,8 @@ pub(crate) mod tests { }, }; + use crate::service::search::datafusion::distributed_plan::remote_scan::RemoteScanExec; + // get the first final aggregate plan from bottom to top pub(crate) fn get_partial_aggregate_plan( plan: Arc, @@ -75,4 +102,42 @@ pub(crate) mod tests { } } } + + pub(crate) fn get_remote_scan(plan: Arc) -> Vec> { + let mut visitor = RemoteScanVisitor::new(); + let _ = plan.visit(&mut visitor); + visitor.get_data() + } + + struct RemoteScanVisitor { + data: Vec>, + } + + impl RemoteScanVisitor { + fn new() -> Self { + Self { data: Vec::new() } + } + + fn get_data(&self) -> Vec> { + self.data.clone() + } + } + + impl Default for RemoteScanVisitor { + fn default() -> Self { + Self::new() + } + } + + impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor { + type Node = Arc; + + fn f_up(&mut self, node: &'n Self::Node) -> Result { + if node.name() == "RemoteScanExec" { + let remote_scan = node.as_any().downcast_ref::().unwrap(); + self.data.push(Arc::new(remote_scan.clone())); + } + Ok(TreeNodeRecursion::Continue) + } + } } diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/remote_scan.rs b/src/service/search/datafusion/optimizer/physical_optimizer/remote_scan.rs index 58c319ff0a..141bab3642 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/remote_scan.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/remote_scan.rs @@ -89,6 +89,7 @@ pub fn generate_remote_scan_rules( #[derive(Debug)] pub struct RemoteScanRule { remote_scan_nodes: Arc, + single_node_optimizer_enable: bool, } impl RemoteScanRule { @@ -110,7 +111,33 @@ impl RemoteScanRule { opentelemetry_context, ); let remote_scan_nodes = Arc::new(remote_scan_nodes); - Self { remote_scan_nodes } + Self { + remote_scan_nodes, + single_node_optimizer_enable: config::get_config() + .common + .feature_single_node_optimize_enabled, + } + } + + #[cfg(test)] + pub fn new_test( + file_id_lists: HashMap>>, + single_node_optimizer_enable: bool, + ) -> Self { + use tracing_opentelemetry::OpenTelemetrySpanExt; + let remote_scan_nodes = RemoteScanNodes::new( + Request::default(), + vec![], + file_id_lists, + HashMap::new(), + false, + tracing::Span::current().context(), + ); + let remote_scan_nodes = Arc::new(remote_scan_nodes); + Self { + remote_scan_nodes, + single_node_optimizer_enable, + } } } @@ -134,11 +161,7 @@ impl PhysicalOptimizerRule for RemoteScanRule { } // if single node and can optimize, add remote scan to top - if config::get_config() - .common - .feature_single_node_optimize_enabled - && is_single_node_optimize(&plan) - { + if self.single_node_optimizer_enable && is_single_node_optimize(&plan) { return remote_scan_to_top_if_needed(plan, self.remote_scan_nodes.clone()); } diff --git a/src/service/search/sql/rewriter/index.rs b/src/service/search/sql/rewriter/index.rs index e9977b502a..5b039fee93 100644 --- a/src/service/search/sql/rewriter/index.rs +++ b/src/service/search/sql/rewriter/index.rs @@ -81,7 +81,10 @@ impl IndexVisitor { let index_fields = if let Some((_, schema)) = schemas.iter().next() { let stream_settings = unwrap_stream_settings(schema.schema()); let index_fields = get_stream_setting_index_fields(&stream_settings); - index_fields.into_iter().collect::>() + index_fields + .into_iter() + .filter_map(|v| schema.contains_field(&v).then_some(v)) + .collect::>() } else { HashSet::new() };