implement agg cost computation

Author: xx01cyx

optd-cost-model/src/common/predicates/attr_ref_pred.rs

@@ -11,6 +11,8 @@ use super::id_pred::IdPred;
 /// Currently, [`AttributeRefPred`] only holds base table attributes, i.e. attributes
 /// that already exist in the table. More complex structures may be introduced in the
 /// future to represent derived attributes (e.g. t.v1 + t.v2).
+///
+/// TODO: Support derived column in `AttributeRefPred`.
 #[derive(Clone, Debug)]
 pub struct AttributeRefPred(pub ArcPredicateNode);
 

optd-cost-model/src/cost/agg.rs

@@ -0,0 +1,61 @@
+use optd_persistent::CostModelStorageLayer;
+
+use crate::{
+    common::{
+        nodes::{ArcPredicateNode, PredicateType, ReprPredicateNode},
+        predicates::{attr_ref_pred::AttributeRefPred, list_pred::ListPred},
+        types::TableId,
+    },
+    cost_model::CostModelImpl,
+    stats::DEFAULT_NUM_DISTINCT,
+    CostModelError, CostModelResult, EstimatedStatistic,
+};
+
+impl<S: CostModelStorageLayer> CostModelImpl<S> {
+    pub fn get_agg_row_cnt(
+        &self,
+        group_by: ArcPredicateNode,
+    ) -> CostModelResult<EstimatedStatistic> {
+        let group_by = ListPred::from_pred_node(group_by).unwrap();
+        if group_by.is_empty() {
+            Ok(EstimatedStatistic(1))
+        } else {
+            // Multiply the n-distinct of all the group by columns.
+            // TODO: improve with multi-dimensional n-distinct
+            let row_cnt = group_by.0.children.iter().try_fold(1, |acc, node| {
+                match node.typ {
+                    PredicateType::AttributeRef => {
+                        let attr_ref =
+                            AttributeRefPred::from_pred_node(node.clone()).ok_or_else(|| {
+                                CostModelError::InvalidPredicate(
+                                    "Expected AttributeRef predicate".to_string(),
+                                )
+                            })?;
+                        if attr_ref.is_derived() {
+                            Ok(acc * DEFAULT_NUM_DISTINCT)
+                        } else {
+                            let table_id = attr_ref.table_id();
+                            let attr_idx = attr_ref.attr_index();
+                            let stats_option =
+                                self.get_attribute_comb_stats(TableId(table_id), &vec![attr_idx])?;
+
+                            let ndistinct = match stats_option {
+                                Some(stats) => stats.ndistinct,
+                                None => {
+                                    // The column type is not supported or stats are missing.
+                                    DEFAULT_NUM_DISTINCT
+                                }
+                            };
+                            Ok(acc * ndistinct)
+                        }
+                    }
+                    _ => {
+                        // TODO: Consider the case where `GROUP BY 1`.
+                        panic!("GROUP BY must have attribute ref predicate")
+                    }
+                }
+            })?;
+            Ok(EstimatedStatistic(row_cnt))
+        }
+    }
+}

optd-cost-model/src/lib.rs

@@ -38,8 +38,8 @@ pub type CostModelResult<T> = Result<T, CostModelError>;
 
 #[derive(Debug)]
 pub enum CostModelError {
-    // TODO: Add more error types
     ORMError(BackendError),
+    InvalidPredicate(String),
 }
 
 pub trait CostModel: 'static + Send + Sync {

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

@@ -8,6 +8,9 @@ use crate::common::values::Value;
 use counter::Counter;
 use serde::{Deserialize, Serialize};
 
+// Default n-distinct estimate for derived columns or columns lacking statistics
+pub const DEFAULT_NUM_DISTINCT: u64 = 200;
+
 pub type AttributeCombValue = Vec<Option<Value>>;
 
 #[derive(Serialize, Deserialize, Debug)]