refactor(core): rm option around cost models (#243)

`Option` was introduced during a transition time where we thought cost
model can compute the cost solely based on the children cost but it
turned out that we need the optimizer for a few derived logical
properties. We can drop them now.

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

docs/src/cost_model.md

@@ -11,7 +11,7 @@ pub trait CostModel<T: RelNodeTyp>: 'static + Send + Sync {
             node: &T,
             data: &Option<Value>,
             children: &[Cost],
-            context: Option<RelNodeContext>,
+            context: RelNodeContext,
     ) -> Cost;
 }
 ```

optd-core/src/cascades/tasks/optimize_inputs.rs

@@ -110,12 +110,12 @@ impl OptimizeInputsTask {
                     .iter()
                     .map(|x| x.expect("child winner should always have statistics?"))
                     .collect::<Vec<_>>(),
-                Some(RelNodeContext {
+                RelNodeContext {
                     group_id,
                     expr_id: self.expr_id,
                     children_group_ids: expr.children.clone(),
-                }),
-                Some(optimizer),
+                },
+                optimizer,
             );
             optimizer.update_group_info(
                 group_id,
@@ -197,8 +197,8 @@ impl<T: NodeType, M: Memo<T>> Task<T, M> for OptimizeInputsTask {
                 &expr.typ,
                 &preds,
                 &input_statistics_ref,
-                Some(context.clone()),
-                Some(optimizer),
+                context.clone(),
+                optimizer,
             );
             let total_cost = cost.sum(&operation_cost, &input_cost);
 

optd-core/src/cost.rs

@@ -16,25 +16,29 @@ pub struct Statistics(pub Box<dyn std::any::Any + Send + Sync + 'static>);
 pub struct Cost(pub Vec<f64>);
 
 pub trait CostModel<T: NodeType, M: Memo<T>>: 'static + Send + Sync {
-    /// Compute the cost of a single operation
+    /// Compute the cost of a single operation. `RelNodeContext` might be
+    /// optional in the future when we implement physical property enforcers.
+    /// If we have not decided the winner for a child group yet, the statistics
+    /// for that group will be `None`.
     #[allow(clippy::too_many_arguments)]
     fn compute_operation_cost(
         &self,
         node: &T,
         predicates: &[ArcPredNode<T>],
         children_stats: &[Option<&Statistics>],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<T, M>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<T, M>,
     ) -> Cost;
 
-    /// Derive the statistics of a single operation
+    /// Derive the statistics of a single operation. `RelNodeContext` might be
+    /// optional in the future when we implement physical property enforcers.
     fn derive_statistics(
         &self,
         node: &T,
         predicates: &[ArcPredNode<T>],
         children_stats: &[&Statistics],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<T, M>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<T, M>,
     ) -> Statistics;
 
     fn explain_cost(&self, cost: &Cost) -> String;

optd-datafusion-repr-adv-cost/src/lib.rs

@@ -61,8 +61,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children_stats: &[Option<&Statistics>],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Cost {
         self.base_model
             .compute_operation_cost(node, predicates, children_stats, context, optimizer)
@@ -73,11 +73,9 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children_stats: &[&Statistics],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Statistics {
-        let context = context.as_ref();
-        let optimizer = optimizer.as_ref();
         let row_cnts = children_stats
             .iter()
             .map(|child| DfCostModel::row_cnt(child))
@@ -100,12 +98,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
                 DfCostModel::stat(row_cnt)
             }
             DfNodeType::PhysicalFilter => {
-                let output_schema = optimizer
-                    .unwrap()
-                    .get_schema_of(context.unwrap().group_id.into());
-                let output_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().group_id.into());
+                let output_schema = optimizer.get_schema_of(context.group_id.into());
+                let output_column_ref = optimizer.get_column_ref_of(context.group_id.into());
                 let row_cnt = self.stats.get_filter_row_cnt(
                     row_cnts[0],
                     output_schema,
@@ -115,18 +109,12 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
                 DfCostModel::stat(row_cnt)
             }
             DfNodeType::PhysicalNestedLoopJoin(join_typ) => {
-                let output_schema = optimizer
-                    .unwrap()
-                    .get_schema_of(context.unwrap().group_id.into());
-                let output_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().group_id.into());
-                let left_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().children_group_ids[0].into());
-                let right_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().children_group_ids[1].into());
+                let output_schema = optimizer.get_schema_of(context.group_id.into());
+                let output_column_ref = optimizer.get_column_ref_of(context.group_id.into());
+                let left_column_ref =
+                    optimizer.get_column_ref_of(context.children_group_ids[0].into());
+                let right_column_ref =
+                    optimizer.get_column_ref_of(context.children_group_ids[1].into());
                 let row_cnt = self.stats.get_nlj_row_cnt(
                     *join_typ,
                     row_cnts[0],
@@ -140,18 +128,12 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
                 DfCostModel::stat(row_cnt)
             }
             DfNodeType::PhysicalHashJoin(join_typ) => {
-                let output_schema = optimizer
-                    .unwrap()
-                    .get_schema_of(context.unwrap().group_id.into());
-                let output_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().group_id.into());
-                let left_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().children_group_ids[0].into());
-                let right_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().children_group_ids[1].into());
+                let output_schema = optimizer.get_schema_of(context.group_id.into());
+                let output_column_ref = optimizer.get_column_ref_of(context.group_id.into());
+                let left_column_ref =
+                    optimizer.get_column_ref_of(context.children_group_ids[0].into());
+                let right_column_ref =
+                    optimizer.get_column_ref_of(context.children_group_ids[1].into());
                 let row_cnt = self.stats.get_hash_join_row_cnt(
                     *join_typ,
                     row_cnts[0],
@@ -166,9 +148,7 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
                 DfCostModel::stat(row_cnt)
             }
             DfNodeType::PhysicalAgg => {
-                let output_column_ref = optimizer
-                    .unwrap()
-                    .get_column_ref_of(context.unwrap().group_id.into());
+                let output_column_ref = optimizer.get_column_ref_of(context.group_id.into());
                 let row_cnt = self
                     .stats
                     .get_agg_row_cnt(predicates[1].clone(), output_column_ref);
@@ -178,8 +158,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdvancedCostModel {
                 node,
                 predicates,
                 children_stats,
-                context.cloned(),
-                optimizer.copied(),
+                context,
+                optimizer,
             ),
         }
     }

optd-datafusion-repr/src/cost/adaptive_cost.rs

@@ -28,11 +28,9 @@ pub struct AdaptiveCostModel {
 }
 
 impl AdaptiveCostModel {
-    fn get_row_cnt(&self, context: &Option<RelNodeContext>) -> f64 {
+    fn get_row_cnt(&self, context: &RelNodeContext) -> f64 {
         let guard = self.runtime_row_cnt.lock().unwrap();
-        if let Some((runtime_row_cnt, iter)) =
-            guard.history.get(&context.as_ref().unwrap().group_id)
-        {
+        if let Some((runtime_row_cnt, iter)) = guard.history.get(&context.group_id) {
             if *iter + self.decay >= guard.iter_cnt {
                 return (*runtime_row_cnt).max(1) as f64;
             }
@@ -67,8 +65,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdaptiveCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children: &[Option<&Statistics>],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Cost {
         if let DfNodeType::PhysicalScan = node {
             let row_cnt = self.get_row_cnt(&context);
@@ -83,8 +81,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for AdaptiveCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children: &[&Statistics],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        context: RelNodeContext,
+        optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Statistics {
         if let DfNodeType::PhysicalScan = node {
             let row_cnt = self.get_row_cnt(&context);

optd-datafusion-repr/src/cost/base_cost.rs

@@ -89,8 +89,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for DfCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children: &[&Statistics],
-        _context: Option<RelNodeContext>,
-        _optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        _context: RelNodeContext,
+        _optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Statistics {
         match node {
             DfNodeType::PhysicalScan => {
@@ -132,8 +132,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for DfCostModel {
         node: &DfNodeType,
         predicates: &[ArcDfPredNode],
         children: &[Option<&Statistics>],
-        _context: Option<RelNodeContext>,
-        _optimizer: Option<&CascadesOptimizer<DfNodeType>>,
+        _context: RelNodeContext,
+        _optimizer: &CascadesOptimizer<DfNodeType>,
     ) -> Cost {
         let row_cnts = children
             .iter()

optd-datafusion-repr/src/testing/dummy_cost.rs

@@ -19,8 +19,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for DummyCostModel {
         _: &DfNodeType,
         _: &[ArcDfPredNode],
         _: &[Option<&Statistics>],
-        _: Option<RelNodeContext>,
-        _: Option<&CascadesOptimizer<DfNodeType>>,
+        _: RelNodeContext,
+        _: &CascadesOptimizer<DfNodeType>,
     ) -> Cost {
         Cost(vec![1.0])
     }
@@ -31,8 +31,8 @@ impl CostModel<DfNodeType, NaiveMemo<DfNodeType>> for DummyCostModel {
         _: &DfNodeType,
         _: &[ArcDfPredNode],
         _: &[&Statistics],
-        _: Option<RelNodeContext>,
-        _: Option<&CascadesOptimizer<DfNodeType>>,
+        _: RelNodeContext,
+        _: &CascadesOptimizer<DfNodeType>,
     ) -> Statistics {
         Statistics(Box::new(()))
     }