refactor(adv-cost): compiles and not rely on opt ctx (#223)

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

Cargo.lock

@@ -58,7 +58,7 @@ tokio = { version = "1.24", features = [
 ] }
 url = "2.2"
 optd-datafusion-bridge = { path = "../optd-datafusion-bridge" }
-# optd-datafusion-repr-adv-cost = { path = "../optd-datafusion-repr-adv-cost" }
+optd-datafusion-repr-adv-cost = { path = "../optd-datafusion-repr-adv-cost" }
 optd-datafusion-repr = { path = "../optd-datafusion-repr" }
 tracing-subscriber = "0.3"
 tracing = "0.1"

datafusion-optd-cli/Cargo.toml

@@ -188,6 +188,17 @@ impl Value {
                 Value::Int64(i64) => (*i64).try_into().unwrap(),
                 _ => panic!("{self} could not be converted into an Int32"),
             }),
+            DataType::Int64 => Value::Int64(match self {
+                Value::Int64(i64) => *i64,
+                Value::Int32(i32) => (*i32).try_into().unwrap(),
+                _ => panic!("{self} could not be converted into an Int64"),
+            }),
+            DataType::UInt64 => Value::UInt64(match self {
+                Value::Int64(i64) => (*i64).try_into().unwrap(),
+                Value::UInt64(i64) => *i64,
+                Value::UInt32(i32) => (*i32).try_into().unwrap(),
+                _ => panic!("{self} could not be converted into an UInt64"),
+            }),
             DataType::Date32 => Value::Date32(match self {
                 Value::Date32(date32) => *date32,
                 Value::String(str) => {
@@ -307,6 +318,12 @@ impl<T: NodeType> From<ArcPlanNode<T>> for PlanNodeOrGroup<T> {
     }
 }
 
+impl<T: NodeType> From<GroupId> for PlanNodeOrGroup<T> {
+    fn from(value: GroupId) -> Self {
+        Self::Group(value)
+    }
+}
+
 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
 pub struct PredNode<T: NodeType> {
     /// A generic predicate node type

optd-core/src/nodes.rs

@@ -71,10 +71,13 @@ impl<
 mod tests {
     use arrow_schema::DataType;
     use itertools::Itertools;
-    use optd_core::nodes::Value;
-    use optd_datafusion_repr::plan_nodes::{
-        BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, Expr, ExprList, InListExpr,
-        LikeExpr, LogOpExpr, LogOpType, OptRelNode, OptRelNodeRef, UnOpExpr, UnOpType,
+    use optd_datafusion_repr::{
+        plan_nodes::{
+            ArcDfPredNode, BinOpPred, BinOpType, CastPred, ColumnRefPred, ConstantPred,
+            DfReprPredNode, InListPred, LikePred, ListPred, LogOpPred, LogOpType, UnOpPred,
+            UnOpType,
+        },
+        Value,
     };
     use serde::{Deserialize, Serialize};
     use std::collections::HashMap;
@@ -285,74 +288,46 @@ mod tests {
         )
     }
 
-    pub fn col_ref(idx: u64) -> OptRelNodeRef {
+    pub fn col_ref(idx: u64) -> ArcDfPredNode {
         // this conversion is always safe because idx was originally a usize
         let idx_as_usize = idx as usize;
-        ColumnRefExpr::new(idx_as_usize).into_rel_node()
+        ColumnRefPred::new(idx_as_usize).into_pred_node()
     }
 
-    pub fn cnst(value: Value) -> OptRelNodeRef {
-        ConstantExpr::new(value).into_rel_node()
+    pub fn cnst(value: Value) -> ArcDfPredNode {
+        ConstantPred::new(value).into_pred_node()
     }
 
-    pub fn cast(child: OptRelNodeRef, cast_type: DataType) -> OptRelNodeRef {
-        CastExpr::new(
-            Expr::from_rel_node(child).expect("child should be an Expr"),
-            cast_type,
-        )
-        .into_rel_node()
+    pub fn cast(child: ArcDfPredNode, cast_type: DataType) -> ArcDfPredNode {
+        CastPred::new(child, cast_type).into_pred_node()
     }
 
-    pub fn bin_op(op_type: BinOpType, left: OptRelNodeRef, right: OptRelNodeRef) -> OptRelNodeRef {
-        BinOpExpr::new(
-            Expr::from_rel_node(left).expect("left should be an Expr"),
-            Expr::from_rel_node(right).expect("right should be an Expr"),
-            op_type,
-        )
-        .into_rel_node()
+    pub fn bin_op(op_type: BinOpType, left: ArcDfPredNode, right: ArcDfPredNode) -> ArcDfPredNode {
+        BinOpPred::new(left, right, op_type).into_pred_node()
     }
 
-    pub fn log_op(op_type: LogOpType, children: Vec<OptRelNodeRef>) -> OptRelNodeRef {
-        LogOpExpr::new(
-            op_type,
-            ExprList::new(
-                children
-                    .into_iter()
-                    .map(|opt_rel_node_ref| {
-                        Expr::from_rel_node(opt_rel_node_ref).expect("all children should be Expr")
-                    })
-                    .collect(),
-            ),
-        )
-        .into_rel_node()
+    pub fn log_op(op_type: LogOpType, children: Vec<ArcDfPredNode>) -> ArcDfPredNode {
+        LogOpPred::new(op_type, children).into_pred_node()
     }
 
-    pub fn un_op(op_type: UnOpType, child: OptRelNodeRef) -> OptRelNodeRef {
-        UnOpExpr::new(
-            Expr::from_rel_node(child).expect("child should be an Expr"),
-            op_type,
-        )
-        .into_rel_node()
+    pub fn un_op(op_type: UnOpType, child: ArcDfPredNode) -> ArcDfPredNode {
+        UnOpPred::new(child, op_type).into_pred_node()
     }
 
-    pub fn in_list(col_ref_idx: u64, list: Vec<Value>, negated: bool) -> InListExpr {
-        InListExpr::new(
-            Expr::from_rel_node(col_ref(col_ref_idx)).unwrap(),
-            ExprList::new(
-                list.into_iter()
-                    .map(|v| Expr::from_rel_node(cnst(v)).unwrap())
-                    .collect_vec(),
-            ),
+    pub fn in_list(col_ref_idx: u64, list: Vec<Value>, negated: bool) -> InListPred {
+        InListPred::new(
+            col_ref(col_ref_idx),
+            ListPred::new(list.into_iter().map(|v| cnst(v)).collect_vec()),
             negated,
         )
     }
 
-    pub fn like(col_ref_idx: u64, pattern: &str, negated: bool) -> LikeExpr {
-        LikeExpr::new(
+    pub fn like(col_ref_idx: u64, pattern: &str, negated: bool) -> LikePred {
+        LikePred::new(
             negated,
             false,
-            Expr::from_rel_node(col_ref(col_ref_idx)).unwrap(),
-            Expr::from_rel_node(cnst(Value::String(pattern.into()))).unwrap(),
+            col_ref(col_ref_idx),
+            cnst(Value::String(pattern.into())),
         )
     }
 

optd-datafusion-repr-adv-cost/src/adv_stats.rs

@@ -1,88 +1,59 @@
-use optd_core::{
-    cascades::{BindingType, CascadesOptimizer, RelNodeContext},
-    cost::Cost,
-};
+use optd_core::cascades::{CascadesOptimizer, RelNodeContext};
 use serde::{de::DeserializeOwned, Serialize};
 
-use crate::adv_cost::{
+use crate::adv_stats::{
     stats::{Distribution, MostCommonValues},
     DEFAULT_NUM_DISTINCT,
 };
 use optd_datafusion_repr::{
-    plan_nodes::{ExprList, OptRelNode, OptRelNodeTyp},
-    properties::column_ref::{BaseTableColumnRef, ColumnRef, ColumnRefPropertyBuilder},
+    plan_nodes::{ArcDfPredNode, DfReprPredNode, ListPred},
+    properties::column_ref::{
+        BaseTableColumnRef, ColumnRef, ColumnRefPropertyBuilder, GroupColumnRefs,
+    },
 };
 
-use super::{OptCostModel, DEFAULT_UNK_SEL};
+use super::{AdvStats, DEFAULT_UNK_SEL};
 
 impl<
         M: MostCommonValues + Serialize + DeserializeOwned,
         D: Distribution + Serialize + DeserializeOwned,
-    > OptCostModel<M, D>
+    > AdvStats<M, D>
 {
-    pub(super) fn get_agg_cost(
-        &self,
-        children: &[Cost],
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
-    ) -> Cost {
-        let child_row_cnt = Self::row_cnt(&children[0]);
-        let row_cnt = self.get_agg_row_cnt(context, optimizer, child_row_cnt);
-        let (_, compute_cost_1, _) = Self::cost_tuple(&children[1]);
-        let (_, compute_cost_2, _) = Self::cost_tuple(&children[2]);
-        Self::cost(
-            row_cnt,
-            child_row_cnt * (compute_cost_1 + compute_cost_2),
-            0.0,
-        )
-    }
-
-    fn get_agg_row_cnt(
+    pub(crate) fn get_agg_row_cnt(
         &self,
-        context: Option<RelNodeContext>,
-        optimizer: Option<&CascadesOptimizer<DfNodeType>>,
-        child_row_cnt: f64,
+        group_by: ArcDfPredNode,
+        output_col_refs: GroupColumnRefs,
     ) -> f64 {
-        if let (Some(context), Some(optimizer)) = (context, optimizer) {
-            let group_by_id = context.children_group_ids[2];
-            let group_by = optimizer
-                .get_predicate_binding(group_by_id)
-                .expect("no expression found?");
-            let group_by = ExprList::from_rel_node(group_by).unwrap();
-            if group_by.is_empty() {
-                1.0
-            } else {
-                // Multiply the n-distinct of all the group by columns.
-                // TODO: improve with multi-dimensional n-distinct
-                let group_col_refs = optimizer
-                    .get_property_by_group::<ColumnRefPropertyBuilder>(context.group_id, 1);
-                group_col_refs
-                    .base_table_column_refs()
-                    .iter()
-                    .take(group_by.len())
-                    .map(|col_ref| match col_ref {
-                        ColumnRef::BaseTableColumnRef(BaseTableColumnRef { table, col_idx }) => {
-                            let table_stats = self.per_table_stats_map.get(table);
-                            let column_stats = table_stats.and_then(|table_stats| {
-                                table_stats.column_comb_stats.get(&vec![*col_idx])
-                            });
+        let group_by = ListPred::from_pred_node(group_by).unwrap();
+        if group_by.is_empty() {
+            1.0
+        } else {
+            // Multiply the n-distinct of all the group by columns.
+            // TODO: improve with multi-dimensional n-distinct
+            output_col_refs
+                .base_table_column_refs()
+                .iter()
+                .take(group_by.len())
+                .map(|col_ref| match col_ref {
+                    ColumnRef::BaseTableColumnRef(BaseTableColumnRef { table, col_idx }) => {
+                        let table_stats = self.per_table_stats_map.get(table);
+                        let column_stats = table_stats.and_then(|table_stats| {
+                            table_stats.column_comb_stats.get(&vec![*col_idx])
+                        });
 
-                            if let Some(column_stats) = column_stats {
-                                column_stats.ndistinct as f64
-                            } else {
-                                // The column type is not supported or stats are missing.
-                                DEFAULT_NUM_DISTINCT as f64
-                            }
+                        if let Some(column_stats) = column_stats {
+                            column_stats.ndistinct as f64
+                        } else {
+                            // The column type is not supported or stats are missing.
+                            DEFAULT_NUM_DISTINCT as f64
                         }
-                        ColumnRef::Derived => DEFAULT_NUM_DISTINCT as f64,
-                        _ => panic!(
-                            "GROUP BY base table column ref must either be derived or base table"
-                        ),
-                    })
-                    .product()
-            }
-        } else {
-            (child_row_cnt * DEFAULT_UNK_SEL).max(1.0)
+                    }
+                    ColumnRef::Derived => DEFAULT_NUM_DISTINCT as f64,
+                    _ => panic!(
+                        "GROUP BY base table column ref must either be derived or base table"
+                    ),
+                })
+                .product()
         }
     }
 }