implement cost model derive statistics

Author: xx01cyx

Cargo.lock

@@ -18,6 +18,7 @@ itertools = "0.13"
 assert_approx_eq = "1.1.0"
 trait-variant = "0.1.2"
 tokio = { version = "1.0.1", features = ["macros", "rt-multi-thread"] }
+async-trait = "0.1"
 
 [dev-dependencies]
 crossbeam = "0.8"

optd-cost-model/Cargo.toml

@@ -904,7 +904,8 @@ mod tests {
             expected_inner_sel
         );
         // check the outer sels
-        assert_outer_selectivities(&cost_model, expr_tree, expr_tree_rev, &attr_refs, 0.25, 0.2);
+        assert_outer_selectivities(&cost_model, expr_tree, expr_tree_rev, &attr_refs, 0.25, 0.2)
+            .await;
     }
 
     /// Non-unique oncond means the column is not unique in either table

optd-cost-model/src/cost/join/core.rs

@@ -9,7 +9,8 @@ use optd_persistent::{
 
 use crate::{
     common::{
-        nodes::{ArcPredicateNode, PhysicalNodeType},
+        nodes::{ArcPredicateNode, PhysicalNodeType, ReprPredicateNode},
+        predicates::list_pred::ListPred,
         types::{AttrId, EpochId, ExprId, TableId},
     },
     memo_ext::MemoExt,
@@ -40,28 +41,83 @@ impl<S: CostModelStorageManager> CostModelImpl<S> {
     }
 }
 
+#[async_trait::async_trait]
 impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModelImpl<S> {
-    fn compute_operation_cost(
+    async fn compute_operation_cost(
         &self,
         node: &PhysicalNodeType,
         predicates: &[ArcPredicateNode],
-        children_stats: &[Option<&EstimatedStatistic>],
+        children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<Cost> {
         todo!()
     }
 
-    fn derive_statistics(
+    async fn derive_statistics(
         &self,
         node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
-        children_statistics: &[Option<&EstimatedStatistic>],
+        children_statistics: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<EstimatedStatistic> {
-        todo!()
+        match node {
+            PhysicalNodeType::PhysicalScan => {
+                let table_id = TableId(predicates[0].data.as_ref().unwrap().as_u64());
+                let row_cnt = self
+                    .storage_manager
+                    .get_table_row_count(table_id)
+                    .await?
+                    .unwrap_or(1) as f64;
+                Ok(EstimatedStatistic(row_cnt))
+            }
+            PhysicalNodeType::PhysicalEmptyRelation => Ok(EstimatedStatistic(0.01)),
+            PhysicalNodeType::PhysicalLimit => {
+                self.get_limit_row_cnt(children_statistics[0].clone(), predicates[1].clone())
+            }
+            PhysicalNodeType::PhysicalFilter => {
+                self.get_filter_row_cnt(
+                    children_statistics[0].clone(),
+                    context.group_id,
+                    predicates[0].clone(),
+                )
+                .await
+            }
+            PhysicalNodeType::PhysicalNestedLoopJoin(join_typ) => {
+                self.get_nlj_row_cnt(
+                    join_typ,
+                    context.group_id,
+                    children_statistics[0].clone(),
+                    children_statistics[1].clone(),
+                    context.children_group_ids[0],
+                    context.children_group_ids[1],
+                    predicates[0].clone(),
+                )
+                .await
+            }
+            PhysicalNodeType::PhysicalHashJoin(join_typ) => {
+                self.get_hash_join_row_cnt(
+                    join_typ,
+                    context.group_id,
+                    children_statistics[0].clone(),
+                    children_statistics[1].clone(),
+                    context.children_group_ids[0],
+                    context.children_group_ids[1],
+                    ListPred::from_pred_node(predicates[0].clone()).unwrap(),
+                    ListPred::from_pred_node(predicates[1].clone()).unwrap(),
+                )
+                .await
+            }
+            PhysicalNodeType::PhysicalAgg => {
+                self.get_agg_row_cnt(context.group_id, predicates[1].clone())
+                    .await
+            }
+            PhysicalNodeType::PhysicalSort | PhysicalNodeType::PhysicalProjection => {
+                Ok(children_statistics[0].clone())
+            }
+        }
     }
 
-    fn update_statistics(
+    async fn update_statistics(
         &self,
         stats: Vec<Stat>,
         source: String,
@@ -70,7 +126,7 @@ impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModel
         todo!()
     }
 
-    fn get_table_statistic_for_analysis(
+    async fn get_table_statistic_for_analysis(
         &self,
         table_id: TableId,
         stat_type: StatType,
@@ -79,7 +135,7 @@ impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModel
         todo!()
     }
 
-    fn get_attribute_statistic_for_analysis(
+    async fn get_attribute_statistic_for_analysis(
         &self,
         attr_ids: Vec<AttrId>,
         stat_type: StatType,
@@ -88,7 +144,7 @@ impl<S: CostModelStorageManager + Send + Sync + 'static> CostModel for CostModel
         todo!()
     }
 
-    fn get_cost_for_analysis(
+    async fn get_cost_for_analysis(
         &self,
         expr_id: ExprId,
         epoch_id: Option<EpochId>,

optd-cost-model/src/cost_model.rs

@@ -33,7 +33,7 @@ pub struct Cost(pub Vec<f64>);
 
 /// Estimated statistic calculated by the cost model.
 /// It is the estimated output row count of the targeted expression.
-#[derive(PartialEq, PartialOrd, Debug)]
+#[derive(PartialEq, PartialOrd, Clone, Debug)]
 pub struct EstimatedStatistic(pub f64);
 
 pub type CostModelResult<T> = Result<T, CostModelError>;
@@ -73,13 +73,14 @@ impl From<serde_json::Error> for CostModelError {
     }
 }
 
+#[async_trait::async_trait]
 pub trait CostModel: 'static + Send + Sync {
     /// TODO: documentation
-    fn compute_operation_cost(
+    async fn compute_operation_cost(
         &self,
         node: &PhysicalNodeType,
         predicates: &[ArcPredicateNode],
-        children_stats: &[Option<&EstimatedStatistic>],
+        children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<Cost>;
 
@@ -88,42 +89,42 @@ pub trait CostModel: 'static + Send + Sync {
     /// statistic calculated by the cost model.
     /// TODO: Consider make it a helper function, so we can store Cost in the
     /// ORM more easily.
-    fn derive_statistics(
+    async fn derive_statistics(
         &self,
         node: PhysicalNodeType,
         predicates: &[ArcPredicateNode],
-        children_stats: &[Option<&EstimatedStatistic>],
+        children_stats: &[EstimatedStatistic],
         context: ComputeCostContext,
     ) -> CostModelResult<EstimatedStatistic>;
 
     /// TODO: documentation
     /// It is for **REAL** statistic updates, not for estimated statistics.
     /// TODO: Change data from String to other types.
-    fn update_statistics(
+    async fn update_statistics(
         &self,
         stats: Vec<Stat>,
         source: String,
         data: String,
     ) -> CostModelResult<()>;
 
     /// TODO: documentation
-    fn get_table_statistic_for_analysis(
+    async fn get_table_statistic_for_analysis(
         &self,
         table_id: TableId,
         stat_type: StatType,
         epoch_id: Option<EpochId>,
     ) -> CostModelResult<Option<StatValue>>;
 
     /// TODO: documentation
-    fn get_attribute_statistic_for_analysis(
+    async fn get_attribute_statistic_for_analysis(
         &self,
         attr_ids: Vec<AttrId>,
         stat_type: StatType,
         epoch_id: Option<EpochId>,
     ) -> CostModelResult<Option<StatValue>>;
 
     /// TODO: documentation
-    fn get_cost_for_analysis(
+    async fn get_cost_for_analysis(
         &self,
         expr_id: ExprId,
         epoch_id: Option<EpochId>,