feat: implement card test for df optimizer (#105)

Implement card test interfaces for Datafusion optimizer. After #104, the
cost of each plan node can be displayed using `explain verbose`.
Utilizing this feature, we can extract the estimated cardinality of the
root plan node. The true cardinality is computed by counting the rows of
the result set.

Also enable users to select queries to execute in perftest by specifying
e.g. `--query-ids 2 3`


![image](https://github.com/cmu-db/optd/assets/52783948/00567813-d10d-4f99-9024-21baf4f5e4f9)

Author: Gun9niR

Cargo.lock

@@ -6,7 +6,17 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-optd-sqlplannertest = { path = "../optd-sqlplannertest" }
+datafusion = { version = "32.0.0", features = [
+    "avro",
+    "crypto_expressions",
+    "encoding_expressions",
+    "regex_expressions",
+    "unicode_expressions",
+    "compression",
+] }
+optd-datafusion-repr = { path = "../optd-datafusion-repr" }
+optd-datafusion-bridge = { path = "../optd-datafusion-bridge" }
+datafusion-optd-cli = { path = "../datafusion-optd-cli" }
 futures = "0.3"
 anyhow = { version = "1", features = ["backtrace"] }
 async-trait = "0.1"
@@ -25,3 +35,4 @@ clap = { version = "4.5", features = [
 ] }
 log = "0.4"
 env_logger = "0.11"
+lazy_static = "1.4.0"

optd-perftest/Cargo.toml

@@ -77,6 +77,12 @@ pub trait CardtestRunnerDBHelper {
     fn get_name(&self) -> &str;
 
     // The order of queries has to be the same between these two functions.
-    async fn eval_benchmark_estcards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>>;
-    async fn eval_benchmark_truecards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>>;
+    async fn eval_benchmark_estcards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>>;
+    async fn eval_benchmark_truecards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>>;
 }

optd-perftest/src/cardtest.rs

@@ -1,32 +1,246 @@
-use crate::{benchmark::Benchmark, cardtest::CardtestRunnerDBHelper};
+use std::{
+    fs,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+
+use crate::{
+    benchmark::Benchmark,
+    cardtest::CardtestRunnerDBHelper,
+    tpch::{TpchConfig, TpchKit},
+};
 use async_trait::async_trait;
-use optd_sqlplannertest::DatafusionDb;
+use datafusion::{
+    arrow::util::display::{ArrayFormatter, FormatOptions},
+    execution::{
+        config::SessionConfig,
+        context::{SessionContext, SessionState},
+        runtime_env::{RuntimeConfig, RuntimeEnv},
+    },
+    sql::{parser::DFParser, sqlparser::dialect::GenericDialect},
+};
+use datafusion_optd_cli::helper::unescape_input;
+use lazy_static::lazy_static;
+use optd_datafusion_bridge::{DatafusionCatalog, OptdQueryPlanner};
+use optd_datafusion_repr::DatafusionOptimizer;
+use regex::Regex;
+
+pub struct DatafusionDb {
+    workspace_dpath: PathBuf,
+    ctx: SessionContext,
+}
 
 #[async_trait]
 impl CardtestRunnerDBHelper for DatafusionDb {
     fn get_name(&self) -> &str {
         "DataFusion"
     }
 
-    async fn eval_benchmark_truecards(&self, _benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
-        Ok(vec![])
+    async fn eval_benchmark_estcards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
+        self.clear_state().await?;
+        self.load_benchmark_data(benchmark).await?;
+        match benchmark {
+            Benchmark::Test => unimplemented!(),
+            Benchmark::Tpch(tpch_config) => self.eval_tpch_estcards(tpch_config).await,
+        }
     }
 
-    async fn eval_benchmark_estcards(&self, _benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
-        Ok(vec![])
+    async fn eval_benchmark_truecards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
+        self.clear_state().await?;
+        self.load_benchmark_data(benchmark).await?;
+        match benchmark {
+            Benchmark::Test => unimplemented!(),
+            Benchmark::Tpch(tpch_config) => self.eval_tpch_truecards(tpch_config).await,
+        }
     }
 }
 
-// helper functions for ```impl CardtestRunnerDBHelper for DatafusionDb```
-// they can't be put in an impl because DatafusionDb is a foreign struct
-async fn _eval_query_truecard(db: &DatafusionDb, sql: &str) -> anyhow::Result<usize> {
-    let rows = db.execute(sql, true).await?;
-    let num_rows = rows.len();
-    Ok(num_rows)
-}
+impl DatafusionDb {
+    pub async fn new<P: AsRef<Path>>(workspace_dpath: P) -> anyhow::Result<Self> {
+        Ok(DatafusionDb {
+            workspace_dpath: workspace_dpath.as_ref().to_path_buf(),
+            ctx: Self::new_session_ctx().await?,
+        })
+    }
+
+    /// Reset data and metadata.
+    async fn clear_state(&mut self) -> anyhow::Result<()> {
+        self.ctx = Self::new_session_ctx().await?;
+        Ok(())
+    }
+
+    async fn new_session_ctx() -> anyhow::Result<SessionContext> {
+        let session_config = SessionConfig::from_env()?.with_information_schema(true);
+        let rn_config = RuntimeConfig::new();
+        let runtime_env = RuntimeEnv::new(rn_config.clone())?;
+        let ctx = {
+            let mut state =
+                SessionState::new_with_config_rt(session_config.clone(), Arc::new(runtime_env));
+            let optimizer = DatafusionOptimizer::new_physical(Arc::new(DatafusionCatalog::new(
+                state.catalog_list(),
+            )));
+            state = state.with_physical_optimizer_rules(vec![]);
+            state = state.with_query_planner(Arc::new(OptdQueryPlanner::new(optimizer)));
+            SessionContext::new_with_state(state)
+        };
+        ctx.refresh_catalogs().await?;
+        Ok(ctx)
+    }
 
-async fn _eval_query_estcard(db: &DatafusionDb, _sql: &str) -> anyhow::Result<usize> {
-    let rows = db.execute("EXPLAIN SELECT * FROM t1;", true).await?;
-    println!("eval_est_card(): rows={:?}", rows);
-    Ok(12)
+    async fn execute(&self, sql: &str) -> anyhow::Result<Vec<Vec<String>>> {
+        let sql = unescape_input(sql)?;
+        let dialect = Box::new(GenericDialect);
+        let statements = DFParser::parse_sql_with_dialect(&sql, dialect.as_ref())?;
+        let mut result = Vec::new();
+        for statement in statements {
+            let df = {
+                let plan = self.ctx.state().statement_to_plan(statement).await?;
+                self.ctx.execute_logical_plan(plan).await?
+            };
+
+            let batches = df.collect().await?;
+
+            let options = FormatOptions::default();
+
+            for batch in batches {
+                let converters = batch
+                    .columns()
+                    .iter()
+                    .map(|a| ArrayFormatter::try_new(a.as_ref(), &options))
+                    .collect::<Result<Vec<_>, _>>()?;
+                for row_idx in 0..batch.num_rows() {
+                    let mut row = Vec::with_capacity(batch.num_columns());
+                    for converter in converters.iter() {
+                        let mut buffer = String::with_capacity(8);
+                        converter.value(row_idx).write(&mut buffer)?;
+                        row.push(buffer);
+                    }
+                    result.push(row);
+                }
+            }
+        }
+        Ok(result)
+    }
+
+    async fn eval_tpch_estcards(&self, tpch_config: &TpchConfig) -> anyhow::Result<Vec<usize>> {
+        let tpch_kit = TpchKit::build(&self.workspace_dpath)?;
+        tpch_kit.gen_queries(tpch_config)?;
+
+        let mut estcards = vec![];
+        for sql_fpath in tpch_kit.get_sql_fpath_ordered_iter(tpch_config)? {
+            let sql = fs::read_to_string(sql_fpath)?;
+            let estcard = self.eval_query_estcard(&sql).await?;
+            estcards.push(estcard);
+        }
+
+        Ok(estcards)
+    }
+
+    async fn eval_tpch_truecards(&self, tpch_config: &TpchConfig) -> anyhow::Result<Vec<usize>> {
+        let tpch_kit = TpchKit::build(&self.workspace_dpath)?;
+        tpch_kit.gen_queries(tpch_config)?;
+
+        let mut truecards = vec![];
+        for sql_fpath in tpch_kit.get_sql_fpath_ordered_iter(tpch_config)? {
+            let sql = fs::read_to_string(sql_fpath)?;
+            let estcard = self.eval_query_truecard(&sql).await?;
+            truecards.push(estcard);
+        }
+
+        Ok(truecards)
+    }
+
+    async fn eval_query_truecard(&self, sql: &str) -> anyhow::Result<usize> {
+        let rows = self.execute(sql).await?;
+        let num_rows = rows.len();
+        Ok(num_rows)
+    }
+
+    async fn eval_query_estcard(&self, sql: &str) -> anyhow::Result<usize> {
+        lazy_static! {
+            static ref ROW_CNT_RE: Regex = Regex::new(r"row_cnt=(\d+\.\d+)").unwrap();
+        }
+        let explains = self.execute(&format!("explain verbose {}", sql)).await?;
+        // Find first occurrence of row_cnt=... in the output.
+        let row_cnt = explains
+            .iter()
+            .find_map(|explain| {
+                // First element is task name, second is the actual explain output.
+                assert!(explain.len() == 2);
+                let explain = &explain[1];
+                if let Some(caps) = ROW_CNT_RE.captures(explain) {
+                    caps.get(1)
+                        .map(|row_cnt| row_cnt.as_str().parse::<f32>().unwrap() as usize)
+                } else {
+                    None
+                }
+            })
+            .unwrap();
+        Ok(row_cnt)
+    }
+
+    async fn load_benchmark_data(&mut self, benchmark: &Benchmark) -> anyhow::Result<()> {
+        match benchmark {
+            Benchmark::Tpch(tpch_config) => self.load_tpch_data(tpch_config).await,
+            _ => unimplemented!(),
+        }
+    }
+
+    async fn load_tpch_data(&mut self, tpch_config: &TpchConfig) -> anyhow::Result<()> {
+        let tpch_kit = TpchKit::build(&self.workspace_dpath)?;
+        tpch_kit.gen_tables(tpch_config)?;
+
+        // Create the tables.
+        let ddls = fs::read_to_string(&tpch_kit.schema_fpath)?;
+        let ddls = ddls
+            .split(';')
+            .map(|s| s.trim())
+            .filter(|s| !s.is_empty())
+            .collect::<Vec<_>>();
+        for ddl in ddls {
+            self.execute(ddl).await?;
+        }
+
+        // Load the data by creating an external table first and copying the data to real tables.
+        let tbl_fpath_iter = tpch_kit.get_tbl_fpath_iter(tpch_config).unwrap();
+        for tbl_fpath in tbl_fpath_iter {
+            let tbl_name = tbl_fpath.file_stem().unwrap().to_str().unwrap();
+            self.execute(&format!(
+                "create external table {}_tbl stored as csv delimiter '|' location '{}';",
+                tbl_name,
+                tbl_fpath.to_str().unwrap()
+            ))
+            .await?;
+
+            // Get the number of columns of this table.
+            let schema = self
+                .ctx
+                .catalog("datafusion")
+                .unwrap()
+                .schema("public")
+                .unwrap()
+                .table(tbl_name)
+                .await
+                .unwrap()
+                .schema();
+            let projection_list = (1..=schema.fields().len())
+                .map(|i| format!("column_{}", i))
+                .collect::<Vec<_>>()
+                .join(", ");
+            self.execute(&format!(
+                "insert into {} select {} from {}_tbl;",
+                tbl_name, projection_list, tbl_name,
+            ))
+            .await?;
+        }
+        Ok(())
+    }
 }
+
+unsafe impl Send for DatafusionDb {}

optd-perftest/src/datafusion_db_cardtest.rs

@@ -6,6 +6,7 @@ use postgres_db::PostgresDb;
 
 use crate::{
     benchmark::Benchmark,
+    datafusion_db_cardtest::DatafusionDb,
     tpch::{TpchConfig, TPCH_KIT_POSTGRES},
 };
 
@@ -37,6 +38,9 @@ enum Commands {
         #[arg(long)]
         #[clap(default_value = "15721")]
         seed: i32,
+        #[arg(long)]
+        #[clap(value_delimiter = ' ', num_args = 1..)]
+        query_ids: Vec<u32>,
     },
 }
 
@@ -50,24 +54,30 @@ async fn main() -> anyhow::Result<()> {
         fs::create_dir(&workspace_dpath)?;
     }
 
-    match &cli.command {
-        Commands::Cardtest { scale_factor, seed } => {
+    match cli.command {
+        Commands::Cardtest {
+            scale_factor,
+            seed,
+            query_ids,
+        } => {
             let tpch_config = TpchConfig {
                 database: String::from(TPCH_KIT_POSTGRES),
-                scale_factor: *scale_factor,
-                seed: *seed,
+                scale_factor,
+                seed,
+                query_ids,
             };
             cardtest(&workspace_dpath, tpch_config).await
         }
     }
 }
 
-async fn cardtest<P: AsRef<Path>>(
+async fn cardtest<P: AsRef<Path> + Clone>(
     workspace_dpath: P,
     tpch_config: TpchConfig,
 ) -> anyhow::Result<()> {
-    let pg_db = PostgresDb::new(workspace_dpath);
-    let databases: Vec<Box<dyn CardtestRunnerDBHelper>> = vec![Box::new(pg_db)];
+    let pg_db = PostgresDb::new(workspace_dpath.clone());
+    let df_db = DatafusionDb::new(workspace_dpath).await?;
+    let databases: Vec<Box<dyn CardtestRunnerDBHelper>> = vec![Box::new(pg_db), Box::new(df_db)];
 
     let tpch_benchmark = Benchmark::Tpch(tpch_config.clone());
     let mut cardtest_runner = CardtestRunner::new(databases).await?;