chore(batch): reuse common logical for index selection (#8531)

risingwavelabs · Mar 14, 2023 · 338ace9 · 338ace9
1 parent eb56b9f
commit 338ace9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 26 deletions.
diff --git a/src/frontend/src/optimizer/plan_node/logical_join.rs b/src/frontend/src/optimizer/plan_node/logical_join.rs
@@ -407,20 +407,10 @@ impl LogicalJoin {
             result_plan = Some(lookup_join);
         }
 
-        let required_col_idx = logical_scan.required_col_idx();
         let indexes = logical_scan.indexes();
         for index in indexes {
-            let p2s_mapping = index.primary_to_secondary_mapping();
-            if required_col_idx.iter().all(|x| p2s_mapping.contains_key(x)) {
-                // Covering index selection
-                let index_scan: PlanRef = logical_scan
-                    .to_index_scan(
-                        &index.name,
-                        index.index_table.table_desc().into(),
-                        p2s_mapping,
-                    )
-                    .into();
-
+            if let Some(index_scan) = logical_scan.to_index_scan_if_index_covered(index) {
+                let index_scan: PlanRef = index_scan.into();
                 let that = self.clone_with_left_right(self.left(), index_scan.clone());
                 let new_logical_join = logical_join.clone_with_left_right(
                     logical_join.left(),

diff --git a/src/frontend/src/optimizer/plan_node/logical_scan.rs b/src/frontend/src/optimizer/plan_node/logical_scan.rs
@@ -664,12 +664,11 @@ impl ToBatch for LogicalScan {
             if let Some(applied) = index_selection_rule.apply(new.clone().into()) {
                 if let Some(scan) = applied.as_logical_scan() {
                     // covering index
-                    return required_order.enforce_if_not_satisfies(scan.to_batch().unwrap());
+                    return required_order.enforce_if_not_satisfies(scan.to_batch()?);
                 } else if let Some(join) = applied.as_logical_join() {
                     // index lookup join
-                    return required_order.enforce_if_not_satisfies(
-                        join.index_lookup_join_to_batch_lookup_join().unwrap(),
-                    );
+                    return required_order
+                        .enforce_if_not_satisfies(join.index_lookup_join_to_batch_lookup_join()?);
                 } else {
                     unreachable!();
                 }

diff --git a/src/frontend/src/optimizer/rule/index_selection_rule.rs b/src/frontend/src/optimizer/rule/index_selection_rule.rs
@@ -106,17 +106,8 @@ impl Rule for IndexSelectionRule {
         let mut final_plan: PlanRef = logical_scan.clone().into();
         let mut min_cost = primary_cost.clone();
 
-        let required_col_idx = logical_scan.required_col_idx();
         for index in indexes {
-            let p2s_mapping = index.primary_to_secondary_mapping();
-            if required_col_idx.iter().all(|x| p2s_mapping.contains_key(x)) {
-                // covering index selection
-                let index_scan = logical_scan.to_index_scan(
-                    &index.name,
-                    index.index_table.table_desc().into(),
-                    p2s_mapping,
-                );
-
+            if let Some(index_scan) = logical_scan.to_index_scan_if_index_covered(index) {
                 let index_cost = self.estimate_table_scan_cost(
                     &index_scan,
                     TableScanIoEstimator::estimate_row_size(&index_scan),