feat(cost-model): migrate pred types (#39)

Author: xx01cyx

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))
+    }
+}