Add get_cost & store_cost in the cost model storage layer

lanlou1554 · lanlou1554 · commit 904f275fbd8d · 2024-11-20T16:48:05.000-05:00
diff --git a/optd-cost-model/src/cost_model.rs b/optd-cost-model/src/cost_model.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     memo_ext::MemoExt,
     stats::AttributeCombValueStats,
-    storage::CostModelStorageManager,
+    storage::{self, CostModelStorageManager},
     ComputeCostContext, Cost, CostModel, CostModelResult, EstimatedStatistic, StatValue,
 };
 
diff --git a/optd-cost-model/src/lib.rs b/optd-cost-model/src/lib.rs
@@ -30,13 +30,52 @@ pub struct ComputeCostContext {
 }
 
 #[derive(Default, Clone, Debug, PartialOrd, PartialEq)]
-pub struct Cost(pub Vec<f64>);
+pub struct Cost {
+    pub compute_cost: f64,
+    pub io_cost: f64,
+}
+
+impl From<Cost> for optd_persistent::cost_model::interface::Cost {
+    fn from(c: Cost) -> optd_persistent::cost_model::interface::Cost {
+        Self {
+            compute_cost: c.compute_cost,
+            io_cost: c.io_cost,
+        }
+    }
+}
+
+impl From<optd_persistent::cost_model::interface::Cost> for Cost {
+    fn from(c: optd_persistent::cost_model::interface::Cost) -> Cost {
+        Self {
+            compute_cost: c.compute_cost,
+            io_cost: c.io_cost,
+        }
+    }
+}
 
 /// Estimated statistic calculated by the cost model.
 /// It is the estimated output row count of the targeted expression.
 #[derive(PartialEq, PartialOrd, Clone, Debug)]
 pub struct EstimatedStatistic(pub f64);
 
+impl From<EstimatedStatistic> for f32 {
+    fn from(e: EstimatedStatistic) -> f32 {
+        e.0 as f32
+    }
+}
+
+impl From<EstimatedStatistic> for f64 {
+    fn from(e: EstimatedStatistic) -> f64 {
+        e.0
+    }
+}
+
+impl From<f32> for EstimatedStatistic {
+    fn from(f: f32) -> EstimatedStatistic {
+        Self(f as f64)
+    }
+}
+
 pub type CostModelResult<T> = Result<T, CostModelError>;
 
 #[derive(Debug)]
diff --git a/optd-cost-model/src/storage/mock.rs b/optd-cost-model/src/storage/mock.rs
@@ -3,7 +3,11 @@ use std::collections::HashMap;
 
 use serde::{Deserialize, Serialize};
 
-use crate::{common::types::TableId, stats::AttributeCombValueStats, CostModelResult};
+use crate::{
+    common::types::{EpochId, ExprId, TableId},
+    stats::AttributeCombValueStats,
+    Cost, CostModelResult, EstimatedStatistic,
+};
 
 use super::CostModelStorageManager;
 
@@ -63,4 +67,23 @@ impl CostModelStorageManager for CostModelStorageMockManagerImpl {
         let table_stats = self.per_table_stats_map.get(&table_id);
         Ok(table_stats.map(|stats| stats.row_cnt))
     }
+
+    /// TODO: finish this when implementing the cost get/store tests
+    async fn get_cost(
+        &self,
+        expr_id: ExprId,
+    ) -> CostModelResult<(Option<crate::Cost>, Option<EstimatedStatistic>)> {
+        todo!()
+    }
+
+    /// TODO: finish this when implementing the cost get/store tests
+    async fn store_cost(
+        &self,
+        expr_id: ExprId,
+        cost: Option<Cost>,
+        estimated_statistic: Option<EstimatedStatistic>,
+        epoch_id: EpochId,
+    ) -> CostModelResult<()> {
+        todo!()
+    }
 }
diff --git a/optd-cost-model/src/storage/mod.rs b/optd-cost-model/src/storage/mod.rs
@@ -1,4 +1,8 @@
-use crate::{common::types::TableId, stats::AttributeCombValueStats, CostModelResult};
+use crate::{
+    common::types::{EpochId, ExprId, TableId},
+    stats::AttributeCombValueStats,
+    Cost, CostModelResult, EstimatedStatistic,
+};
 
 pub mod mock;
 pub mod persistent;
@@ -12,4 +16,17 @@ pub trait CostModelStorageManager {
     ) -> CostModelResult<Option<AttributeCombValueStats>>;
 
     async fn get_table_row_count(&self, table_id: TableId) -> CostModelResult<Option<u64>>;
+
+    async fn get_cost(
+        &self,
+        expr_id: ExprId,
+    ) -> CostModelResult<(Option<Cost>, Option<EstimatedStatistic>)>;
+
+    async fn store_cost(
+        &self,
+        expr_id: ExprId,
+        cost: Option<Cost>,
+        estimated_statistic: Option<EstimatedStatistic>,
+        epoch_id: EpochId,
+    ) -> CostModelResult<()>;
 }
diff --git a/optd-cost-model/src/storage/persistent.rs b/optd-cost-model/src/storage/persistent.rs
@@ -4,9 +4,9 @@ use std::sync::Arc;
 use optd_persistent::{cost_model::interface::StatType, CostModelStorageLayer};
 
 use crate::{
-    common::types::TableId,
+    common::types::{EpochId, ExprId, TableId},
     stats::{utilities::counter::Counter, AttributeCombValueStats, Distribution, MostCommonValues},
-    CostModelResult,
+    Cost, CostModelResult, EstimatedStatistic,
 };
 
 use super::CostModelStorageManager;
@@ -125,5 +125,38 @@ impl<S: CostModelStorageLayer + Send + Sync> CostModelStorageManager
             .transpose()?)
     }
 
+    /// TODO: The name is misleading, since we can also get the estimated statistic. We should
+    /// rename it.
+    async fn get_cost(
+        &self,
+        expr_id: ExprId,
+    ) -> CostModelResult<(Option<Cost>, Option<EstimatedStatistic>)> {
+        let (cost, estimated_statistic) = self.backend_manager.get_cost(expr_id.into()).await?;
+        Ok((
+            cost.map(|c| c.into()),
+            estimated_statistic.map(|x| x.into()),
+        ))
+    }
+
+    /// TODO: The name is misleading, since we can also get the estimated statistic. We should
+    /// rename it.
+    async fn store_cost(
+        &self,
+        expr_id: ExprId,
+        cost: Option<Cost>,
+        estimated_statistic: Option<EstimatedStatistic>,
+        epoch_id: EpochId,
+    ) -> CostModelResult<()> {
+        self.backend_manager
+            .store_cost(
+                expr_id.into(),
+                cost.map(|c| c.into()),
+                estimated_statistic.map(|x| x.into()),
+                epoch_id.into(),
+            )
+            .await?;
+        Ok(())
+    }
+
     // TODO: Support querying for a specific type of statistics.
 }
diff --git a/optd-persistent/src/bin/init.rs b/optd-persistent/src/bin/init.rs
@@ -356,7 +356,7 @@ async fn init_all_tables() -> Result<(), sea_orm::error::DbErr> {
         physical_expression_id: Set(1),
         epoch_id: Set(1),
         cost: Set(Some(json!({"compute_cost":10, "io_cost":10}))),
-        estimated_statistic: Set(Some(10)),
+        estimated_statistic: Set(Some(10.0)),
         is_valid: Set(true),
     };
     plan_cost::Entity::insert(plan_cost)
diff --git a/optd-persistent/src/cost_model/interface.rs b/optd-persistent/src/cost_model/interface.rs
@@ -89,8 +89,8 @@ pub struct Stat {
 /// TODO: documentation
 #[derive(Clone, Debug, PartialEq)]
 pub struct Cost {
-    pub compute_cost: i32,
-    pub io_cost: i32,
+    pub compute_cost: f64,
+    pub io_cost: f64,
 }
 
 #[derive(Clone, Debug)]
@@ -120,7 +120,7 @@ pub trait CostModelStorageLayer {
         &self,
         expr_id: ExprId,
         cost: Option<Cost>,
-        estimated_statistic: Option<i32>,
+        estimated_statistic: Option<f32>,
         epoch_id: EpochId,
     ) -> StorageResult<()>;
 
@@ -165,9 +165,9 @@ pub trait CostModelStorageLayer {
         &self,
         expr_id: ExprId,
         epoch_id: EpochId,
-    ) -> StorageResult<(Option<Cost>, Option<i32>)>;
+    ) -> StorageResult<(Option<Cost>, Option<f32>)>;
 
-    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<(Option<Cost>, Option<i32>)>;
+    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<(Option<Cost>, Option<f32>)>;
 
     async fn get_attribute(
         &self,
diff --git a/optd-persistent/src/cost_model/orm.rs b/optd-persistent/src/cost_model/orm.rs
@@ -474,7 +474,7 @@ impl CostModelStorageLayer for BackendManager {
         &self,
         expr_id: ExprId,
         epoch_id: EpochId,
-    ) -> StorageResult<(Option<Cost>, Option<i32>)> {
+    ) -> StorageResult<(Option<Cost>, Option<f32>)> {
         let cost = PlanCost::find()
             .filter(plan_cost::Column::PhysicalExpressionId.eq(expr_id))
             .filter(plan_cost::Column::EpochId.eq(epoch_id))
@@ -486,8 +486,8 @@ impl CostModelStorageLayer for BackendManager {
         }
 
         let real_cost = cost.as_ref().and_then(|c| c.cost.as_ref()).map(|c| Cost {
-            compute_cost: c.get("compute_cost").unwrap().as_i64().unwrap() as i32,
-            io_cost: c.get("io_cost").unwrap().as_i64().unwrap() as i32,
+            compute_cost: c.get("compute_cost").unwrap().as_f64().unwrap(),
+            io_cost: c.get("io_cost").unwrap().as_f64().unwrap(),
         });
 
         Ok((real_cost, cost.unwrap().estimated_statistic))
@@ -498,7 +498,7 @@ impl CostModelStorageLayer for BackendManager {
     /// Each record in the `plan_cost` table can contain either the cost or the estimated statistic
     /// or both, but never neither.
     /// The name can be misleading, since it can also return the estimated statistic.
-    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<(Option<Cost>, Option<i32>)> {
+    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<(Option<Cost>, Option<f32>)> {
         let cost = PlanCost::find()
             .filter(plan_cost::Column::PhysicalExpressionId.eq(expr_id))
             .order_by_desc(plan_cost::Column::EpochId)
@@ -510,8 +510,8 @@ impl CostModelStorageLayer for BackendManager {
         }
 
         let real_cost = cost.as_ref().and_then(|c| c.cost.as_ref()).map(|c| Cost {
-            compute_cost: c.get("compute_cost").unwrap().as_i64().unwrap() as i32,
-            io_cost: c.get("io_cost").unwrap().as_i64().unwrap() as i32,
+            compute_cost: c.get("compute_cost").unwrap().as_f64().unwrap(),
+            io_cost: c.get("io_cost").unwrap().as_f64().unwrap(),
         });
 
         Ok((real_cost, cost.unwrap().estimated_statistic))
@@ -524,7 +524,7 @@ impl CostModelStorageLayer for BackendManager {
         &self,
         physical_expression_id: ExprId,
         cost: Option<Cost>,
-        estimated_statistic: Option<i32>,
+        estimated_statistic: Option<f32>,
         epoch_id: EpochId,
     ) -> StorageResult<()> {
         assert!(cost.is_some() || estimated_statistic.is_some());
@@ -801,10 +801,10 @@ mod tests {
             .store_cost(
                 expr_id,
                 Some(Cost {
-                    compute_cost: 42,
-                    io_cost: 42,
+                    compute_cost: 42.0,
+                    io_cost: 42.0,
                 }),
-                Some(42),
+                Some(42.0),
                 versioned_stat_res[0].epoch_id,
             )
             .await
@@ -1004,10 +1004,10 @@ mod tests {
             .unwrap();
         let physical_expression_id = 1;
         let cost = Cost {
-            compute_cost: 42,
-            io_cost: 42,
+            compute_cost: 42.0,
+            io_cost: 42.0,
         };
-        let mut estimated_statistic = 42;
+        let mut estimated_statistic = 42.0;
         backend_manager
             .store_cost(
                 physical_expression_id,
@@ -1028,12 +1028,9 @@ mod tests {
             costs[1].cost,
             Some(json!({"compute_cost": cost.compute_cost, "io_cost": cost.io_cost}))
         );
-        assert_eq!(
-            costs[1].estimated_statistic.unwrap() as i32,
-            estimated_statistic
-        );
+        assert_eq!(costs[1].estimated_statistic.unwrap(), estimated_statistic);
 
-        estimated_statistic = 50;
+        estimated_statistic = 50.0;
         backend_manager
             .store_cost(
                 physical_expression_id,
@@ -1055,7 +1052,7 @@ mod tests {
             Some(json!({"compute_cost": cost.compute_cost, "io_cost": cost.io_cost}))
         );
         assert_eq!(
-            costs[1].estimated_statistic.unwrap() as i32,
+            costs[1].estimated_statistic.unwrap(),
             estimated_statistic // The estimated_statistic should be update
         );
 
@@ -1074,8 +1071,8 @@ mod tests {
             .unwrap();
         let physical_expression_id = 1;
         let cost = Cost {
-            compute_cost: 42,
-            io_cost: 42,
+            compute_cost: 42.0,
+            io_cost: 42.0,
         };
         let _ = backend_manager
             .store_cost(physical_expression_id, Some(cost.clone()), None, epoch_id)
@@ -1114,7 +1111,7 @@ mod tests {
             .await
             .unwrap();
         let physical_expression_id = 1;
-        let estimated_statistic = 42;
+        let estimated_statistic = 42.0;
         let _ = backend_manager
             .store_cost(
                 physical_expression_id,
@@ -1131,10 +1128,7 @@ mod tests {
         assert_eq!(costs[1].epoch_id, epoch_id);
         assert_eq!(costs[1].physical_expression_id, physical_expression_id);
         assert_eq!(costs[1].cost, None);
-        assert_eq!(
-            costs[1].estimated_statistic.unwrap() as i32,
-            estimated_statistic
-        );
+        assert_eq!(costs[1].estimated_statistic.unwrap(), estimated_statistic);
         println!("{:?}", costs);
 
         // Retrieve physical_expression_id 1 and epoch_id 1
@@ -1147,11 +1141,11 @@ mod tests {
         assert_eq!(
             res.0.unwrap(),
             Cost {
-                compute_cost: 10,
-                io_cost: 10,
+                compute_cost: 10.0,
+                io_cost: 10.0,
             }
         );
-        assert_eq!(res.1.unwrap(), 10);
+        assert_eq!(res.1.unwrap(), 10.0);
 
         remove_db_file(DATABASE_FILE);
     }
diff --git a/optd-persistent/src/entities/plan_cost.rs b/optd-persistent/src/entities/plan_cost.rs
@@ -2,15 +2,16 @@
 
 use sea_orm::entity::prelude::*;
 
-#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
 #[sea_orm(table_name = "plan_cost")]
 pub struct Model {
     #[sea_orm(primary_key)]
     pub id: i32,
     pub physical_expression_id: i32,
     pub epoch_id: i32,
     pub cost: Option<Json>,
-    pub estimated_statistic: Option<i32>,
+    #[sea_orm(column_type = "Float", nullable)]
+    pub estimated_statistic: Option<f32>,
     pub is_valid: bool,
 }
 
diff --git a/optd-persistent/src/migrator/cost_model/m20241029_000001_plan_cost.rs b/optd-persistent/src/migrator/cost_model/m20241029_000001_plan_cost.rs
@@ -58,7 +58,7 @@ impl MigrationTrait for Migration {
                             .on_update(ForeignKeyAction::Cascade),
                     )
                     .col(json_null(PlanCost::Cost))
-                    .col(integer_null(PlanCost::EstimatedStatistic))
+                    .col(float_null(PlanCost::EstimatedStatistic))
                     .col(boolean(PlanCost::IsValid))
                     .to_owned(),
             )
diff --git a/schema/all_tables.dbml b/schema/all_tables.dbml

Original file line number	Diff line number	Diff line change
`@@ -58,7 +58,7 @@ impl MigrationTrait for Migration {`
`58`	`58`	`.on_update(ForeignKeyAction::Cascade),`
`59`	`59`	`)`
`60`	`60`	`.col(json_null(PlanCost::Cost))`
`61`		`- .col(integer_null(PlanCost::EstimatedStatistic))`
	`61`	`+ .col(float_null(PlanCost::EstimatedStatistic))`
`62`	`62`	`.col(boolean(PlanCost::IsValid))`
`63`	`63`	`.to_owned(),`
`64`	`64`	`)`