refactor: simplify parameterized index scan runtime probe flow

Author: KKould

src/db.rs

@@ -463,12 +463,12 @@ impl<S: Storage> State<S> {
         &self.view_cache
     }
 
-    fn build_plan_and_runtime_params<A: AsRef<[(&'static str, DataValue)]>>(
+    fn build_plan<A: AsRef<[(&'static str, DataValue)]>>(
         &self,
         stmt: &Statement,
         params: A,
         transaction: &<S as Storage>::TransactionType<'_>,
-    ) -> Result<(LogicalPlan, usize), DatabaseError> {
+    ) -> Result<LogicalPlan, DatabaseError> {
         let mut binder = Binder::new(
             BinderContext::new(
                 self.table_cache(),
@@ -491,7 +491,6 @@ impl<S: Storage> State<S> {
         let source_plan = binder.bind(stmt)?;
         let mut optimizer = self.optimizer_pipeline.instantiate(source_plan);
         optimizer.optimize(Some(&transaction.meta_loader(self.meta_cache())))?;
-        let runtime_param_count = optimizer.runtime_param_count();
         let mut best_plan = optimizer.into_plan();
 
         if let Operator::Analyze(op) = &mut best_plan.operator {
@@ -500,7 +499,7 @@ impl<S: Storage> State<S> {
             }
         }
 
-        Ok((best_plan, runtime_param_count))
+        Ok(best_plan)
     }
 
     fn execute<'a, 'txn, A: AsRef<[(&'static str, DataValue)]>>(
@@ -512,11 +511,9 @@ impl<S: Storage> State<S> {
     where
         S: 'txn,
     {
-        let (mut plan, runtime_param_count) =
-            self.build_plan_and_runtime_params(stmt, params, transaction)?;
+        let mut plan = self.build_plan(stmt, params, transaction)?;
         let schema = plan.output_schema().clone();
         let mut arena = ExecArena::default();
-        arena.init_runtime_params(runtime_param_count);
         let root = build_write(
             &mut arena,
             plan,
@@ -1165,10 +1162,7 @@ pub(crate) mod test {
             None,
         );
         let source_plan = binder.bind(&stmt)?;
-        let (best_plan, _) =
-            kite_sql
-                .state
-                .build_plan_and_runtime_params(&stmt, [], &transaction)?;
+        let best_plan = kite_sql.state.build_plan(&stmt, [], &transaction)?;
 
         let join_plan = match source_plan.operator {
             Operator::Project(_) => source_plan.childrens.pop_only(),
@@ -1252,10 +1246,7 @@ pub(crate) mod test {
             None,
         );
         let source_plan = binder.bind(&stmt)?;
-        let (best_plan, _) =
-            kite_sql
-                .state
-                .build_plan_and_runtime_params(&stmt, [], &transaction)?;
+        let best_plan = kite_sql.state.build_plan(&stmt, [], &transaction)?;
 
         let join_plan = match source_plan.operator {
             Operator::Project(_) => source_plan.childrens.pop_only(),
@@ -1353,10 +1344,7 @@ pub(crate) mod test {
             "SELECT o.x, t.y FROM onecolumn o INNER JOIN twocolumn t ON (o.x=t.x AND t.y=53)",
         )?;
         let transaction = kite_sql.storage.transaction()?;
-        let (best_plan, _) =
-            kite_sql
-                .state
-                .build_plan_and_runtime_params(&stmt, [], &transaction)?;
+        let best_plan = kite_sql.state.build_plan(&stmt, [], &transaction)?;
         let join_plan = match best_plan.operator {
             Operator::Project(_) => best_plan.childrens.pop_only(),
             Operator::Join(_) => best_plan,
@@ -1556,7 +1544,7 @@ pub(crate) mod test {
                     "unexpected explain plan: {explain_plan}"
                 );
                 assert!(
-                    explain_plan.contains(&format!("IndexScan By {index_name} => Probe $0")),
+                    explain_plan.contains(&format!("IndexScan By {index_name} => Probe")),
                     "unexpected explain plan: {explain_plan}"
                 );
                 Ok(())
@@ -1727,7 +1715,7 @@ pub(crate) mod test {
                 "unexpected explain plan: {explain_plan}"
             );
             assert!(
-                explain_plan.contains("IndexScan By exists_inner_v_index => Probe $0"),
+                explain_plan.contains("IndexScan By exists_inner_v_index => Probe"),
                 "unexpected explain plan: {explain_plan}"
             );
             Ok(())

src/execution/dql/index_scan.rs

@@ -16,35 +16,17 @@ use crate::errors::DatabaseError;
 use crate::execution::{ExecArena, ExecId, ExecNode, ExecutionCaches, ExecutorNode, ReadExecutor};
 use crate::expression::range_detacher::Range;
 use crate::planner::operator::table_scan::TableScanOperator;
-use crate::storage::{IndexIter, Iter, Transaction};
+use crate::storage::{IndexIter, IndexRanges, Iter, Transaction};
 use crate::types::index::{IndexLookup, IndexMetaRef, RuntimeIndexProbe};
 use crate::types::serialize::TupleValueSerializableImpl;
-use std::array;
-use std::vec;
-
-enum IndexLookupRanges {
-    One(array::IntoIter<Range, 1>),
-    Many(vec::IntoIter<Range>),
-}
-
-impl Iterator for IndexLookupRanges {
-    type Item = Range;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match self {
-            IndexLookupRanges::One(iter) => iter.next(),
-            IndexLookupRanges::Many(iter) => iter.next(),
-        }
-    }
-}
 
 pub(crate) struct IndexScan<'a, T: Transaction + 'a> {
     op: Option<TableScanOperator>,
     index_by: IndexMetaRef,
     lookup: Option<IndexLookup>,
     covered_deserializers: Option<Vec<TupleValueSerializableImpl>>,
     cover_mapping: Option<Vec<usize>>,
-    iter: Option<IndexIter<'a, T, IndexLookupRanges>>,
+    iter: Option<IndexIter<'a, T>>,
 }
 
 impl<'a, T: Transaction + 'a>
@@ -111,23 +93,13 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for IndexScan<'a, T> {
 }
 
 impl<'a, T: Transaction + 'a> IndexScan<'a, T> {
-    fn ranges_from_lookup(lookup: IndexLookup, arena: &ExecArena<'a, T>) -> IndexLookupRanges {
+    fn ranges_from_lookup(lookup: IndexLookup, arena: &mut ExecArena<'a, T>) -> IndexRanges {
         match lookup {
-            IndexLookup::Static(Range::SortedRanges(ranges)) => {
-                IndexLookupRanges::Many(ranges.into_iter())
-            }
-            IndexLookup::Static(range) => IndexLookupRanges::One([range].into_iter()),
-            IndexLookup::Probe(param) => match arena.runtime_param(param) {
-                RuntimeIndexProbe::Eq(value) => {
-                    IndexLookupRanges::One([Range::Eq(value.clone())].into_iter())
-                }
-                RuntimeIndexProbe::Scope { min, max } => IndexLookupRanges::One(
-                    [Range::Scope {
-                        min: min.clone(),
-                        max: max.clone(),
-                    }]
-                    .into_iter(),
-                ),
+            IndexLookup::Static(Range::SortedRanges(ranges)) => ranges.into(),
+            IndexLookup::Static(range) => range.into(),
+            IndexLookup::Probe => match arena.pop_runtime_probe() {
+                RuntimeIndexProbe::Eq(value) => Range::Eq(value).into(),
+                RuntimeIndexProbe::Scope { min, max } => Range::Scope { min, max }.into(),
             },
         }
     }