feat(cost-model): migration of compute cost parts

optd-cost-model/Cargo.toml

@@ -2,6 +2,7 @@
 name = "optd-cost-model"
 version = "0.1.0"
 edition = "2021"
+authors = ["Yuanxin Cao", "Lan Lou", "Kunle Li"]
 
 [dependencies]
 optd-persistent = { path = "../optd-persistent", version = "0.1" }
@@ -10,10 +11,15 @@ serde_json = "1.0"
 serde_with = { version = "3.7.0", features = ["json"] }
 arrow-schema = "53.2.0"
 datafusion-expr = "32.0.0"
+datafusion = "32.0.0"
 ordered-float = "4.0"
 chrono = "0.4"
 itertools = "0.13"
+assert_approx_eq = "1.1.0"
+trait-variant = "0.1.2"
+tokio = { version = "1.0.1", features = ["macros", "rt-multi-thread"] }
 
 [dev-dependencies]
 crossbeam = "0.8"
 rand = "0.8"
+test-case = "3.3"

optd-cost-model/src/common/nodes.rs

@@ -1,4 +1,5 @@
-use std::sync::Arc;
+use core::fmt;
+use std::{fmt::Display, sync::Arc};
 
 use arrow_schema::DataType;
 
@@ -24,6 +25,12 @@ pub enum JoinType {
     RightAnti,
 }
 
+impl Display for JoinType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", self)
+    }
+}
+
 /// TODO: documentation
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum PhysicalNodeType {
@@ -49,8 +56,7 @@ impl std::fmt::Display for PhysicalNodeType {
 pub enum PredicateType {
     List,
     Constant(ConstantType),
-    AttributeRef,
-    ExternAttributeRef,
+    AttrIndex,
     UnOp(UnOpType),
     BinOp(BinOpType),
     LogOp(LogOpType),
@@ -77,7 +83,7 @@ pub struct PredicateNode {
     /// A generic predicate node type
     pub typ: PredicateType,
     /// Child predicate nodes, always materialized
-    pub children: Vec<PredicateNode>,
+    pub children: Vec<ArcPredicateNode>,
     /// Data associated with the predicate, if any
     pub data: Option<Value>,
 }
@@ -94,3 +100,28 @@ impl std::fmt::Display for PredicateNode {
         write!(f, ")")
     }
 }
+
+impl PredicateNode {
+    pub fn child(&self, idx: usize) -> ArcPredicateNode {
+        self.children[idx].clone()
+    }
+
+    pub fn unwrap_data(&self) -> Value {
+        self.data.clone().unwrap()
+    }
+}
+pub trait ReprPredicateNode: 'static + Clone {
+    fn into_pred_node(self) -> ArcPredicateNode;
+
+    fn from_pred_node(pred_node: ArcPredicateNode) -> Option<Self>;
+}
+
+impl ReprPredicateNode for ArcPredicateNode {
+    fn into_pred_node(self) -> ArcPredicateNode {
+        self
+    }
+
+    fn from_pred_node(pred_node: ArcPredicateNode) -> Option<Self> {
+        Some(pred_node)
+    }
+}

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

@@ -0,0 +1,42 @@
+use crate::common::{
+    nodes::{ArcPredicateNode, PredicateNode, PredicateType, ReprPredicateNode},
+    values::Value,
+};
+
+/// [`AttributeIndexPred`] represents the position of an attribute in a schema or
+/// [`GroupAttrRefs`].
+///
+/// The `data` field holds the index of the attribute in the schema or [`GroupAttrRefs`].
+#[derive(Clone, Debug)]
+pub struct AttrIndexPred(pub ArcPredicateNode);
+
+impl AttrIndexPred {
+    pub fn new(attr_idx: u64) -> AttrIndexPred {
+        AttrIndexPred(
+            PredicateNode {
+                typ: PredicateType::AttrIndex,
+                children: vec![],
+                data: Some(Value::UInt64(attr_idx)),
+            }
+            .into(),
+        )
+    }
+
+    /// Gets the attribute index.
+    pub fn attr_index(&self) -> u64 {
+        self.0.data.as_ref().unwrap().as_u64()
+    }
+}
+
+impl ReprPredicateNode for AttrIndexPred {
+    fn into_pred_node(self) -> ArcPredicateNode {
+        self.0
+    }
+
+    fn from_pred_node(pred_node: ArcPredicateNode) -> Option<Self> {
+        if pred_node.typ != PredicateType::AttrIndex {
+            return None;
+        }
+        Some(Self(pred_node))
+    }
+}

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

@@ -1,3 +1,5 @@
+use crate::common::nodes::{ArcPredicateNode, PredicateNode, PredicateType, ReprPredicateNode};
+
 /// TODO: documentation
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 pub enum BinOpType {
@@ -38,3 +40,48 @@ impl BinOpType {
         )
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct BinOpPred(pub ArcPredicateNode);
+
+impl BinOpPred {
+    pub fn new(left: ArcPredicateNode, right: ArcPredicateNode, op_type: BinOpType) -> Self {
+        BinOpPred(
+            PredicateNode {
+                typ: PredicateType::BinOp(op_type),
+                children: vec![left, right],
+                data: None,
+            }
+            .into(),
+        )
+    }
+
+    pub fn left_child(&self) -> ArcPredicateNode {
+        self.0.child(0)
+    }
+
+    pub fn right_child(&self) -> ArcPredicateNode {
+        self.0.child(1)
+    }
+
+    pub fn op_type(&self) -> BinOpType {
+        if let PredicateType::BinOp(op_type) = self.0.typ {
+            op_type
+        } else {
+            panic!("not a bin op")
+        }
+    }
+}
+
+impl ReprPredicateNode for BinOpPred {
+    fn into_pred_node(self) -> ArcPredicateNode {
+        self.0
+    }
+
+    fn from_pred_node(pred_node: ArcPredicateNode) -> Option<Self> {
+        if !matches!(pred_node.typ, PredicateType::BinOp(_)) {
+            return None;
+        }
+        Some(Self(pred_node))
+    }
+}

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

@@ -0,0 +1,49 @@
+use arrow_schema::DataType;
+
+use crate::common::nodes::{ArcPredicateNode, PredicateNode, PredicateType, ReprPredicateNode};
+
+use super::data_type_pred::DataTypePred;
+
+/// [`CastPred`] casts a column from one data type to another.
+///
+/// A [`CastPred`] has two children:
+/// 1. The original data to cast
+/// 2. The target data type to cast to
+#[derive(Clone, Debug)]
+pub struct CastPred(pub ArcPredicateNode);
+
+impl CastPred {
+    pub fn new(child: ArcPredicateNode, cast_to: DataType) -> Self {
+        CastPred(
+            PredicateNode {
+                typ: PredicateType::Cast,
+                children: vec![child, DataTypePred::new(cast_to).into_pred_node()],
+                data: None,
+            }
+            .into(),
+        )
+    }
+
+    pub fn child(&self) -> ArcPredicateNode {
+        self.0.child(0)
+    }
+
+    pub fn cast_to(&self) -> DataType {
+        DataTypePred::from_pred_node(self.0.child(1))
+            .unwrap()
+            .data_type()
+    }
+}
+
+impl ReprPredicateNode for CastPred {
+    fn into_pred_node(self) -> ArcPredicateNode {
+        self.0
+    }
+
+    fn from_pred_node(pred_node: ArcPredicateNode) -> Option<Self> {
+        if !matches!(pred_node.typ, PredicateType::Cast) {
+            return None;
+        }
+        Some(Self(pred_node))
+    }
+}