Own SQL benchmark suite

.github/workflows/sql-benchmarks.yml

@@ -298,6 +298,25 @@ on:
                 {"engine": "duckdb", "format": "duckdb"}
               ],
               "iterations": "10"
+            },
+            {
+              "id": "vortex-queries",
+              "subcommand": "vortex",
+              "name": "Vortex queries",
+              "data_formats": ["parquet", "vortex"],
+              "pr_targets": [
+                {"engine": "datafusion", "format": "parquet"},
+                {"engine": "datafusion", "format": "vortex"},
+                {"engine": "duckdb", "format": "parquet"},
+                {"engine": "duckdb", "format": "vortex"}
+              ],
+              "develop_targets": [
+                {"engine": "datafusion", "format": "parquet"},
+                {"engine": "datafusion", "format": "vortex"},
+                {"engine": "duckdb", "format": "parquet"},
+                {"engine": "duckdb", "format": "vortex"}
+              ],
+              "iterations": "100"
             }
           ]
       base_benchmark_matrix:
@@ -499,6 +518,25 @@ on:
                 {"engine": "datafusion", "format": "vortex"}
               ],
               "scale_factor": "1"
+            },
+            {
+              "id": "vortex-queries",
+              "subcommand": "vortex",
+              "name": "Vortex queries",
+              "data_formats": ["parquet", "vortex"],
+              "pr_targets": [
+                {"engine": "datafusion", "format": "parquet"},
+                {"engine": "datafusion", "format": "vortex"},
+                {"engine": "duckdb", "format": "parquet"},
+                {"engine": "duckdb", "format": "vortex"}
+              ],
+              "develop_targets": [
+                {"engine": "datafusion", "format": "parquet"},
+                {"engine": "datafusion", "format": "vortex"},
+                {"engine": "duckdb", "format": "parquet"},
+                {"engine": "duckdb", "format": "vortex"}
+              ],
+              "iterations": "100"
             }
           ]
       benchmark_profile:

bench-orchestrator/bench_orchestrator/config.py

@@ -54,6 +54,7 @@ class Benchmark(Enum):
     PUBLIC_BI = "public-bi"
     STATPOPGEN = "statpopgen"
     SPATIALBENCH = "spatialbench"
+    VORTEX_QUERIES = "vortex"
 
 
 # Engine to supported formats mapping.

vortex-bench/sql/vortex/0_sum-with-filter.sql

@@ -0,0 +1,4 @@
+-- As this aggregation has a filter, Vortex has to use a linear scan. Once stats
+-- are propagated to arrays, this should use zone maps to aggregate instead of
+-- decoding and processing each row.
+SELECT sum(col) FROM test WHERE col2 > 0 AND col2 < 1000;

vortex-bench/sql/vortex/1_sum.sql

@@ -0,0 +1,3 @@
+-- When Footer changes land, vortex-duckdb should populate statistics from
+-- Footer without loading and decoding the data.
+SELECT sum(col) FROM test;

vortex-bench/sql/vortex/init.sql

@@ -0,0 +1,8 @@
+-- Script that prepares Parquet data for our SQL microbenchmarks.
+
+COPY (
+    SELECT
+        i % 1000 AS col,
+        (i * 2654435761) % 100000 AS col2
+    FROM range(25000000) t(i)
+) TO 'test.parquet' (FORMAT parquet);

vortex-bench/src/datasets/mod.rs

@@ -81,6 +81,8 @@ pub enum BenchmarkDataset {
     Fineweb,
     #[serde(rename = "gharchive")]
     GhArchive,
+    #[serde(rename = "vortex")]
+    VortexQueries,
 }
 
 impl BenchmarkDataset {
@@ -97,6 +99,7 @@ impl BenchmarkDataset {
             BenchmarkDataset::PolarSignals { .. } => "polarsignals",
             BenchmarkDataset::Fineweb => "fineweb",
             BenchmarkDataset::GhArchive => "gharchive",
+            BenchmarkDataset::VortexQueries => "vortex",
         }
     }
 }
@@ -122,6 +125,7 @@ impl Display for BenchmarkDataset {
             }
             BenchmarkDataset::Fineweb => write!(f, "fineweb"),
             BenchmarkDataset::GhArchive => write!(f, "gharchive"),
+            BenchmarkDataset::VortexQueries => write!(f, "vortex"),
         }
     }
 }
@@ -179,6 +183,8 @@ impl BenchmarkDataset {
             BenchmarkDataset::PolarSignals { .. } => &["stacktraces"],
             BenchmarkDataset::Fineweb => &["fineweb"],
             BenchmarkDataset::GhArchive => &["events"],
+            // See VortexBenchmark::table_specs
+            BenchmarkDataset::VortexQueries => &[],
         }
     }
 }

vortex-bench/src/lib.rs

@@ -36,6 +36,7 @@ use vortex::file::WriteStrategyBuilder;
 use vortex::utils::aliases::hash_map::HashMap;
 
 use crate::spatialbench::SpatialBenchBenchmark;
+use crate::vortex_queries::VortexBenchmark;
 
 pub mod appian;
 pub mod benchmark;
@@ -61,6 +62,7 @@ pub mod tpch;
 pub mod utils;
 pub mod v3;
 pub mod vector_dataset;
+pub mod vortex_queries;
 
 pub use benchmark::Benchmark;
 pub use benchmark::TableSpec;
@@ -281,6 +283,8 @@ pub enum BenchmarkArg {
     PublicBi,
     #[clap(name = "spatialbench")]
     SpatialBench,
+    #[clap(name = "vortex")]
+    VortexQueries,
 }
 
 /// Default scale factor for TPC-related benchmarks
@@ -353,6 +357,13 @@ pub fn create_benchmark(b: BenchmarkArg, opts: &Opts) -> anyhow::Result<Box<dyn
             let benchmark = SpatialBenchBenchmark::new(scale_factor.to_string(), remote_data_dir)?;
             Ok(Box::new(benchmark) as _)
         }
+        BenchmarkArg::VortexQueries => {
+            let mut benchmark = VortexBenchmark::new()?;
+            if let Some(query) = opts.get("query") {
+                benchmark = benchmark.with_query(query)?;
+            }
+            Ok(Box::new(benchmark) as _)
+        }
     }
 }
 

vortex-bench/src/utils/file.rs

@@ -56,8 +56,12 @@ pub trait IdempotentPath {
     fn to_data_path(&self) -> PathBuf;
 }
 
+pub fn bench_dir() -> PathBuf {
+    workspace_root().join("vortex-bench")
+}
+
 pub fn data_dir() -> PathBuf {
-    workspace_root().join("vortex-bench").join("data")
+    bench_dir().join("data")
 }
 
 /// Find the workspace's root by looking for Cargo's lock file

vortex-bench/src/v3.rs

@@ -296,6 +296,7 @@ fn canonical_tpc_scale_factor(scale_factor: &str) -> String {
 /// | `Appian`                    | `appian`       | `None`              | `None`                                              | Static dataset; no scale factor. |
 /// | `PublicBi { name }`         | `public-bi`    | dataset name (e.g. `cms-provider`) | `None`               | Sub-dataset name lives in `dataset_variant`. |
 /// | `SpatialBench { scale_factor }` | `spatialbench` | `None`         | SF as string | Same canonicalization as TPC-H; no historical v2 records to merge with. |
+/// | `VortexQueries` | `vortex` | `None` | `None` | Own microbenchmarks |
 pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option<String>, Option<String>) {
     match d {
         BenchmarkDataset::TpcH { scale_factor } => (
@@ -331,6 +332,7 @@ pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option<String>,
         BenchmarkDataset::Fineweb => ("fineweb".to_string(), None, None),
         BenchmarkDataset::GhArchive => ("gharchive".to_string(), None, None),
         BenchmarkDataset::Appian => ("appian".to_string(), None, None),
+        BenchmarkDataset::VortexQueries => ("vortex".to_string(), None, None),
     }
 }
 

vortex-bench/src/vortex_queries.rs

@@ -0,0 +1,175 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Benchmark that runs queries which aren't part of any existing benchmark
+//! suite but which performance we want to track.
+
+use std::fs;
+use std::path::PathBuf;
+use std::process::Command;
+
+use anyhow::Context;
+use anyhow::Result;
+use anyhow::anyhow;
+use anyhow::bail;
+use glob::Pattern;
+use tracing::debug;
+use url::Url;
+use vortex::error::VortexExpect;
+
+use crate::Benchmark;
+use crate::BenchmarkDataset;
+use crate::Format;
+use crate::TableSpec;
+use crate::bench_dir;
+
+// Path to script that creates Parquet data
+const INIT_SQL: &str = "init.sql";
+
+pub struct VortexBenchmark {
+    data_url: Url,
+    queries_dir: PathBuf,
+    query: Option<PathBuf>,
+}
+
+impl VortexBenchmark {
+    pub fn new() -> Result<Self> {
+        let queries_dir = bench_dir().join("sql").join("vortex");
+        let data_url = Url::from_directory_path(&queries_dir)
+            .map_err(|_| anyhow!("cannot build URL for {}", queries_dir.display()))?;
+        Ok(Self {
+            data_url,
+            queries_dir,
+            query: None,
+        })
+    }
+
+    // Same as new(), but run only for "query" SQL file
+    pub fn with_query(mut self, query: &str) -> Result<Self> {
+        let as_path = PathBuf::from(query);
+        let path = if as_path.is_file() {
+            as_path
+        } else if query.ends_with(".sql") {
+            self.queries_dir.join(query)
+        } else {
+            self.queries_dir.join(format!("{query}.sql"))
+        };
+        if !path.is_file() {
+            bail!("{query} file not found in {}", self.queries_dir.display());
+        }
+        self.query = Some(path);
+        Ok(self)
+    }
+
+    fn query_files(&self) -> Result<Vec<PathBuf>> {
+        if let Some(query) = &self.query {
+            return Ok(vec![query.clone()]);
+        }
+
+        let entries = fs::read_dir(&self.queries_dir)
+            .with_context(|| format!("cannot list queries in {}", self.queries_dir.display()))?
+            .collect::<std::io::Result<Vec<_>>>()?;
+
+        let mut files: Vec<PathBuf> = entries
+            .into_iter()
+            .map(|entry| entry.path())
+            .filter(|path| {
+                path.extension().is_some_and(|ext| ext == "sql")
+                    && path.file_name().is_some_and(|name| name != INIT_SQL)
+            })
+            .collect();
+        files.sort();
+
+        if files.is_empty() {
+            bail!("no query files found in {}", self.queries_dir.display());
+        }
+        Ok(files)
+    }
+}
+
+#[async_trait::async_trait]
+impl Benchmark for VortexBenchmark {
+    fn queries(&self) -> Result<Vec<(usize, String)>> {
+        self.query_files()?
+            .iter()
+            .map(|path| {
+                let idx = path
+                    .file_name()
+                    .vortex_expect("no file name")
+                    .to_str()
+                    .vortex_expect("not utf-8")
+                    .split_once("_")
+                    .vortex_expect("query without a number")
+                    .0
+                    .parse::<usize>()?;
+                let query = fs::read_to_string(path)
+                    .with_context(|| format!("cannot read query {}", path.display()))?;
+                debug!(idx, file = %path.display(), "Loaded vortex query");
+                Ok((idx, query))
+            })
+            .collect()
+    }
+
+    async fn generate_base_data(&self) -> Result<()> {
+        let parquet_dir = self.queries_dir.join(Format::Parquet.name());
+        fs::create_dir_all(&parquet_dir)?;
+        let parquet_file = parquet_dir.join("test.parquet");
+
+        if parquet_file.exists() {
+            debug!("Parquet data present in {}", parquet_dir.display());
+            return Ok(());
+        }
+
+        let init_path = self.queries_dir.join(INIT_SQL);
+        let script = fs::read_to_string(&init_path)
+            .with_context(|| format!("cannot read {}", init_path.display()))?;
+
+        let output = Command::new("duckdb")
+            .current_dir(&parquet_dir)
+            .arg("-c")
+            .arg(&script)
+            .output()
+            .context("cannot run duckdb")?;
+        if !output.status.success() {
+            bail!(
+                "duckdb {INIT_SQL} failed: stdout={:?} stderr={:?}",
+                String::from_utf8_lossy(&output.stdout),
+                String::from_utf8_lossy(&output.stderr),
+            );
+        }
+
+        if !parquet_file.exists() {
+            bail!("{INIT_SQL} did not create Parquet files");
+        }
+
+        debug!("Parquet data generated in {}", parquet_dir.display());
+        Ok(())
+    }
+
+    fn dataset(&self) -> BenchmarkDataset {
+        BenchmarkDataset::VortexQueries
+    }
+
+    fn dataset_name(&self) -> &str {
+        "vortex"
+    }
+
+    fn dataset_display(&self) -> String {
+        "vortex".to_owned()
+    }
+
+    fn data_url(&self) -> &Url {
+        &self.data_url
+    }
+
+    fn table_specs(&self) -> Vec<TableSpec> {
+        vec![TableSpec::new("test", None)]
+    }
+
+    fn pattern(&self, table_name: &str, format: Format) -> Option<Pattern> {
+        Some(
+            Pattern::new(&format!("{table_name}.{}", format.ext()))
+                .expect("table name is a valid identifier"),
+        )
+    }
+}