split plan node representation into a crate

Cargo.toml

@@ -1,3 +1,8 @@
 [workspace]
-members = ["datafusion-optd-cli", "optd-core", "datafusion-optd-bridge"]
+members = [
+    "datafusion-optd-cli",
+    "optd-core",
+    "optd-datafusion-bridge",
+    "optd-datafusion-repr",
+]
 resolver = "2"

README.md

@@ -1,3 +1,10 @@
 # optd
 
-A query optimizer. Currently aimed for a cascades optimizer for Apache datafusion.
+A query optimizer. Currently aimed for a cascades optimizer for Apache Arrow Datafusion.
+
+## Structure
+
+* `datafusion-optd-cli`: patched Apache Arrow Datafusion cli that calls into optd
+* `datafusion-optd-bridge`: implementation of Apache Arrow Datafusion query planner as a bridge between optd and Apache Arrow Datafusion.
+* `optd-core`: the core framework of optd.
+* `optd-datafusion-repr`: representation of Apache Arrow Datafusion plan nodes in optd.

datafusion-optd-cli/Cargo.toml

@@ -57,7 +57,7 @@ tokio = { version = "1.24", features = [
     "parking_lot",
 ] }
 url = "2.2"
-datafusion-optd-bridge = { path = "../datafusion-optd-bridge" }
+optd-datafusion-bridge = { path = "../optd-datafusion-bridge" }
 
 [dev-dependencies]
 assert_cmd = "2.0"

datafusion-optd-cli/src/main.rs

@@ -21,7 +21,6 @@ use datafusion::execution::context::{SessionConfig, SessionState};
 use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool};
 use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
 use datafusion::prelude::SessionContext;
-use datafusion_optd_bridge::OptdQueryPlanner;
 use datafusion_optd_cli::catalog::DynamicFileCatalog;
 use datafusion_optd_cli::{
     exec,
@@ -30,6 +29,7 @@ use datafusion_optd_cli::{
     DATAFUSION_CLI_VERSION,
 };
 use mimalloc::MiMalloc;
+use optd_datafusion_bridge::OptdQueryPlanner;
 use std::collections::HashMap;
 use std::env;
 use std::path::Path;

optd-core/src/cascades/tasks/apply_rule.rs

@@ -11,7 +11,7 @@ use crate::{
         tasks::{OptimizeExpressionTask, OptimizeInputsTask},
         GroupId,
     },
-    rel_node::{RelNode, RelNodeRef, RelNodeTyp},
+    rel_node::{RelNode, RelNodeTyp},
     rules::{OneOrMany, RuleMatcher},
 };
 

optd-core/src/cost.rs

@@ -1,26 +1,6 @@
-use std::collections::HashMap;
-
 use itertools::Itertools;
 
-use crate::{
-    plan_nodes::OptRelNodeTyp,
-    rel_node::{RelNode, RelNodeTyp, Value},
-};
-
-fn compute_plan_node_cost<T: RelNodeTyp, C: CostModel<T>>(
-    model: &C,
-    node: &RelNode<T>,
-    total_cost: &mut Cost,
-) -> Cost {
-    let children = node
-        .children
-        .iter()
-        .map(|child| compute_plan_node_cost(model, child, total_cost))
-        .collect_vec();
-    let cost = model.compute_cost(&node.typ, &node.data, &children);
-    model.accumulate(total_cost, &cost);
-    cost
-}
+use crate::rel_node::{RelNode, RelNodeTyp, Value};
 
 #[derive(Default, Clone, Debug, PartialOrd, PartialEq)]
 pub struct Cost(pub Vec<f64>);
@@ -36,108 +16,3 @@ pub trait CostModel<T: RelNodeTyp>: 'static + Send + Sync {
 
     fn zero(&self) -> Cost;
 }
-
-pub struct OptCostModel {
-    table_stat: HashMap<String, usize>,
-}
-
-pub const ROW_COUNT: usize = 1;
-pub const COMPUTE_COST: usize = 2;
-pub const IO_COST: usize = 3;
-
-impl OptCostModel {
-    fn row_cnt(Cost(cost): &Cost) -> f64 {
-        cost[ROW_COUNT]
-    }
-
-    fn compute_cost(Cost(cost): &Cost) -> f64 {
-        cost[COMPUTE_COST]
-    }
-
-    fn io_cost(Cost(cost): &Cost) -> f64 {
-        cost[IO_COST]
-    }
-
-    fn cost_tuple(Cost(cost): &Cost) -> (f64, f64, f64) {
-        (cost[ROW_COUNT], cost[COMPUTE_COST], cost[IO_COST])
-    }
-
-    fn weighted_cost(row_cnt: f64, compute_cost: f64, io_cost: f64) -> f64 {
-        let _ = row_cnt;
-        compute_cost + io_cost * 10.0
-    }
-
-    pub fn cost(row_cnt: f64, compute_cost: f64, io_cost: f64) -> Cost {
-        Cost(vec![
-            Self::weighted_cost(row_cnt, compute_cost, io_cost),
-            row_cnt,
-            compute_cost,
-            io_cost,
-        ])
-    }
-}
-
-impl CostModel<OptRelNodeTyp> for OptCostModel {
-    fn explain(&self, cost: &Cost) -> String {
-        format!(
-            "weighted={},row_cnt={},compute={},io={}",
-            cost.0[0],
-            Self::row_cnt(cost),
-            Self::compute_cost(cost),
-            Self::io_cost(cost)
-        )
-    }
-
-    fn accumulate(&self, total_cost: &mut Cost, cost: &Cost) {
-        total_cost.0[ROW_COUNT] += Self::row_cnt(cost);
-        total_cost.0[COMPUTE_COST] += Self::compute_cost(cost);
-        total_cost.0[IO_COST] += Self::io_cost(cost);
-        total_cost.0[0] = Self::weighted_cost(
-            total_cost.0[ROW_COUNT],
-            total_cost.0[COMPUTE_COST],
-            total_cost.0[IO_COST],
-        );
-    }
-
-    fn zero(&self) -> Cost {
-        Self::cost(0.0, 0.0, 0.0)
-    }
-
-    fn compute_cost(&self, node: &OptRelNodeTyp, data: &Option<Value>, children: &[Cost]) -> Cost {
-        match node {
-            OptRelNodeTyp::PhysicalScan => {
-                let table_name = data.as_ref().unwrap().as_str();
-                let row_cnt = self.table_stat.get(table_name.as_ref()).copied().unwrap() as f64;
-                Self::cost(row_cnt, 0.0, row_cnt)
-            }
-            OptRelNodeTyp::PhysicalFilter => {
-                let (row_cnt, _, _) = Self::cost_tuple(&children[0]);
-                let selectivity = 0.1;
-                Self::cost(row_cnt * selectivity, row_cnt, 0.0)
-            }
-            OptRelNodeTyp::PhysicalNestedLoopJoin(_) => {
-                let (row_cnt_1, _, _) = Self::cost_tuple(&children[0]);
-                let (row_cnt_2, _, _) = Self::cost_tuple(&children[1]);
-                let selectivity = 0.1;
-                Self::cost(
-                    row_cnt_1 * row_cnt_2 * selectivity,
-                    row_cnt_1 * row_cnt_2,
-                    0.0,
-                )
-            }
-            _ => Self::cost(1.0, 0.0, 0.0),
-        }
-    }
-
-    fn compute_plan_node_cost(&self, node: &RelNode<OptRelNodeTyp>) -> Cost {
-        let mut cost = self.zero();
-        compute_plan_node_cost(self, node, &mut cost);
-        cost
-    }
-}
-
-impl OptCostModel {
-    pub fn new(table_stat: HashMap<String, usize>) -> Self {
-        Self { table_stat }
-    }
-}

optd-core/src/lib.rs

@@ -1,5 +1,4 @@
 pub mod cascades;
 pub mod cost;
-pub mod plan_nodes;
 pub mod rel_node;
 pub mod rules;

optd-core/src/rules.rs

@@ -1,18 +1,10 @@
-mod filter_join;
 mod ir;
-mod join_assoc;
-mod join_commute;
-mod physical;
 
 use std::collections::HashMap;
 
 use crate::rel_node::{RelNode, RelNodeTyp};
 
-pub use filter_join::FilterJoinPullUpRule;
 pub use ir::{OneOrMany, RuleMatcher};
-pub use join_assoc::{JoinAssocLeftRule, JoinAssocRightRule};
-pub use join_commute::JoinCommuteRule;
-pub use physical::PhysicalConversionRule;
 
 pub trait Rule<T: RelNodeTyp> {
     fn matcher(&self) -> &RuleMatcher<T>;

datafusion-optd-bridge/Cargo.toml → optd-datafusion-bridge/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "datafusion-optd-bridge"
+name = "optd-datafusion-bridge"
 version = "0.1.0"
 edition = "2021"
 

optd-datafusion-repr/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "optd-datafusion-repr"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1"
+num-traits = "0.2"
+num-derive = "0.2"
+tracing = "0.1"
+ordered-float = "4"
+tracing-subscriber = "0.3"
+pretty-xmlish = "0.1"
+itertools = "0.11"
+optd-core = { path = "../optd-core" }

optd-core/src/bin/test_optimize.rs → optd-datafusion-repr/src/bin/test_optimize.rs

@@ -1,13 +1,12 @@
 use std::sync::Arc;
 
-use optd_core::{
-    cascades::CascadesOptimizer,
+use optd_core::{cascades::CascadesOptimizer, rel_node::Value};
+use optd_datafusion_repr::{
     cost::OptCostModel,
     plan_nodes::{
         BinOpExpr, BinOpType, ColumnRefExpr, ConstantExpr, JoinType, LogicalFilter, LogicalJoin,
         LogicalScan, OptRelNode, OptRelNodeTyp, PlanNode,
     },
-    rel_node::Value,
     rules::{
         FilterJoinPullUpRule, JoinAssocLeftRule, JoinAssocRightRule, JoinCommuteRule,
         PhysicalConversionRule,

optd-datafusion-repr/src/cost.rs

@@ -0,0 +1,128 @@
+use std::collections::HashMap;
+
+use crate::plan_nodes::OptRelNodeTyp;
+use itertools::Itertools;
+use optd_core::{
+    cost::{Cost, CostModel},
+    rel_node::{RelNode, RelNodeTyp, Value},
+};
+
+fn compute_plan_node_cost<T: RelNodeTyp, C: CostModel<T>>(
+    model: &C,
+    node: &RelNode<T>,
+    total_cost: &mut Cost,
+) -> Cost {
+    let children = node
+        .children
+        .iter()
+        .map(|child| compute_plan_node_cost(model, child, total_cost))
+        .collect_vec();
+    let cost = model.compute_cost(&node.typ, &node.data, &children);
+    model.accumulate(total_cost, &cost);
+    cost
+}
+
+pub struct OptCostModel {
+    table_stat: HashMap<String, usize>,
+}
+
+pub const ROW_COUNT: usize = 1;
+pub const COMPUTE_COST: usize = 2;
+pub const IO_COST: usize = 3;
+
+impl OptCostModel {
+    fn row_cnt(Cost(cost): &Cost) -> f64 {
+        cost[ROW_COUNT]
+    }
+
+    fn compute_cost(Cost(cost): &Cost) -> f64 {
+        cost[COMPUTE_COST]
+    }
+
+    fn io_cost(Cost(cost): &Cost) -> f64 {
+        cost[IO_COST]
+    }
+
+    fn cost_tuple(Cost(cost): &Cost) -> (f64, f64, f64) {
+        (cost[ROW_COUNT], cost[COMPUTE_COST], cost[IO_COST])
+    }
+
+    fn weighted_cost(row_cnt: f64, compute_cost: f64, io_cost: f64) -> f64 {
+        let _ = row_cnt;
+        compute_cost + io_cost * 10.0
+    }
+
+    pub fn cost(row_cnt: f64, compute_cost: f64, io_cost: f64) -> Cost {
+        Cost(vec![
+            Self::weighted_cost(row_cnt, compute_cost, io_cost),
+            row_cnt,
+            compute_cost,
+            io_cost,
+        ])
+    }
+}
+
+impl CostModel<OptRelNodeTyp> for OptCostModel {
+    fn explain(&self, cost: &Cost) -> String {
+        format!(
+            "weighted={},row_cnt={},compute={},io={}",
+            cost.0[0],
+            Self::row_cnt(cost),
+            Self::compute_cost(cost),
+            Self::io_cost(cost)
+        )
+    }
+
+    fn accumulate(&self, total_cost: &mut Cost, cost: &Cost) {
+        total_cost.0[ROW_COUNT] += Self::row_cnt(cost);
+        total_cost.0[COMPUTE_COST] += Self::compute_cost(cost);
+        total_cost.0[IO_COST] += Self::io_cost(cost);
+        total_cost.0[0] = Self::weighted_cost(
+            total_cost.0[ROW_COUNT],
+            total_cost.0[COMPUTE_COST],
+            total_cost.0[IO_COST],
+        );
+    }
+
+    fn zero(&self) -> Cost {
+        Self::cost(0.0, 0.0, 0.0)
+    }
+
+    fn compute_cost(&self, node: &OptRelNodeTyp, data: &Option<Value>, children: &[Cost]) -> Cost {
+        match node {
+            OptRelNodeTyp::PhysicalScan => {
+                let table_name = data.as_ref().unwrap().as_str();
+                let row_cnt = self.table_stat.get(table_name.as_ref()).copied().unwrap() as f64;
+                Self::cost(row_cnt, 0.0, row_cnt)
+            }
+            OptRelNodeTyp::PhysicalFilter => {
+                let (row_cnt, _, _) = Self::cost_tuple(&children[0]);
+                let selectivity = 0.1;
+                Self::cost(row_cnt * selectivity, row_cnt, 0.0)
+            }
+            OptRelNodeTyp::PhysicalNestedLoopJoin(_) => {
+                let (row_cnt_1, _, _) = Self::cost_tuple(&children[0]);
+                let (row_cnt_2, _, _) = Self::cost_tuple(&children[1]);
+                let selectivity = 0.1;
+                Self::cost(
+                    row_cnt_1 * row_cnt_2 * selectivity,
+                    row_cnt_1 * row_cnt_2,
+                    0.0,
+                )
+            }
+            _ => Self::cost(1.0, 0.0, 0.0),
+        }
+    }
+
+    fn compute_plan_node_cost(&self, node: &RelNode<OptRelNodeTyp>) -> Cost {
+        let mut cost = self.zero();
+        compute_plan_node_cost(self, node, &mut cost);
+        cost
+    }
+}
+
+impl OptCostModel {
+    pub fn new(table_stat: HashMap<String, usize>) -> Self {
+        Self { table_stat }
+    }
+}

optd-datafusion-repr/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod cost;
+pub mod plan_nodes;
+pub mod rules;