Finish most tests for filter

Author: lanlou1554

optd-cost-model/src/cost/filter/controller.rs

@@ -136,10 +136,9 @@ pub mod tests {
             values::Value,
         },
         stats::{
-            counter::Counter, tdigest::TDigest, AttributeCombValueStats, Distribution,
-            MostCommonValues,
+            utilities::counter::Counter, AttributeCombValueStats, Distribution, MostCommonValues,
         },
-        storage::mock::{CostModelStorageMockManagerImpl, TableStats},
+        storage::mock::{BaseTableAttrInfo, CostModelStorageMockManagerImpl, TableStats},
     };
 
     use super::*;
@@ -152,6 +151,7 @@ pub mod tests {
         table_id: Vec<TableId>,
         per_attribute_stats: Vec<TestPerAttributeStats>,
         row_counts: Vec<Option<usize>>,
+        per_table_attr_infos: BaseTableAttrInfo,
     ) -> TestOptCostModelMock {
         let storage_manager = CostModelStorageMockManagerImpl::new(
             table_id
@@ -168,6 +168,7 @@ pub mod tests {
                     )
                 })
                 .collect(),
+            per_table_attr_infos,
         );
         CostModelImpl::new(storage_manager, CatalogSource::Mock)
     }
@@ -222,9 +223,6 @@ pub mod tests {
         )
     }
 
-    /// The reason this isn't an associated function of PerAttributeStats is because that would require
-    ///   adding an empty type to the enum definitions of MostCommonValues and Distribution,
-    ///   which I wanted to avoid
     pub(crate) fn get_empty_per_attr_stats() -> TestPerAttributeStats {
         TestPerAttributeStats::new(MostCommonValues::Counter(Counter::default()), 0, 0.0, None)
     }

optd-cost-model/src/cost_model.rs

@@ -1,12 +1,15 @@
 #![allow(unused)]
 
 mod arith_encoder;
-pub mod counter;
-pub mod tdigest;
+pub mod utilities;
 
 use crate::common::values::Value;
-use counter::Counter;
 use serde::{Deserialize, Serialize};
+use utilities::counter::Counter;
+use utilities::{
+    simple_map::{self, SimpleMap},
+    tdigest::TDigest,
+};
 
 // Default n-distinct estimate for derived columns or columns lacking statistics
 pub const DEFAULT_NUM_DISTINCT: u64 = 200;
@@ -27,7 +30,8 @@ pub const FIXED_CHAR_SEL_FACTOR: f64 = 0.2;
 
 pub type AttributeCombValue = Vec<Option<Value>>;
 
-#[derive(Serialize, Deserialize, Debug)]
+// TODO: remove the clone, see the comment in the [`AttributeCombValueStats`]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(tag = "type")]
 pub enum MostCommonValues {
     Counter(Counter<AttributeCombValue>),
@@ -71,10 +75,12 @@ impl MostCommonValues {
     }
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+// TODO: remove the clone, see the comment in the [`AttributeCombValueStats`]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(tag = "type")]
 pub enum Distribution {
-    TDigest(tdigest::TDigest<Value>),
+    TDigest(TDigest<Value>),
+    SimpleDistribution(SimpleMap),
     // Add more types here...
 }
 
@@ -89,11 +95,21 @@ impl Distribution {
                     tdigest.centroids.len() as f64 * tdigest.cdf(value) / nb_rows as f64
                 }
             }
+            Distribution::SimpleDistribution(simple_distribution) => {
+                *simple_distribution.m.get(value).unwrap_or(&0.0)
+            }
         }
     }
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+// TODO: Remove the clone. Now I have to add this because
+// persistent.rs doesn't have a memory cache, so we have to
+// return AttributeCombValueStats rather than &AttributeCombValueStats.
+// But this poses a problem for mock.rs when testing, since mock storage
+// only has memory hash map, so we need to return a clone of AttributeCombValueStats.
+// Later, if memory cache is added, we should change this to return a reference.
+// **and** remove the clone.
+#[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct AttributeCombValueStats {
     pub mcvs: MostCommonValues,      // Does NOT contain full nulls.
     pub distr: Option<Distribution>, // Does NOT contain mcvs; optional.

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

@@ -5,8 +5,9 @@ use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 
 /// The Counter structure to track exact frequencies of fixed elements.
+/// TODO: remove the clone, see the comment in the [`AttributeCombValueStats`]
 #[serde_with::serde_as]
-#[derive(Default, Serialize, Deserialize, Debug)]
+#[derive(Default, Serialize, Deserialize, Debug, Clone)]
 pub struct Counter<T: PartialEq + Eq + Hash + Clone + Serialize + DeserializeOwned> {
     #[serde_as(as = "HashMap<serde_with::json::JsonString, _>")]
     counts: HashMap<T, i32>, // The exact counts of an element T.

optd-cost-model/src/stats/counter.rs renamed to optd-cost-model/src/stats/utilities/counter.rs

@@ -0,0 +1,3 @@
+pub mod counter;
+pub mod simple_map;
+pub mod tdigest;

optd-cost-model/src/stats/utilities/mod.rs

@@ -0,0 +1,20 @@
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+
+use crate::common::values::Value;
+
+/// TODO: documentation
+/// Now it is mainly for testing purposes.
+#[derive(Clone, Serialize, Deserialize, Debug)]
+pub struct SimpleMap {
+    pub(crate) m: HashMap<Value, f64>,
+}
+
+impl SimpleMap {
+    pub fn new(v: Vec<(Value, f64)>) -> Self {
+        Self {
+            m: v.into_iter().collect(),
+        }
+    }
+}

optd-cost-model/src/stats/utilities/simple_map.rs

@@ -15,9 +15,7 @@ use std::marker::PhantomData;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
-use crate::common::values::Value;
-
-use super::arith_encoder;
+use crate::{common::values::Value, stats::arith_encoder};
 
 pub const DEFAULT_COMPRESSION: f64 = 200.0;
 

optd-cost-model/src/stats/tdigest.rs renamed to optd-cost-model/src/stats/utilities/tdigest.rs

@@ -30,15 +30,21 @@ impl TableStats {
 }
 
 pub type BaseTableStats = HashMap<TableId, TableStats>;
+pub type BaseTableAttrInfo = HashMap<TableId, HashMap<i32, Attribute>>;
 
 pub struct CostModelStorageMockManagerImpl {
     pub(crate) per_table_stats_map: BaseTableStats,
+    pub(crate) per_table_attr_infos_map: BaseTableAttrInfo,
 }
 
 impl CostModelStorageMockManagerImpl {
-    pub fn new(per_table_stats_map: BaseTableStats) -> Self {
+    pub fn new(
+        per_table_stats_map: BaseTableStats,
+        per_table_attr_infos_map: BaseTableAttrInfo,
+    ) -> Self {
         Self {
             per_table_stats_map,
+            per_table_attr_infos_map,
         }
     }
 }
@@ -49,14 +55,28 @@ impl CostModelStorageManager for CostModelStorageMockManagerImpl {
         table_id: TableId,
         attr_base_index: i32,
     ) -> CostModelResult<Option<Attribute>> {
-        todo!()
+        let table_attr_infos = self.per_table_attr_infos_map.get(&table_id);
+        match table_attr_infos {
+            None => Ok(None),
+            Some(table_attr_infos) => match table_attr_infos.get(&attr_base_index) {
+                None => Ok(None),
+                Some(attr) => Ok(Some(attr.clone())),
+            },
+        }
     }
 
     async fn get_attributes_comb_statistics(
         &self,
         table_id: TableId,
         attr_base_indices: &[usize],
     ) -> CostModelResult<Option<AttributeCombValueStats>> {
-        todo!()
+        let table_stats = self.per_table_stats_map.get(&table_id);
+        match table_stats {
+            None => Ok(None),
+            Some(table_stats) => match table_stats.column_comb_stats.get(attr_base_indices) {
+                None => Ok(None),
+                Some(stats) => Ok(Some(stats.clone())),
+            },
+        }
     }
 }

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

@@ -5,7 +5,7 @@ use optd_persistent::{cost_model::interface::StatType, CostModelStorageLayer};
 
 use crate::{
     common::{predicates::constant_pred::ConstantType, types::TableId},
-    stats::{counter::Counter, AttributeCombValueStats, Distribution, MostCommonValues},
+    stats::{utilities::counter::Counter, AttributeCombValueStats, Distribution, MostCommonValues},
     CostModelResult,
 };