modify MemoExt interface

Author: xx01cyx

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

@@ -28,9 +28,9 @@ impl<S: CostModelStorageManager> CostModelImpl<S> {
         right_keys: ListPred,
     ) -> CostModelResult<EstimatedStatistic> {
         let selectivity = {
-            let output_attr_refs = self.memo.get_attribute_ref(group_id);
-            let left_attr_refs = self.memo.get_attribute_ref(left_group_id);
-            let right_attr_refs = self.memo.get_attribute_ref(right_group_id);
+            let output_attr_refs = self.memo.get_attribute_refs(group_id);
+            let left_attr_refs = self.memo.get_attribute_refs(left_group_id);
+            let right_attr_refs = self.memo.get_attribute_refs(right_group_id);
             let left_attr_cnt = left_attr_refs.attr_refs().len();
             // there may be more than one expression tree in a group.
             // see comment in PredicateType::PhysicalFilter(_) for more information

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

@@ -25,9 +25,9 @@ impl<S: CostModelStorageManager> CostModelImpl<S> {
         join_cond: ArcPredicateNode,
     ) -> CostModelResult<EstimatedStatistic> {
         let selectivity = {
-            let output_attr_refs = self.memo.get_attribute_ref(group_id);
-            let left_attr_refs = self.memo.get_attribute_ref(left_group_id);
-            let right_attr_refs = self.memo.get_attribute_ref(right_group_id);
+            let output_attr_refs = self.memo.get_attribute_refs(group_id);
+            let left_attr_refs = self.memo.get_attribute_refs(left_group_id);
+            let right_attr_refs = self.memo.get_attribute_refs(right_group_id);
             let input_correlation = get_input_correlation(left_attr_refs, right_attr_refs);
 
             self.get_join_selectivity_from_expr_tree(

optd-cost-model/src/cost_model.rs

@@ -155,15 +155,15 @@ pub mod tests {
 
     use super::*;
 
-    const TEST_TABLE1_ID: TableId = TableId(0);
-    const TEST_TABLE2_ID: TableId = TableId(1);
-    const TEST_TABLE3_ID: TableId = TableId(2);
-    const TEST_TABLE4_ID: TableId = TableId(3);
+    pub const TEST_TABLE1_ID: TableId = TableId(0);
+    pub const TEST_TABLE2_ID: TableId = TableId(1);
+    pub const TEST_TABLE3_ID: TableId = TableId(2);
+    pub const TEST_TABLE4_ID: TableId = TableId(3);
 
-    const TEST_GROUP1_ID: GroupId = GroupId(0);
-    const TEST_GROUP2_ID: GroupId = GroupId(1);
-    const TEST_GROUP3_ID: GroupId = GroupId(2);
-    const TEST_GROUP4_ID: GroupId = GroupId(3);
+    pub const TEST_GROUP1_ID: GroupId = GroupId(0);
+    pub const TEST_GROUP2_ID: GroupId = GroupId(1);
+    pub const TEST_GROUP3_ID: GroupId = GroupId(2);
+    pub const TEST_GROUP4_ID: GroupId = GroupId(3);
 
     pub type TestPerAttributeStats = AttributeCombValueStats;
     // TODO: add tests for non-mock storage manager
@@ -495,7 +495,7 @@ pub mod tests {
         }
 
         pub fn get_attr_refs(&self, group_id: GroupId) -> GroupAttrRefs {
-            self.memo.get_attribute_ref(group_id)
+            self.memo.get_attribute_refs(group_id)
         }
     }
 

optd-cost-model/src/memo_ext.rs

@@ -1,5 +1,9 @@
 use crate::common::{
-    properties::{attr_ref::GroupAttrRefs, schema::Schema, Attribute},
+    properties::{
+        attr_ref::{AttrRef, GroupAttrRefs},
+        schema::Schema,
+        Attribute,
+    },
     types::GroupId,
 };
 
@@ -13,10 +17,12 @@ use crate::common::{
 pub trait MemoExt: Send + Sync + 'static {
     /// Get the schema of a group in the memo.
     fn get_schema(&self, group_id: GroupId) -> Schema;
-    /// Get the attribute reference of a group in the memo.
-    fn get_attribute_ref(&self, group_id: GroupId) -> GroupAttrRefs;
-    /// Get the attribute information of a given attribute in a group in the memo.
+    /// Get the attribute info of a given attribute in a group in the memo.
     fn get_attribute_info(&self, group_id: GroupId, attr_ref_idx: u64) -> Attribute;
+    /// Get the attribute reference of a group in the memo.
+    fn get_attribute_refs(&self, group_id: GroupId) -> GroupAttrRefs;
+    /// Get the attribute reference of a given attribute in a group in the memo.
+    fn get_attribute_ref(&self, group_id: GroupId, attr_ref_idx: u64) -> AttrRef;
 
     // TODO: Figure out what other information is needed to compute the cost...
 }
@@ -26,18 +32,22 @@ pub mod tests {
     use std::collections::HashMap;
 
     use crate::common::{
-        properties::{attr_ref::GroupAttrRefs, schema::Schema, Attribute},
+        properties::{
+            attr_ref::{AttrRef, GroupAttrRefs},
+            schema::Schema,
+            Attribute,
+        },
         types::GroupId,
     };
 
     pub struct MemoGroupInfo {
         pub schema: Schema,
-        pub attr_ref: GroupAttrRefs,
+        pub attr_refs: GroupAttrRefs,
     }
 
     impl MemoGroupInfo {
-        pub fn new(schema: Schema, attr_ref: GroupAttrRefs) -> Self {
-            Self { schema, attr_ref }
+        pub fn new(schema: Schema, attr_refs: GroupAttrRefs) -> Self {
+            Self { schema, attr_refs }
         }
     }
 
@@ -63,13 +73,17 @@ pub mod tests {
             self.memo.get(&group_id).unwrap().schema.clone()
         }
 
-        fn get_attribute_ref(&self, group_id: GroupId) -> GroupAttrRefs {
-            self.memo.get(&group_id).unwrap().attr_ref.clone()
-        }
-
         fn get_attribute_info(&self, group_id: GroupId, attr_ref_idx: u64) -> Attribute {
             self.memo.get(&group_id).unwrap().schema.attributes[attr_ref_idx as usize].clone()
         }
+
+        fn get_attribute_refs(&self, group_id: GroupId) -> GroupAttrRefs {
+            self.memo.get(&group_id).unwrap().attr_refs.clone()
+        }
+
+        fn get_attribute_ref(&self, group_id: GroupId, attr_ref_idx: u64) -> AttrRef {
+            self.memo.get(&group_id).unwrap().attr_refs.attr_refs()[attr_ref_idx as usize].clone()
+        }
     }
 
     impl From<HashMap<GroupId, MemoGroupInfo>> for MockMemoExtImpl {