Separate advanced cost model from the previous basic cost model (#200)

Tested only briefly, quick and dirty

Author: jurplel

Cargo.lock

@@ -8,5 +8,6 @@ members = [
     "optd-adaptive-demo",
     "optd-gungnir",
     "optd-perfbench",
+    "optd-datafusion-repr-adv-cost",
 ]
 resolver = "2"

Cargo.toml

@@ -58,6 +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 = { path = "../optd-datafusion-repr" }
 tracing-subscriber = "0.3"
 tracing = "0.1"

datafusion-optd-cli/Cargo.toml

@@ -30,8 +30,8 @@ use datafusion_optd_cli::{
 };
 use mimalloc::MiMalloc;
 use optd_datafusion_bridge::{DatafusionCatalog, OptdQueryPlanner};
-use optd_datafusion_repr::cost::BaseTableStats;
 use optd_datafusion_repr::DatafusionOptimizer;
+use optd_datafusion_repr_adv_cost::adv_cost::stats::BaseTableStats;
 use std::collections::HashMap;
 use std::env;
 use std::path::Path;
@@ -141,6 +141,9 @@ struct Args {
 
     #[clap(long, help = "Turn on adaptive optimization")]
     enable_adaptive: bool,
+
+    #[clap(long, help = "Use advanced cost model")]
+    adv_cost: bool,
 }
 
 #[tokio::main]
@@ -204,11 +207,18 @@ pub async fn main() -> Result<()> {
             state = state.with_physical_optimizer_rules(vec![]);
         }
         // use optd-bridge query planner
-        let optimizer = DatafusionOptimizer::new_physical(
-            Arc::new(DatafusionCatalog::new(state.catalog_list())),
-            BaseTableStats::default(),
-            args.enable_adaptive,
-        );
+        let optimizer = if args.adv_cost {
+            optd_datafusion_repr_adv_cost::new_physical_adv_cost(
+                Arc::new(DatafusionCatalog::new(state.catalog_list())),
+                BaseTableStats::default(),
+                args.enable_adaptive,
+            )
+        } else {
+            DatafusionOptimizer::new_physical(
+                Arc::new(DatafusionCatalog::new(state.catalog_list())),
+                args.enable_adaptive,
+            )
+        };
         state = state.with_query_planner(Arc::new(OptdQueryPlanner::new(optimizer)));
         SessionContext::new_with_state(state)
     };

datafusion-optd-cli/src/main.rs

@@ -10,7 +10,6 @@ use datafusion_optd_cli::{
 };
 use mimalloc::MiMalloc;
 use optd_datafusion_bridge::{DatafusionCatalog, OptdQueryPlanner};
-use optd_datafusion_repr::cost::BaseTableStats;
 use optd_datafusion_repr::DatafusionOptimizer;
 use rand::{thread_rng, Rng};
 use std::sync::Arc;
@@ -30,7 +29,6 @@ async fn main() -> Result<()> {
             SessionState::new_with_config_rt(session_config.clone(), Arc::new(runtime_env));
         let mut optimizer: DatafusionOptimizer = DatafusionOptimizer::new_physical(
             Arc::new(DatafusionCatalog::new(state.catalog_list())),
-            BaseTableStats::default(),
             true,
         );
         optimizer.optd_optimizer_mut().prop.partial_explore_iter = None;

optd-adaptive-demo/src/bin/optd-adaptive-three-join.rs

@@ -10,7 +10,6 @@ use datafusion_optd_cli::{
 };
 use mimalloc::MiMalloc;
 use optd_datafusion_bridge::{DatafusionCatalog, OptdQueryPlanner};
-use optd_datafusion_repr::cost::BaseTableStats;
 use optd_datafusion_repr::DatafusionOptimizer;
 use std::sync::Arc;
 use std::time::Duration;
@@ -31,7 +30,6 @@ async fn main() -> Result<()> {
             SessionState::new_with_config_rt(session_config.clone(), Arc::new(runtime_env));
         let optimizer: DatafusionOptimizer = DatafusionOptimizer::new_physical(
             Arc::new(DatafusionCatalog::new(state.catalog_list())),
-            BaseTableStats::default(),
             true,
         );
         // clean up optimizer rules so that we can plug in our own optimizer

optd-adaptive-demo/src/bin/optd-adaptive-tpch-q8.rs

@@ -0,0 +1,23 @@
+[package]
+name = "optd-datafusion-repr-adv-cost"
+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"
+arrow-schema = "47.0.0"
+assert_approx_eq = "1.1.0"
+datafusion = "32.0.0"
+ordered-float = "4"
+optd-datafusion-repr = { path = "../optd-datafusion-repr" }
+optd-core = { path = "../optd-core" }
+serde = { version = "1.0", features = ["derive"] }
+rayon = "1.10"
+itertools = "0.11"
+test-case = "3.3"
+tracing = "0.1"
+tracing-subscriber = "0.3"
+optd-gungnir = { path = "../optd-gungnir" }
+serde_with = {version = "3.7.0", features = ["json"]}

optd-datafusion-repr-adv-cost/Cargo.toml

@@ -0,0 +1,108 @@
+use std::sync::{Arc, Mutex};
+
+use crate::adv_cost::OptCostModel;
+use optd_core::{
+    cascades::{CascadesOptimizer, RelNodeContext},
+    cost::{Cost, CostModel},
+    rel_node::{RelNode, Value},
+};
+use optd_datafusion_repr::{
+    cost::adaptive_cost::RuntimeAdaptionStorageInner, plan_nodes::OptRelNodeTyp,
+};
+use serde::{de::DeserializeOwned, Serialize};
+
+use super::adv_cost::stats::{
+    BaseTableStats, DataFusionDistribution, DataFusionMostCommonValues, Distribution,
+    MostCommonValues,
+};
+
+pub type RuntimeAdaptionStorage = Arc<Mutex<RuntimeAdaptionStorageInner>>;
+pub type DataFusionAdaptiveCostModel =
+    AdaptiveCostModel<DataFusionMostCommonValues, DataFusionDistribution>;
+
+pub const DEFAULT_DECAY: usize = 50;
+
+pub struct AdaptiveCostModel<
+    M: MostCommonValues + Serialize + DeserializeOwned,
+    D: Distribution + Serialize + DeserializeOwned,
+> {
+    runtime_row_cnt: RuntimeAdaptionStorage,
+    base_model: OptCostModel<M, D>,
+    decay: usize,
+}
+
+impl<
+        M: MostCommonValues + Serialize + DeserializeOwned,
+        D: Distribution + Serialize + DeserializeOwned,
+    > CostModel<OptRelNodeTyp> for AdaptiveCostModel<M, D>
+{
+    fn explain(&self, cost: &Cost) -> String {
+        self.base_model.explain(cost)
+    }
+
+    fn accumulate(&self, total_cost: &mut Cost, cost: &Cost) {
+        self.base_model.accumulate(total_cost, cost)
+    }
+
+    fn zero(&self) -> Cost {
+        self.base_model.zero()
+    }
+
+    fn compute_cost(
+        &self,
+        node: &OptRelNodeTyp,
+        data: &Option<Value>,
+        children: &[Cost],
+        context: Option<RelNodeContext>,
+        optimizer: Option<&CascadesOptimizer<OptRelNodeTyp>>,
+    ) -> Cost {
+        if let OptRelNodeTyp::PhysicalScan = node {
+            let guard = self.runtime_row_cnt.lock().unwrap();
+            if let Some((runtime_row_cnt, iter)) =
+                guard.history.get(&context.as_ref().unwrap().group_id)
+            {
+                if *iter + self.decay >= guard.iter_cnt {
+                    let runtime_row_cnt = (*runtime_row_cnt).max(1) as f64;
+                    return OptCostModel::<M, D>::cost(runtime_row_cnt, 0.0, runtime_row_cnt);
+                }
+            }
+        }
+        let (mut row_cnt, compute_cost, io_cost) = OptCostModel::<M, D>::cost_tuple(
+            &self
+                .base_model
+                .compute_cost(node, data, children, context.clone(), optimizer),
+        );
+        if let Some(context) = context {
+            let guard = self.runtime_row_cnt.lock().unwrap();
+            if let Some((runtime_row_cnt, iter)) = guard.history.get(&context.group_id) {
+                if *iter + self.decay >= guard.iter_cnt {
+                    let runtime_row_cnt = (*runtime_row_cnt).max(1) as f64;
+                    row_cnt = runtime_row_cnt;
+                }
+            }
+        }
+        OptCostModel::<M, D>::cost(row_cnt, compute_cost, io_cost)
+    }
+
+    fn compute_plan_node_cost(&self, node: &RelNode<OptRelNodeTyp>) -> Cost {
+        self.base_model.compute_plan_node_cost(node)
+    }
+}
+
+impl<
+        M: MostCommonValues + Serialize + DeserializeOwned,
+        D: Distribution + Serialize + DeserializeOwned,
+    > AdaptiveCostModel<M, D>
+{
+    pub fn new(decay: usize, stats: BaseTableStats<M, D>) -> Self {
+        Self {
+            runtime_row_cnt: RuntimeAdaptionStorage::default(),
+            base_model: OptCostModel::new(stats),
+            decay,
+        }
+    }
+
+    pub fn get_runtime_map(&self) -> RuntimeAdaptionStorage {
+        self.runtime_row_cnt.clone()
+    }
+}