Separate memo and cost model storage_layer trait

optd-persistent/src/cost_model_orm_manager_impl.rs

@@ -0,0 +1,86 @@
+#![allow(dead_code, unused_imports, unused_variables)]
+
+use crate::memo_storage_layer::MemoStorageLayer;
+use crate::orm_manager::ORMManager;
+
+impl MemoStorageLayer for ORMManager {
+    async fn get_group_winner_from_group_id(
+        &self,
+        group_id: i32,
+    ) -> crate::StorageResult<Option<crate::entities::physical_expression::ActiveModel>> {
+        todo!()
+    }
+
+    async fn add_new_expr(
+        &mut self,
+        expr: crate::memo_storage_layer::Expression,
+    ) -> crate::StorageResult<(crate::GroupId, crate::ExprId)> {
+        todo!()
+    }
+
+    async fn add_expr_to_group(
+        &mut self,
+        expr: crate::memo_storage_layer::Expression,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_group_id(&self, expr_id: crate::ExprId) -> crate::StorageResult<crate::GroupId> {
+        todo!()
+    }
+
+    async fn get_expr_memoed(
+        &self,
+        expr_id: crate::ExprId,
+    ) -> crate::StorageResult<crate::memo_storage_layer::Expression> {
+        todo!()
+    }
+
+    async fn get_all_group_ids(&self) -> crate::StorageResult<Vec<crate::GroupId>> {
+        todo!()
+    }
+
+    async fn get_group(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<crate::entities::cascades_group::ActiveModel> {
+        todo!()
+    }
+
+    async fn update_group_winner(
+        &mut self,
+        group_id: crate::GroupId,
+        latest_winner: Option<crate::ExprId>,
+    ) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn get_all_exprs_in_group(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Vec<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_group_info(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<&Option<crate::ExprId>> {
+        todo!()
+    }
+
+    async fn get_predicate_binding(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::memo_storage_layer::Expression>> {
+        todo!()
+    }
+
+    async fn try_get_predicate_binding(
+        &self,
+        group_id: crate::GroupId,
+    ) -> crate::StorageResult<Option<crate::memo_storage_layer::Expression>> {
+        todo!()
+    }
+}

optd-persistent/src/cost_model_storage_layer.rs

@@ -0,0 +1,75 @@
+#![allow(dead_code, unused_imports)]
+
+use crate::entities::cascades_group;
+use crate::entities::event::Model as event_model;
+use crate::entities::logical_expression;
+use crate::entities::physical_expression;
+use crate::{EpochId, ExprId, StatId, StorageResult};
+use sea_orm::*;
+use sea_orm_migration::prelude::*;
+use serde_json::json;
+use std::sync::Arc;
+
+pub enum CatalogSource {
+    Iceberg(),
+}
+
+pub trait CostModelStorageLayer {
+    // TODO: Change EpochId to event::Model::epoch_id
+    async fn create_new_epoch(&mut self, source: String, data: String) -> StorageResult<EpochId>;
+
+    async fn update_stats_from_catalog(
+        &self,
+        c: CatalogSource,
+        epoch_id: EpochId,
+    ) -> StorageResult<()>;
+
+    // i32 in `stats:i32` is a placeholder for the stats type
+    async fn update_stats(&self, stats: i32, epoch_id: EpochId) -> StorageResult<()>;
+
+    async fn store_cost(&self, expr_id: ExprId, cost: i32, epoch_id: EpochId) -> StorageResult<()>;
+
+    async fn store_expr_stats_mappings(
+        &self,
+        expr_id: ExprId,
+        stat_ids: Vec<StatId>,
+    ) -> StorageResult<()>;
+
+    /// Get the statistics for a given table.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_table(
+        &self,
+        table_id: i32,
+        stat_type: i32,
+        epoch_id: Option<EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    /// Get the statistics for a given attribute.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_attr(
+        &self,
+        attr_id: i32,
+        stat_type: i32,
+        epoch_id: Option<EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    /// Get the joint statistics for a list of attributes.
+    ///
+    /// If `epoch_id` is None, it will return the latest statistics.
+    async fn get_stats_for_attrs(
+        &self,
+        attr_ids: Vec<i32>,
+        stat_type: i32,
+        epoch_id: Option<EpochId>,
+    ) -> StorageResult<Option<f32>>;
+
+    async fn get_cost_analysis(
+        &self,
+        expr_id: ExprId,
+        epoch_id: EpochId,
+    ) -> StorageResult<Option<i32>>;
+
+    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<Option<i32>>;
+}

optd-persistent/src/lib.rs

@@ -1,12 +1,23 @@
 use sea_orm::*;
 use sea_orm_migration::prelude::*;
 
+mod cost_model_orm_manager_impl;
+mod cost_model_storage_layer;
 mod entities;
+mod memo_orm_manager_impl;
+mod memo_storage_layer;
 mod migrator;
 mod orm_manager;
-mod storage_layer;
+
 use migrator::Migrator;
 
+pub type GroupId = i32;
+pub type ExprId = i32;
+pub type EpochId = i32;
+pub type StatId = i32;
+
+pub type StorageResult<T> = Result<T, DbErr>;
+
 pub const DATABASE_URL: &str = "sqlite:./sqlite.db?mode=rwc";
 pub const DATABASE_FILE: &str = "./sqlite.db";
 pub const TEST_DATABASE_URL: &str = "sqlite:./test.db?mode=rwc";

optd-persistent/src/memo_orm_manager_impl.rs

@@ -0,0 +1,83 @@
+#![allow(dead_code, unused_imports, unused_variables)]
+
+use crate::cost_model_storage_layer::CostModelStorageLayer;
+use crate::memo_storage_layer::MemoStorageLayer;
+use crate::orm_manager::ORMManager;
+
+impl CostModelStorageLayer for ORMManager {
+    async fn create_new_epoch(
+        &mut self,
+        source: String,
+        data: String,
+    ) -> crate::StorageResult<crate::EpochId> {
+        todo!()
+    }
+
+    async fn update_stats_from_catalog(
+        &self,
+        c: crate::cost_model_storage_layer::CatalogSource,
+        epoch_id: crate::EpochId,
+    ) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn update_stats(&self, stats: i32, epoch_id: crate::EpochId) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn store_cost(
+        &self,
+        expr_id: crate::ExprId,
+        cost: i32,
+        epoch_id: crate::EpochId,
+    ) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn store_expr_stats_mappings(
+        &self,
+        expr_id: crate::ExprId,
+        stat_ids: Vec<crate::StatId>,
+    ) -> crate::StorageResult<()> {
+        todo!()
+    }
+
+    async fn get_stats_for_table(
+        &self,
+        table_id: i32,
+        stat_type: i32,
+        epoch_id: Option<crate::EpochId>,
+    ) -> crate::StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_stats_for_attr(
+        &self,
+        attr_id: i32,
+        stat_type: i32,
+        epoch_id: Option<crate::EpochId>,
+    ) -> crate::StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_stats_for_attrs(
+        &self,
+        attr_ids: Vec<i32>,
+        stat_type: i32,
+        epoch_id: Option<crate::EpochId>,
+    ) -> crate::StorageResult<Option<f32>> {
+        todo!()
+    }
+
+    async fn get_cost_analysis(
+        &self,
+        expr_id: crate::ExprId,
+        epoch_id: crate::EpochId,
+    ) -> crate::StorageResult<Option<i32>> {
+        todo!()
+    }
+
+    async fn get_cost(&self, expr_id: crate::ExprId) -> crate::StorageResult<Option<i32>> {
+        todo!()
+    }
+}

optd-persistent/src/storage_layer.rs → optd-persistent/src/memo_storage_layer.rs

@@ -4,21 +4,12 @@ use crate::entities::cascades_group;
 use crate::entities::event::Model as event_model;
 use crate::entities::logical_expression;
 use crate::entities::physical_expression;
+use crate::{ExprId, GroupId, StorageResult};
 use sea_orm::*;
 use sea_orm_migration::prelude::*;
 use serde_json::json;
 use std::sync::Arc;
 
-pub type GroupId = i32;
-pub type ExprId = i32;
-pub type EpochId = i32;
-
-pub type StorageResult<T> = Result<T, DbErr>;
-
-pub enum CatalogSource {
-    Iceberg(),
-}
-
 pub enum Expression {
     LogicalExpression(logical_expression::Model),
     PhysicalExpression(physical_expression::Model),
@@ -37,59 +28,7 @@ pub enum Expression {
 // The optd WinnerInfo struct makes everything too coupled.
 pub struct WinnerInfo {}
 
-pub trait StorageLayer {
-    // TODO: Change EpochId to event::Model::epoch_id
-    async fn create_new_epoch(&mut self, source: String, data: String) -> StorageResult<EpochId>;
-
-    async fn update_stats_from_catalog(
-        &self,
-        c: CatalogSource,
-        epoch_id: EpochId,
-    ) -> StorageResult<()>;
-
-    // i32 in `stats:i32` is a placeholder for the stats type
-    async fn update_stats(&self, stats: i32, epoch_id: EpochId) -> StorageResult<()>;
-
-    async fn store_cost(&self, expr_id: ExprId, cost: i32, epoch_id: EpochId) -> StorageResult<()>;
-
-    /// Get the statistics for a given table.
-    ///
-    /// If `epoch_id` is None, it will return the latest statistics.
-    async fn get_stats_for_table(
-        &self,
-        table_id: i32,
-        stat_type: i32,
-        epoch_id: Option<EpochId>,
-    ) -> StorageResult<Option<f32>>;
-
-    /// Get the statistics for a given attribute.
-    ///
-    /// If `epoch_id` is None, it will return the latest statistics.
-    async fn get_stats_for_attr(
-        &self,
-        attr_id: i32,
-        stat_type: i32,
-        epoch_id: Option<EpochId>,
-    ) -> StorageResult<Option<f32>>;
-
-    /// Get the joint statistics for a list of attributes.
-    ///
-    /// If `epoch_id` is None, it will return the latest statistics.
-    async fn get_stats_for_attrs(
-        &self,
-        attr_ids: Vec<i32>,
-        stat_type: i32,
-        epoch_id: Option<EpochId>,
-    ) -> StorageResult<Option<f32>>;
-
-    async fn get_cost_analysis(
-        &self,
-        expr_id: ExprId,
-        epoch_id: EpochId,
-    ) -> StorageResult<Option<i32>>;
-
-    async fn get_cost(&self, expr_id: ExprId) -> StorageResult<Option<i32>>;
-
+pub trait MemoStorageLayer {
     async fn get_group_winner_from_group_id(
         &self,
         group_id: i32,