Finish compute_operation_cost

Author: lanlou1554

optd-cost-model/src/cost_model.rs

@@ -43,24 +43,111 @@ impl<S: CostModelStorageManager> CostModelImpl<S> {
 
 #[async_trait::async_trait]
 impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModelImpl<S> {
+    /// TODO: should we add epoch_id?
     async fn compute_operation_cost(
         &self,
-        node: &PhysicalNodeType,
+        node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
+        children_costs: &[Cost],
         children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<Cost> {
-        todo!()
+        let res = self.storage_manager.get_cost(context.expr_id).await;
+        if let Ok((Some(cost), _)) = res {
+            return Ok(cost);
+        };
+        let mut output_statistic = None;
+        if let Ok((_, Some(statistic))) = res {
+            output_statistic = Some(statistic);
+        };
+        let output_cost = match node {
+            PhysicalNodeType::PhysicalScan => {
+                let output_statistic_data = output_statistic.unwrap_or(
+                    self.derive_statistics(
+                        node,
+                        predicates,
+                        children_stats,
+                        context.clone(),
+                        false,
+                    )
+                    .await?,
+                );
+                output_statistic = Some(output_statistic_data.clone());
+                Cost {
+                    compute_cost: 0.0,
+                    io_cost: output_statistic_data.0,
+                }
+            }
+            PhysicalNodeType::PhysicalEmptyRelation => Cost {
+                compute_cost: 0.1,
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalLimit => Cost {
+                compute_cost: children_costs[0].compute_cost,
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalFilter => Cost {
+                // TODO: now this equation is specific to optd, and try to make this equation more general
+                compute_cost: children_costs[1].compute_cost * children_stats[0].0,
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalNestedLoopJoin(join_typ) => {
+                let child_compute_cost = children_costs[2].compute_cost;
+                Cost {
+                    compute_cost: children_stats[0].0 * children_stats[1].0 * child_compute_cost
+                        + children_stats[0].0,
+                    io_cost: 0.0,
+                }
+            }
+            // TODO: we should document that the first child is the left table, which is used to build
+            // the hash table.
+            PhysicalNodeType::PhysicalHashJoin(join_typ) => Cost {
+                compute_cost: children_stats[0].0 * 2.0 + children_stats[1].0,
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalAgg => Cost {
+                compute_cost: children_stats[0].0
+                    * (children_costs[1].compute_cost + children_costs[2].compute_cost),
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalProjection => Cost {
+                compute_cost: children_stats[0].0 * children_costs[1].compute_cost,
+                io_cost: 0.0,
+            },
+            PhysicalNodeType::PhysicalSort => Cost {
+                compute_cost: children_stats[0].0 * children_stats[0].0.ln_1p().max(1.0),
+                io_cost: 0.0,
+            },
+        };
+        let res = self
+            .storage_manager
+            .store_cost(
+                context.expr_id,
+                Some(output_cost.clone()),
+                output_statistic,
+                None,
+            )
+            .await;
+        if res.is_err() {
+            eprintln!("Failed to store output cost");
+        }
+        Ok(output_cost)
     }
 
+    /// TODO: should we add epoch_id?
     async fn derive_statistics(
         &self,
         node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
         children_statistics: &[EstimatedStatistic],
         context: ComputeCostContext,
+        store_output_statistic: bool,
     ) -> CostModelResult<EstimatedStatistic> {
-        match node {
+        let res = self.storage_manager.get_cost(context.expr_id).await;
+        if let Ok((_, Some(statistic))) = res {
+            return Ok(statistic);
+        }
+        let output_statistic = match node {
             PhysicalNodeType::PhysicalScan => {
                 let table_id = TableId(predicates[0].data.as_ref().unwrap().as_u64());
                 let row_cnt = self
@@ -114,7 +201,17 @@ impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModel
             PhysicalNodeType::PhysicalSort | PhysicalNodeType::PhysicalProjection => {
                 Ok(children_statistics[0].clone())
             }
-        }
+        }?;
+        if store_output_statistic {
+            let res = self
+                .storage_manager
+                .store_cost(context.expr_id, None, Some(output_statistic.clone()), None)
+                .await;
+            if res.is_err() {
+                eprintln!("Failed to store output statistic");
+            }
+        };
+        Ok(output_statistic)
     }
 
     async fn update_statistics(
@@ -167,3 +264,5 @@ impl<S: CostModelStorageManager> CostModelImpl<S> {
             .await
     }
 }
+
+// TODO: Add tests for `derive_statistic`` and `compute_operation_cost`.

optd-cost-model/src/lib.rs

@@ -118,23 +118,41 @@ pub trait CostModel: 'static + Send + Sync {
     /// TODO: documentation
     async fn compute_operation_cost(
         &self,
-        node: &PhysicalNodeType,
+        node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
+        children_costs: &[Cost],
         children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<Cost>;
 
     /// TODO: documentation
     /// It is for cardinality estimation. The output should be the estimated
     /// statistic calculated by the cost model.
+    /// If this method is called by `compute_operation_cost`, please set
+    /// `store_output_statistic` to `false`; if it is called by the optimizer,
+    /// please set `store_output_statistic` to `true`. Since we can store the
+    /// estimated statistic and cost by calling the ORM method once.
+    ///
+    /// TODO: I am not sure whether to introduce `store_output_statistic`, since
+    /// it add complexity to the interface, considering currently only Scan needs
+    /// the output row count to calculate the costs. So updating the database twice
+    /// seems cheap. But in the future, maybe more cost computations rely on the output
+    /// row count.
+    ///
     /// TODO: Consider make it a helper function, so we can store Cost in the
     /// ORM more easily.
+    ///
+    /// TODO: I would suggest to rename this method to `derive_row_count`, since
+    /// statistic is easily to be confused with the real statistic.
+    /// Also we need to update other places to use estimated statistic to row count,
+    /// either in this crate or in optd-persistent.
     async fn derive_statistics(
         &self,
         node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
         children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
+        store_output_statistic: bool,
     ) -> CostModelResult<EstimatedStatistic>;
 
     /// TODO: documentation

optd-cost-model/src/storage/mock.rs

@@ -82,7 +82,7 @@ impl CostModelStorageManager for CostModelStorageMockManagerImpl {
         expr_id: ExprId,
         cost: Option<Cost>,
         estimated_statistic: Option<EstimatedStatistic>,
-        epoch_id: EpochId,
+        epoch_id: Option<EpochId>,
     ) -> CostModelResult<()> {
         todo!()
     }

optd-cost-model/src/storage/mod.rs

@@ -27,6 +27,6 @@ pub trait CostModelStorageManager {
         expr_id: ExprId,
         cost: Option<Cost>,
         estimated_statistic: Option<EstimatedStatistic>,
-        epoch_id: EpochId,
+        epoch_id: Option<EpochId>,
     ) -> CostModelResult<()>;
 }

optd-cost-model/src/storage/persistent.rs

@@ -145,14 +145,14 @@ impl<S: CostModelStorageLayer + Send + Sync> CostModelStorageManager
         expr_id: ExprId,
         cost: Option<Cost>,
         estimated_statistic: Option<EstimatedStatistic>,
-        epoch_id: EpochId,
+        epoch_id: Option<EpochId>,
     ) -> CostModelResult<()> {
         self.backend_manager
             .store_cost(
                 expr_id.into(),
                 cost.map(|c| c.into()),
                 estimated_statistic.map(|x| x.into()),
-                epoch_id.into(),
+                epoch_id.map(|id| id.into()),
             )
             .await?;
         Ok(())

optd-persistent/src/cost_model/interface.rs

@@ -121,7 +121,7 @@ pub trait CostModelStorageLayer {
         expr_id: ExprId,
         cost: Option<Cost>,
         estimated_statistic: Option<f32>,
-        epoch_id: EpochId,
+        epoch_id: Option<EpochId>,
     ) -> StorageResult<()>;
 
     async fn store_expr_stats_mappings(

optd-persistent/src/cost_model/orm.rs

@@ -519,13 +519,14 @@ impl CostModelStorageLayer for BackendManager {
 
     /// This method should handle the case when the cost is already stored.
     /// The name maybe misleading, since it can also store the estimated statistic.
+    /// If epoch_id is none, we pick the latest epoch_id.
     /// TODO: documentation
     async fn store_cost(
         &self,
         physical_expression_id: ExprId,
         cost: Option<Cost>,
         estimated_statistic: Option<f32>,
-        epoch_id: EpochId,
+        epoch_id: Option<EpochId>,
     ) -> StorageResult<()> {
         assert!(cost.is_some() || estimated_statistic.is_some());
         // TODO: should we do the following checks in the production environment?
@@ -542,17 +543,32 @@ impl CostModelStorageLayer for BackendManager {
             ));
         }
         // Check if epoch_id exists in Event table
-        let epoch_exists = Event::find()
-            .filter(event::Column::EpochId.eq(epoch_id))
-            .one(&self.db)
-            .await
-            .unwrap();
-        if epoch_exists.is_none() {
-            return Err(BackendError::CostModel(
-                format!("epoch id {} not found when storing cost", epoch_id).into(),
-            ));
+        if epoch_id.is_some() {
+            let epoch_exists = Event::find()
+                .filter(event::Column::EpochId.eq(epoch_id.unwrap()))
+                .one(&self.db)
+                .await
+                .unwrap();
+            if epoch_exists.is_none() {
+                return Err(BackendError::CostModel(
+                    format!("epoch id {} not found when storing cost", epoch_id.unwrap()).into(),
+                ));
+            }
         }
 
+        let epoch_id = match epoch_id {
+            Some(id) => id,
+            None => {
+                // When init, please make sure there is at least one epoch in the Event table.
+                let latest_epoch_id = Event::find()
+                    .order_by_desc(event::Column::EpochId)
+                    .one(&self.db)
+                    .await?
+                    .unwrap();
+                latest_epoch_id.epoch_id
+            }
+        };
+
         let transaction = self.db.begin().await?;
 
         let valid_cost = PlanCost::find()
@@ -805,7 +821,7 @@ mod tests {
                     io_cost: 42.0,
                 }),
                 Some(42.0),
-                versioned_stat_res[0].epoch_id,
+                Some(versioned_stat_res[0].epoch_id),
             )
             .await
             .unwrap();
@@ -871,7 +887,7 @@ mod tests {
         assert_eq!(cost_res.len(), 1);
         assert_eq!(
             cost_res[0].cost,
-            Some(json!({"compute_cost": 42, "io_cost": 42}))
+            Some(json!({"compute_cost": 42.0, "io_cost": 42.0}))
         );
         assert_eq!(cost_res[0].epoch_id, epoch_id1);
         assert!(!cost_res[0].is_valid);
@@ -1013,7 +1029,7 @@ mod tests {
                 physical_expression_id,
                 Some(cost.clone()),
                 Some(estimated_statistic),
-                epoch_id,
+                Some(epoch_id),
             )
             .await
             .unwrap();
@@ -1036,7 +1052,7 @@ mod tests {
                 physical_expression_id,
                 None,
                 Some(estimated_statistic),
-                epoch_id,
+                None,
             )
             .await
             .unwrap();
@@ -1075,7 +1091,12 @@ mod tests {
             io_cost: 42.0,
         };
         let _ = backend_manager
-            .store_cost(physical_expression_id, Some(cost.clone()), None, epoch_id)
+            .store_cost(
+                physical_expression_id,
+                Some(cost.clone()),
+                None,
+                Some(epoch_id),
+            )
             .await;
         let costs = super::PlanCost::find()
             .all(&backend_manager.db)
@@ -1117,7 +1138,7 @@ mod tests {
                 physical_expression_id,
                 None,
                 Some(estimated_statistic),
-                epoch_id,
+                Some(epoch_id),
             )
             .await;
         let costs = super::PlanCost::find()