feat: support index optimizer for join/subuqery in multi stream (#8575)

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<IndexOptimizeMode> {
+        self.remote_scan_node
+            .index_info
+            .index_optimize_mode
+            .clone()
+            .map(|x| x.into())
+    }
 }
 
 impl DisplayAs for RemoteScanExec {

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 <http://www.gnu.org/licenses/>.
 
-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<TableReference, HashSet<String>> = 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::<HashSet<_>>();
+        index_fields.insert(stream_name.clone(), idx_fields);
     }
+    rules.push(Arc::new(LeaderIndexOptimizerRule::new(index_fields)) as _);
 
     rules.push(Arc::new(LimitPushdown::new()) as _);
 

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);

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);

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);

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 <http://www.gnu.org/licenses/>.
 
-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<dyn ExecutionPlan>;
 
     fn f_up(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
+        if is_complex_plan(&plan) {
+            return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop));
+        }
+
         if plan
             .as_any()
             .downcast_ref::<SortPreservingMergeExec>()
@@ -148,11 +156,11 @@ impl TreeNodeRewriter for FollowerIndexOptimizer {
 /// like order and limit
 #[derive(Default, Debug)]
 pub struct LeaderIndexOptimizerRule {
-    index_fields: HashSet<String>,
+    index_fields: HashMap<TableReference, HashSet<String>>,
 }
 
 impl LeaderIndexOptimizerRule {
-    pub fn new(index_fields: HashSet<String>) -> Self {
+    pub fn new(index_fields: HashMap<TableReference, HashSet<String>>) -> Self {
         Self { index_fields }
     }
 }
@@ -182,11 +190,11 @@ impl PhysicalOptimizerRule for LeaderIndexOptimizerRule {
 }
 
 struct LeaderIndexOptimizer {
-    index_fields: HashSet<String>,
+    index_fields: HashMap<TableReference, HashSet<String>>,
 }
 
 impl LeaderIndexOptimizer {
-    pub fn new(index_fields: HashSet<String>) -> Self {
+    pub fn new(index_fields: HashMap<TableReference, HashSet<String>>) -> Self {
         Self { index_fields }
     }
 }
@@ -195,28 +203,45 @@ impl TreeNodeRewriter for LeaderIndexOptimizer {
     type Node = Arc<dyn ExecutionPlan>;
 
     fn f_up(&mut self, plan: Self::Node) -> Result<Transformed<Self::Node>> {
+        if is_complex_plan(&plan) {
+            return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop));
+        }
+
         if plan
             .as_any()
             .downcast_ref::<SortPreservingMergeExec>()
             .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<TableReference>,
+}
+
+impl TableNameVisitor {
+    pub fn new() -> Self {
+        Self { table_name: None }
+    }
+}
+
+impl<'n> TreeNodeVisitor<'n> for TableNameVisitor {
+    type Node = Arc<dyn ExecutionPlan>;
+
+    fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
+        let name = node.name();
+        if name == "NewEmptyExec" {
+            let table = node.as_any().downcast_ref::<NewEmptyExec>().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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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::<RemoteScanExec>()
+            .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<TableReference, HashSet<String>> = 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<TableReference, HashSet<String>> = 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<TableReference, HashSet<String>> = 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
+            ))
+        )
+    }
+}

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);

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);

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 <http://www.gnu.org/licenses/>.
 
+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<dyn ExecutionPlan>) -> 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<dyn ExecutionPlan>,
@@ -75,4 +102,42 @@ pub(crate) mod tests {
             }
         }
     }
+
+    pub(crate) fn get_remote_scan(plan: Arc<dyn ExecutionPlan>) -> Vec<Arc<RemoteScanExec>> {
+        let mut visitor = RemoteScanVisitor::new();
+        let _ = plan.visit(&mut visitor);
+        visitor.get_data()
+    }
+
+    struct RemoteScanVisitor {
+        data: Vec<Arc<RemoteScanExec>>,
+    }
+
+    impl RemoteScanVisitor {
+        fn new() -> Self {
+            Self { data: Vec::new() }
+        }
+
+        fn get_data(&self) -> Vec<Arc<RemoteScanExec>> {
+            self.data.clone()
+        }
+    }
+
+    impl Default for RemoteScanVisitor {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor {
+        type Node = Arc<dyn ExecutionPlan>;
+
+        fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
+            if node.name() == "RemoteScanExec" {
+                let remote_scan = node.as_any().downcast_ref::<RemoteScanExec>().unwrap();
+                self.data.push(Arc::new(remote_scan.clone()));
+            }
+            Ok(TreeNodeRecursion::Continue)
+        }
+    }
 }

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<RemoteScanNodes>,
+    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<TableReference, Vec<Vec<i64>>>,
+        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());
         }
 

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::<HashSet<_>>()
+            index_fields
+                .into_iter()
+                .filter_map(|v| schema.contains_field(&v).then_some(v))
+                .collect::<HashSet<_>>()
         } else {
             HashSet::new()
         };