refactor(sql): improve agg index rewrite matching

src/query/service/src/physical_plans/format/format_table_scan.rs

@@ -137,16 +137,21 @@ impl<'a> PhysicalFormat for TableScanFormatter<'a> {
 
         // Aggregating index
         if let Some(agg_index) = agg_index {
-            let (_, agg_index_sql, _) = ctx
+            let table_index = self
+                .inner
+                .table_index
+                .expect("agg index should only exist for bound table scans");
+            let index = ctx
                 .metadata
-                .get_agg_indices(&table_name)
+                .get_agg_indices(table_index)
                 .unwrap()
                 .iter()
-                .find(|(index, _, _)| *index == agg_index.index_id)
+                .find(|index| index.index_id == agg_index.index_id)
                 .unwrap();
 
             children.push(FormatTreeNode::new(format!(
-                "aggregating index: [{agg_index_sql}]"
+                "aggregating index: [{}]",
+                index.sql
             )));
 
             let agg_sel = agg_index

src/query/service/tests/it/sql/planner/optimizer/optimizers/rule/agg_rules/agg_index_query_rewrite.rs

@@ -24,6 +24,7 @@ use databend_common_meta_app::schema::CreateOption;
 use databend_common_sql::BindContext;
 use databend_common_sql::MetadataRef;
 use databend_common_sql::optimizer::OptimizerContext;
+use databend_common_sql::optimizer::build_agg_index_plan_for_table;
 use databend_common_sql::optimizer::ir::SExpr;
 use databend_common_sql::optimizer::optimizers::recursive::RecursiveRuleOptimizer;
 use databend_common_sql::optimizer::optimizers::rule::DEFAULT_REWRITE_RULES;
@@ -227,6 +228,13 @@ fn get_test_suites() -> Vec<TestSuite> {
             index_selection: vec!["index_col_0 (#0)", "index_col_1 (#1)"],
             rewritten_predicates: vec![],
         },
+        TestSuite {
+            query: "select avg(a) from t",
+            index: "select count(), count(a), sum(a), sum(b) from t",
+            is_matched: true,
+            index_selection: vec!["index_col_2 (#2)", "index_col_1 (#1)"],
+            rewritten_predicates: vec![],
+        },
         // query: eval-agg-eval-filter-scan, index: eval-agg-eval-scan
         TestSuite {
             query: "select sum(a) + 1 from t where b > 1 group by b",
@@ -378,18 +386,32 @@ async fn test_query_rewrite_impl(format: &str) -> Result<()> {
     let test_suites = get_test_suites();
     for suite in test_suites.into_iter() {
         let (index, _, _) = plan_sql(ctx.clone(), suite.index, false).await?;
-        let (mut query, _, metadata) = plan_sql(ctx.clone(), suite.query, true).await?;
+        let (mut query, _, metadata_ref) = plan_sql(ctx.clone(), suite.query, true).await?;
+        let table_index = {
+            let metadata = metadata_ref.read();
+            metadata
+                .tables()
+                .iter()
+                .find(|table| table.name() == "t")
+                .map(|table| table.index())
+                .expect("query metadata should contain table t")
+        };
+        let agg_index_plan = build_agg_index_plan_for_table(
+            ctx.clone(),
+            None,
+            metadata_ref.clone(),
+            table_index,
+            0,
+            suite.index.to_string(),
+            index,
+        )?;
         {
-            let mut metadata = metadata.write();
-            metadata.add_agg_indices("default.default.t".to_string(), vec![(
-                0,
-                suite.index.to_string(),
-                index,
-            )]);
+            let mut metadata = metadata_ref.write();
+            metadata.add_agg_indices(table_index, vec![agg_index_plan]);
         }
         query.clear_applied_rules();
 
-        let opt_ctx = OptimizerContext::new(ctx.clone(), metadata.clone());
+        let opt_ctx = OptimizerContext::new(ctx.clone(), metadata_ref.clone());
         let result = RecursiveRuleOptimizer::new(opt_ctx.clone(), &[RuleID::TryApplyAggIndex])
             .optimize_sync(&query)?;
         let agg_index = find_push_down_index_info(&result)?;

src/query/sql/src/planner/binder/bind_table_reference/bind_table.rs

@@ -136,6 +136,38 @@ impl Binder {
 
         let navigation = self.resolve_temporal_clause(bind_context, temporal)?;
 
+        if bind_context.planning_agg_index {
+            let source_table_index = {
+                let metadata = self.metadata.read();
+                metadata
+                    .tables()
+                    .iter()
+                    .find(|table| {
+                        table.is_source_of_index()
+                            && table.catalog() == catalog
+                            && table.database() == database
+                            && table.name() == table_name
+                    })
+                    .map(|table| table.index())
+            };
+            if let Some(table_index) = source_table_index {
+                let (s_expr, mut bind_context) = self.bind_base_table(
+                    bind_context,
+                    database.as_str(),
+                    table_index,
+                    None,
+                    sample,
+                    true,
+                    true,
+                )?;
+
+                if let Some(alias) = alias {
+                    bind_context.apply_table_alias(alias, &self.name_resolution_ctx)?;
+                }
+                return Ok((s_expr, bind_context));
+            }
+        }
+
         // Resolve table with catalog
         let table_meta = {
             let table_name = if let Some(cte_suffix_name) = cte_suffix_name.as_ref() {
@@ -206,6 +238,7 @@ impl Binder {
                     change_type,
                     sample,
                     true,
+                    false,
                 )?;
 
                 if let Some(alias) = alias {
@@ -327,6 +360,7 @@ impl Binder {
                     None,
                     sample,
                     true,
+                    false,
                 )?;
                 if let Some(alias) = alias {
                     bind_context.apply_table_alias(alias, &self.name_resolution_ctx)?;

src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs

@@ -309,8 +309,15 @@ impl Binder {
             false,
             None,
         );
-        let (s_expr, mut bind_context) =
-            self.bind_base_table(bind_context, "system", table_index, None, sample, true)?;
+        let (s_expr, mut bind_context) = self.bind_base_table(
+            bind_context,
+            "system",
+            table_index,
+            None,
+            sample,
+            true,
+            false,
+        )?;
         if let Some(alias) = alias {
             bind_context.apply_table_alias(alias, &self.name_resolution_ctx)?;
         }
@@ -371,8 +378,15 @@ impl Binder {
                 None,
             );
 
-            let (s_expr, mut bind_context) =
-                self.bind_base_table(bind_context, "system", table_index, None, &None, true)?;
+            let (s_expr, mut bind_context) = self.bind_base_table(
+                bind_context,
+                "system",
+                table_index,
+                None,
+                &None,
+                true,
+                false,
+            )?;
             if let Some(alias) = alias {
                 bind_context.apply_table_alias(alias, &self.name_resolution_ctx)?;
             }