benchmarks/src/sql_benchmark.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use arrow::array::{Array, RecordBatch};
use arrow::datatypes::*;
use arrow::error::ArrowError;
use arrow::util::display::{ArrayFormatter, FormatOptions};
use datafusion::dataframe::DataFrameWriteOptions;
use datafusion::datasource::MemTable;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::{ExecutionPlan, collect, execute_stream};
use datafusion::prelude::{CsvReadOptions, DataFrame, SessionContext};
use datafusion_common::config::CsvOptions;
use datafusion_common::{DataFusionError, Result, exec_datafusion_err};
use futures::StreamExt;
use log::{debug, info, trace, warn};
use regex::Regex;
use std::collections::HashMap;
use std::fmt::Debug;
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock};

/// A collection of benchmark configurations and state used by the DataFusion
/// sql test harness.  Each benchmark is defined by a file that can contain
/// directives such as `load`, `run`, `assert`, `result`, etc.  The
/// `SqlBenchmark` struct holds the parsed data from that file and
/// the impl provides methods to run, assert, persist, verify and cleanup benchmark
/// results.
#[derive(Debug, Clone)]
pub struct SqlBenchmark {
    /// Human‑readable name of the benchmark.
    name: String,
    /// Top‑level group name (derived from the file path or defined in a benchmark).
    group: String,
    /// Subgroup name, often a logical grouping.
    subgroup: String,
    /// Full path to the benchmark file.
    benchmark_path: PathBuf,
    /// Mapping of placeholder keys to concrete values (e.g. `"BENCHMARK_DIR"`).
    replacement_mapping: HashMap<String, String>,
    /// Expected string that must appear in the physical plan of the queries.
    expect: Vec<String>,
    /// All SQL queries grouped by directive (`load`, `run`, etc.).
    queries: HashMap<QueryDirective, Vec<String>>,
    /// Queries whose results are persisted to disk for later comparison.
    result_queries: Vec<BenchmarkQuery>,
    /// Queries whose results are asserted against an expected table.
    assert_queries: Vec<BenchmarkQuery>,
    /// Flag indicating whether the benchmark has been fully loaded
    is_loaded: bool,
    /// Flag indicating whether the `expect` strings were checked against the
    /// physical plan. The plan does not change between iterations, so the
    /// check runs on the first iteration only and is not repeated in the
    /// measured region.
    plans_validated: bool,
    /// Stores the last run results if needed so they can be compared or persisted.
    last_results: Option<Vec<RecordBatch>>,
    /// Print plans and execution metrics for each run query.
    debug: bool,
    /// Whether diagnostics have been printed for a successful run.
    diagnostics_printed: bool,
    /// echo statements
    echo: Vec<String>,
}

impl SqlBenchmark {
    pub async fn new(
        ctx: &SessionContext,
        full_path: impl AsRef<Path>,
        benchmark_directory: impl AsRef<Path>,
    ) -> Result<Self> {
        Self::new_with_replacements(ctx, full_path, benchmark_directory, HashMap::new())
            .await
    }

    /// Creates a benchmark using caller-provided template replacements.
    ///
    /// Caller values take precedence over environment variables during
    /// `${...}` substitution; `BENCHMARK_DIR` is still set internally.
    pub async fn new_with_replacements(
        ctx: &SessionContext,
        full_path: impl AsRef<Path>,
        benchmark_directory: impl AsRef<Path>,
        replacement_mapping: HashMap<String, String>,
    ) -> Result<Self> {
        let full_path = full_path.as_ref();
        let benchmark_directory = benchmark_directory.as_ref();
        let group_name = parse_group_from_path(full_path, benchmark_directory);
        let mut bm = Self {
            name: String::new(),
            group: group_name,
            subgroup: String::new(),
            benchmark_path: full_path.to_path_buf(),
            replacement_mapping,
            expect: vec![],
            plans_validated: false,
            queries: HashMap::new(),
            result_queries: vec![],
            assert_queries: vec![],
            is_loaded: false,
            last_results: None,
            debug: false,
            diagnostics_printed: false,
            echo: vec![],
        };
        insert_replacement(
            &mut bm.replacement_mapping,
            "BENCHMARK_DIR",
            benchmark_directory.to_string_lossy().into_owned(),
        );

        let path = bm.benchmark_path.clone();
        bm.process_file(ctx, &path).await?;

        Ok(bm)
    }

    /// Enable plan diagnostics for run queries, including post-execution metrics.
    /// Output occurs on the first successful run, including preparation or warm-up,
    /// and affects that execution's timing.
    pub fn set_debug(&mut self, debug: bool) {
        self.debug = debug;
    }

    /// Initializes the benchmark by executing `load` and `init` queries.
    ///
    /// Registers any required tables or sets up state in the provided
    /// `SessionContext` before running queries.  This method is idempotent:
    /// calling it multiple times on the same instance returns
    /// immediately after the first successful initialization.
    ///
    /// # Errors
    /// Returns an error if any `load` or `init` query fails, or if the
    /// benchmark file does not contain a `run` query.
    pub async fn initialize(&mut self, ctx: &SessionContext) -> Result<()> {
        if self.is_loaded {
            return Ok(());
        }

        let path = self.benchmark_path.to_string_lossy().into_owned();

        // validate there was a run query
        if !self.queries.contains_key(&QueryDirective::Run) {
            return Err(exec_datafusion_err!(
                "Invalid benchmark file: no \"run\" query specified: {path}"
            ));
        }

        // display any echo's
        self.echo.iter().for_each(|txt| println!("{txt}"));

        let load_queries = self.queries.get(&QueryDirective::Load);

        if let Some(queries) = load_queries {
            for query in queries {
                debug!("Executing load query {query}");
                ctx.sql(query).await?.collect().await?;
            }
        }

        let init_queries = self.queries.get(&QueryDirective::Init);

        if let Some(queries) = init_queries {
            for query in queries {
                debug!("Executing init query {query}");
                ctx.sql(query).await?.collect().await?;
            }
        }

        self.is_loaded = true;

        Ok(())
    }

    /// Executes the `assert` queries and compares actual results against
    /// expected values.
    ///
    /// Each `assert` query must be followed by a result table (separated by
    /// `----`) in the benchmark file.  The assertion passes only if the
    /// returned record batches exactly match the expected rows.
    ///
    /// # Errors
    /// Returns an error if any `assert` query fails, or if the actual and
    /// expected results differ in row count or cell values.
    pub async fn assert(&mut self, ctx: &SessionContext) -> Result<()> {
        info!("Running assertions...");

        for assert_query in &self.assert_queries {
            let query = &assert_query.query;

            info!("Executing assert query {query}");

            let result = ctx.sql(query).await?.collect().await?;
            let formatted_actual_results = format_record_batches(&result)?;

            Self::compare_results(
                assert_query,
                &formatted_actual_results,
                &assert_query.expected_result,
            )?;
        }

        Ok(())
    }

    /// Executes the `run` queries, optionally saving results for later
    /// verification. If there are multiple queries only the results for
    /// the last query are saved.
    ///
    /// When `save_results` is `true`, it collects `SELECT`/`WITH` query
    /// results and stores them in `last_results`.
    ///
    /// When `save_results` is `false`, it streams results and counts rows
    /// without buffering them.
    ///
    /// If an 'expect' string is defined this method also validates that
    /// the physical plan contains that string.
    ///
    /// # Errors
    /// Returns an error if a `run` query fails or if expected plan strings
    /// are not found.
    pub async fn run(
        &mut self,
        ctx: &SessionContext,
        save_results: bool,
    ) -> Result<usize> {
        let run_queries = self
            .queries
            .get(&QueryDirective::Run)
            .ok_or_else(|| exec_datafusion_err!("Run query should be loaded by now"))?;

        let mut result_count = 0;

        let result: Vec<RecordBatch> = {
            let mut local_result = vec![];

            for query in run_queries {
                match save_results {
                    true => {
                        debug!(
                            "Running query (saving results) {}-{}: {query}",
                            self.group, self.subgroup
                        );

                        let df = ctx.sql(query).await?;
                        let result_schema = Arc::new(df.schema().as_arrow().clone());
                        let physical_plan = self.create_physical_plan(df).await?;

                        if !self.plans_validated {
                            self.validate_expected_plan(&physical_plan)?;
                        }

                        let mut batches =
                            collect(Arc::clone(&physical_plan), ctx.task_ctx()).await?;

                        self.print_plan_metrics(&physical_plan);

                        // save the output for select/with queries
                        if is_result_statement(query) {
                            if batches.is_empty() {
                                batches.push(RecordBatch::new_empty(result_schema));
                            }
                            let row_count_for_query =
                                batches.iter().map(RecordBatch::num_rows).sum::<usize>();
                            debug!(
                                "Persisting {} batches ({} rows)...",
                                batches.len(),
                                row_count_for_query
                            );

                            result_count = row_count_for_query;
                            local_result = batches;
                        }
                    }
                    false => {
                        debug!(
                            "Running query (ignoring results) {}-{}: {query}",
                            self.group, self.subgroup
                        );

                        let row_count = self
                            .execute_sql_without_result_buffering(
                                query,
                                ctx,
                                !self.plans_validated,
                            )
                            .await?;

                        if is_result_statement(query) {
                            result_count = row_count;
                        }
                    }
                }
            }

            Ok::<Vec<RecordBatch>, DataFusionError>(local_result)
        }?;

        debug!("Results have {result_count} rows");

        // Every `run` query was planned and checked above, and a plan does not
        // change between iterations, so later iterations skip the check.
        self.plans_validated = true;
        self.diagnostics_printed |= self.debug;

        // Store results for verification
        self.last_results = Some(result);

        Ok(result_count)
    }

    /// Calls run and persists results to disk as a CSV file.
    ///
    /// Requires that the benchmark defines a `result` or `result_query`.
    /// Registers the results in a memory table and writes them to disk with
    /// pipe delimiters and a header row.
    ///
    /// # Errors
    /// Returns an error if no results are available or if writing to the
    /// target path fails.
    pub async fn persist(&mut self, ctx: &SessionContext) -> Result<()> {
        let _ = self.run(ctx, true).await?;

        // Check if we have result queries to persist for
        if self.result_queries.is_empty() {
            info!("No result paths to persist");
            return Ok(());
        }

        let results = self
            .last_results
            .as_ref()
            .expect("run should store last_results after successful execution");

        let query = &self.result_queries[0];
        let path = query.path.as_ref().ok_or_else(|| {
            exec_datafusion_err!(
                "Unable to persist results from query '{}', no result specified",
                query.query
            )
        })?;

        info!("Persisting results for query to {path}");

        let first_batch = results
            .first()
            .ok_or_else(|| exec_datafusion_err!("Results should be loaded"))?;

        let schema = first_batch.schema();
        let provider = MemTable::try_new(schema, vec![results.clone()])?;

        ctx.register_table("persist_data", Arc::new(provider))?;

        let df = ctx.table("persist_data").await?;
        df.write_csv(
            path,
            DataFrameWriteOptions::new(),
            Some(
                CsvOptions::default()
                    .with_delimiter(b'|')
                    .with_has_header(true),
            ),
        )
        .await?;

        ctx.deregister_table("persist_data")?;

        Ok(())
    }

    /// Verifies persisted results against expected values.
    ///
    /// Executes the `result_query` or uses the stored last run results, then
    /// compares actual output rows to the expected values defined in the
    /// benchmark file.
    ///
    /// # Errors
    /// Returns an error if no results are available or if the actual and
    /// expected results differ in count or content.
    pub async fn verify(&mut self, ctx: &SessionContext) -> Result<()> {
        // Check if we have result queries to verify
        if self.result_queries.is_empty() {
            return Ok(());
        }

        if self.last_results.is_none() {
            return Err(exec_datafusion_err!(
                "No results available for verification. Run the benchmark first."
            ));
        }

        info!("Verifying results...");

        self.load_expected_result_files(ctx).await?;

        // Get the first result query (assuming only one for now)
        let query = &self.result_queries[0];
        let formatted_actual_results = if !query.query.trim().is_empty() {
            let results = ctx.sql(&query.query).await?.collect().await?;
            format_record_batches(&results)
        } else {
            let actual_results = self
                .last_results
                .as_ref()
                .expect("last_results should be present after successful run");
            format_record_batches(actual_results)
        }?;

        Self::compare_results(query, &formatted_actual_results, &query.expected_result)
    }

    /// Runs `cleanup` queries to reset state after the benchmark run.
    pub async fn cleanup(&mut self, ctx: &SessionContext) -> Result<()> {
        info!("Running cleanup...");

        let cleanup_queries = self.queries.get(&QueryDirective::Cleanup);

        if let Some(queries) = cleanup_queries {
            for query in queries {
                let _ = ctx.sql(query).await?.collect().await?;
            }
        }

        Ok(())
    }

    async fn load_expected_result_files(&mut self, ctx: &SessionContext) -> Result<()> {
        for query in &mut self.result_queries {
            if query.query.trim().is_empty() {
                let Some(path) = query.path.clone() else {
                    continue;
                };

                let loaded_query =
                    read_query_from_file(ctx, path, &HashMap::new()).await?;
                query.column_count = loaded_query.column_count;
                query.expected_result = loaded_query.expected_result;
            }
        }

        Ok(())
    }

    fn compare_results(
        query: &BenchmarkQuery,
        actual_results: &[Vec<String>],
        expected_results: &[Vec<String>],
    ) -> Result<()> {
        if actual_results.is_empty() && expected_results.is_empty() {
            return Ok(());
        }

        // Compare row count
        if actual_results.len() != expected_results.len() {
            return Err(exec_datafusion_err!(
                "Error in result: expected {} rows but got {} for query {}",
                expected_results.len(),
                actual_results.len(),
                query.query
            ));
        }

        // Compare values
        let zipped = actual_results
            .iter()
            .enumerate()
            .zip(expected_results.iter());

        for ((row_idx, actual), expected) in zipped {
            trace!(
                "row {}\nactual: {actual:?}\nexpected: {expected:?}",
                row_idx + 1
            );

            // Compare column count
            if actual.len() != expected.len() {
                return Err(exec_datafusion_err!(
                    "Error in result: expected {} columns but got {} for query {}",
                    expected.len(),
                    actual.len(),
                    query.query
                ));
            }

            for (col_idx, expected_val) in
                expected.iter().enumerate().take(query.column_count)
            {
                // The row-width check above guarantees this index exists.
                let actual_val = &actual[col_idx];

                trace!("actual_val = {actual_val:?}\nexpected_val = {expected_val:?}");

                if (expected_val == "NULL" && actual_val.is_empty())
                    || (expected_val == actual_val)
                    || (expected_val == "(empty)"
                        && (actual_val.is_empty() || actual_val == "NULL"))
                {
                    continue;
                }

                return Err(exec_datafusion_err!(
                    "Error in result on row {}, column {} running query \"{}\": expected value \
                    \"{expected_val}\" but got value \"{actual_val}\" in row: {actual:?}",
                    row_idx + 1,
                    col_idx + 1,
                    query.query
                ));
            }
        }

        Ok(())
    }

    async fn process_file(&mut self, ctx: &SessionContext, path: &Path) -> Result<()> {
        debug!("Processing file {}", path.display());

        let mut replacement_mapping = self.replacement_mapping.clone();
        insert_replacement(
            &mut replacement_mapping,
            "FILE_PATH",
            path.to_string_lossy().into_owned(),
        );

        let mut reader = BenchmarkFileReader::new(path, replacement_mapping)?;
        let mut line = String::with_capacity(1024);
        let mut reader_result = reader.read_line(&mut line);

        while let Some(result) = reader_result {
            match result {
                Ok(_) => {
                    if !is_blank_or_comment_line(&line) {
                        // boxing required because of recursion
                        Box::pin(self.process_line(ctx, &mut reader, &mut line)).await?;
                    }
                }
                Err(e) => return Err(e),
            }

            // Clear the line buffer for the next iteration.
            line.clear();
            reader_result = reader.read_line(&mut line);
        }

        Ok(())
    }

    async fn process_line(
        &mut self,
        ctx: &SessionContext,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
    ) -> Result<()> {
        // Split the line into directive and arguments.
        let cloned_line = line.trim_start().to_string();
        let splits: Vec<&str> = cloned_line.split_whitespace().collect();

        BenchmarkDirective::select(reader, splits[0])?
            .process(ctx, self, reader, line, &splits)
            .await
    }

    fn process_query(&mut self, splits: &[&str], mut query: String) -> Result<()> {
        debug!("Processing query {query}");

        // Trim and validate.
        query = query.trim().to_string();
        if query.is_empty() {
            return Ok(());
        }

        // remove comments
        query = query
            .lines()
            .filter(|line| !is_comment_line(line))
            .collect::<Vec<_>>()
            .join("\n");

        if query.trim().is_empty() {
            return Ok(());
        }

        query = process_replacements(&query, self.replacement_mapping())?;

        let directive = QueryDirective::parse(splits[0]).ok_or_else(|| {
            exec_datafusion_err!("Invalid query directive: {}", splits[0])
        })?;

        self.queries.entry(directive).or_default().push(query);

        Ok(())
    }

    /// Checks the `expect` strings against the plan as `EXPLAIN` displays it.
    ///
    /// `{:#?}` would dump every `RecordBatch` an in-memory source holds, which
    /// is megabytes for a benchmark that builds its tables with `CREATE TABLE
    /// ... AS SELECT`.
    fn validate_expected_plan(
        &self,
        physical_plan: &Arc<dyn ExecutionPlan>,
    ) -> Result<()> {
        if self.expect.is_empty() {
            return Ok(());
        }

        let plan_string = DisplayableExecutionPlan::new(physical_plan.as_ref())
            .indent(true)
            .to_string();

        for exp_str in &self.expect {
            if !plan_string.contains(exp_str) {
                return Err(exec_datafusion_err!(
                    "The query physical plan does not contain the expected string '{exp_str}'. Physical plan: {plan_string}"
                ));
            }
        }

        Ok(())
    }

    async fn create_physical_plan(
        &self,
        df: DataFrame,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if !self.debug || self.diagnostics_printed {
            return df.create_physical_plan().await;
        }

        let (state, plan) = df.into_parts();
        println!("=== Logical plan ===\n{plan}\n");
        let plan = state.optimize(&plan)?;
        println!("=== Optimized logical plan ===\n{plan}\n");
        let physical_plan = state.create_physical_plan(&plan).await?;
        println!(
            "=== Physical plan ===\n{}\n",
            DisplayableExecutionPlan::new(physical_plan.as_ref()).indent(true)
        );
        Ok(physical_plan)
    }

    fn print_plan_metrics(&self, physical_plan: &Arc<dyn ExecutionPlan>) {
        if self.debug && !self.diagnostics_printed {
            println!(
                "=== Physical plan with metrics ===\n{}\n",
                DisplayableExecutionPlan::with_metrics(physical_plan.as_ref())
                    .indent(true)
            );
        }
    }

    async fn execute_sql_without_result_buffering(
        &self,
        sql: &str,
        ctx: &SessionContext,
        validate_plan: bool,
    ) -> Result<usize> {
        let mut row_count = 0;

        let df = ctx.sql(sql).await?;
        let physical_plan = self.create_physical_plan(df).await?;

        if validate_plan {
            self.validate_expected_plan(&physical_plan)?;
        }
        let mut stream = execute_stream(Arc::clone(&physical_plan), ctx.task_ctx())?;

        while let Some(batch) = stream.next().await {
            row_count += batch?.num_rows();

            // Evaluate the result and do nothing, the result will be dropped
            // to reduce memory pressure
        }

        self.print_plan_metrics(&physical_plan);
        Ok(row_count)
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn group(&self) -> &str {
        &self.group
    }

    pub fn subgroup(&self) -> &str {
        &self.subgroup
    }

    pub fn benchmark_path(&self) -> &Path {
        &self.benchmark_path
    }

    pub fn replacement_mapping(&self) -> &HashMap<String, String> {
        &self.replacement_mapping
    }

    pub fn queries(&self) -> &HashMap<QueryDirective, Vec<String>> {
        &self.queries
    }

    pub fn result_queries(&self) -> &[BenchmarkQuery] {
        &self.result_queries
    }

    pub fn assert_queries(&self) -> &[BenchmarkQuery] {
        &self.assert_queries
    }

    pub fn is_loaded(&self) -> bool {
        self.is_loaded
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QueryDirective {
    Load,
    Run,
    Init,
    Cleanup,
}

impl QueryDirective {
    fn parse(value: &str) -> Option<Self> {
        if value.eq_ignore_ascii_case("load") {
            Some(Self::Load)
        } else if value.eq_ignore_ascii_case("init") {
            Some(Self::Init)
        } else if value.eq_ignore_ascii_case("run") {
            Some(Self::Run)
        } else if value.eq_ignore_ascii_case("cleanup") {
            Some(Self::Cleanup)
        } else {
            None
        }
    }

    fn as_str(self) -> &'static str {
        match self {
            Self::Load => "load",
            Self::Run => "run",
            Self::Init => "init",
            Self::Cleanup => "cleanup",
        }
    }
}

enum BenchmarkDirective {
    Load,
    Run,
    Init,
    Cleanup,
    Name,
    Group,
    Subgroup,
    Expect,
    Assert,
    ResultQuery,
    Results,
    Template,
    Include,
    Echo,
}

impl BenchmarkDirective {
    fn select(
        reader: &BenchmarkFileReader,
        directive: &str,
    ) -> Result<BenchmarkDirective> {
        if directive.eq_ignore_ascii_case("load") {
            Ok(BenchmarkDirective::Load)
        } else if directive.eq_ignore_ascii_case("run") {
            Ok(BenchmarkDirective::Run)
        } else if directive.eq_ignore_ascii_case("init") {
            Ok(BenchmarkDirective::Init)
        } else if directive.eq_ignore_ascii_case("cleanup") {
            Ok(BenchmarkDirective::Cleanup)
        } else if directive.eq_ignore_ascii_case("name") {
            Ok(BenchmarkDirective::Name)
        } else if directive.eq_ignore_ascii_case("group") {
            Ok(BenchmarkDirective::Group)
        } else if directive.eq_ignore_ascii_case("subgroup") {
            Ok(BenchmarkDirective::Subgroup)
        } else if directive.eq_ignore_ascii_case("expect_plan") {
            Ok(BenchmarkDirective::Expect)
        } else if directive.eq_ignore_ascii_case("assert") {
            Ok(BenchmarkDirective::Assert)
        } else if directive.eq_ignore_ascii_case("result_query") {
            Ok(BenchmarkDirective::ResultQuery)
        } else if directive.eq_ignore_ascii_case("result") {
            Ok(BenchmarkDirective::Results)
        } else if directive.eq_ignore_ascii_case("template") {
            Ok(BenchmarkDirective::Template)
        } else if directive.eq_ignore_ascii_case("include") {
            Ok(BenchmarkDirective::Include)
        } else if directive.eq_ignore_ascii_case("echo") {
            Ok(BenchmarkDirective::Echo)
        } else {
            Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(&format!("Unrecognized command: {directive}"))
            ))
        }
    }

    async fn process(
        &self,
        ctx: &SessionContext,
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
        splits: &[&str],
    ) -> Result<()> {
        trace!("-- handling {}", splits[0]);

        match self {
            BenchmarkDirective::Load
            | BenchmarkDirective::Run
            | BenchmarkDirective::Init
            | BenchmarkDirective::Cleanup => {
                Self::process_query_directive(bench, reader, line, splits)
            }
            BenchmarkDirective::Name => Self::process_metadata_value(
                bench,
                reader,
                line,
                "name",
                "BENCH_NAME",
                "name must be followed by a value",
            ),
            BenchmarkDirective::Group => Self::process_metadata_value(
                bench,
                reader,
                line,
                "group",
                "BENCH_GROUP",
                "group must be followed by a value",
            ),
            BenchmarkDirective::Subgroup => Self::process_metadata_value(
                bench,
                reader,
                line,
                "subgroup",
                "BENCH_SUBGROUP",
                "subgroup must be followed by a value",
            ),
            BenchmarkDirective::Expect => Self::process_expect(bench, reader, splits),
            BenchmarkDirective::Assert => {
                Self::process_assert(bench, reader, line, splits)
            }
            BenchmarkDirective::ResultQuery => {
                Self::process_result_query(bench, reader, line, splits)
            }
            BenchmarkDirective::Results => Self::process_results(bench, reader, splits),
            BenchmarkDirective::Template => {
                Self::process_template(ctx, bench, reader, line, splits).await
            }
            BenchmarkDirective::Include => {
                Self::process_include(ctx, bench, reader, splits).await
            }
            BenchmarkDirective::Echo => Self::process_echo(bench, reader, splits),
        }
    }

    fn process_query_directive(
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
        splits: &[&str],
    ) -> Result<()> {
        let directive = QueryDirective::parse(splits[0]).ok_or_else(|| {
            exec_datafusion_err!("Invalid query directive: {}", splits[0])
        })?;

        if directive == QueryDirective::Run && bench.queries.contains_key(&directive) {
            return Err(exec_datafusion_err!(
                "Multiple calls to run in the same benchmark file"
            ));
        }

        line.clear();

        // Read the query body until a blank line or EOF.
        let mut query = String::new();
        let mut reader_result = reader.read_line(line);

        loop {
            match reader_result {
                Some(Ok(_)) => {
                    if is_comment_line(line) {
                        // comment, ignore
                    } else if is_blank_line(line) {
                        break;
                    } else {
                        query.push_str(line);
                        query.push('\n');
                    }
                }
                Some(Err(e)) => return Err(e),
                None => break,
            }

            // Clear the line buffer for the next iteration.
            line.clear();
            reader_result = reader.read_line(line);
        }

        // Optional file parameter.
        if splits.len() > 1 && !splits[1].is_empty() {
            if !query.trim().is_empty() {
                return Err(exec_datafusion_err!(
                    "{}",
                    reader.format_exception(&format!(
                        "{} directive must use either a query file or inline SQL, not both",
                        directive.as_str()
                    ))
                ));
            }

            let query_path = resolve_benchmark_file_path(splits[1]);
            debug!("Processing {} file: {}", splits[0], query_path.display());

            let query_file = fs::read_to_string(&query_path).map_err(|e| {
                exec_datafusion_err!(
                    "Failed to read query file {}: {e}",
                    query_path.display()
                )
            })?;
            let query_file = query_file.replace("\r\n", "\n");

            // some files have multiple queries, split apart
            for query in split_query_statements(&query_file) {
                bench.process_query(splits, query.to_string())?;
            }
        } else if directive == QueryDirective::Run {
            for query in split_query_statements(&query) {
                bench.process_query(splits, query.to_string())?;
            }
        } else {
            bench.process_query(splits, query)?;
        }

        Ok(())
    }

    fn process_metadata_value(
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &str,
        directive: &str,
        replacement_key: &str,
        message: &str,
    ) -> Result<()> {
        let value =
            directive_value(reader, line.trim_start(), directive, message)?.to_string();

        match directive {
            "name" => bench.name.clone_from(&value),
            "group" => bench.group.clone_from(&value),
            "subgroup" => bench.subgroup.clone_from(&value),
            _ => unreachable!("unsupported metadata directive: {directive}"),
        }

        insert_replacement(
            &mut bench.replacement_mapping,
            replacement_key,
            value.clone(),
        );
        insert_replacement(&mut reader.replacements, replacement_key, value);

        Ok(())
    }

    fn process_expect(
        bench: &mut SqlBenchmark,
        reader: &BenchmarkFileReader,
        splits: &[&str],
    ) -> Result<()> {
        trace!("-- handling {}", splits[0]);

        if splits.len() <= 1 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "expect_plan must be followed by a string to search in the physical plan"
                )
            ));
        }

        bench.expect.push(splits[1..].join(" ").to_string());

        Ok(())
    }

    fn process_assert(
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
        splits: &[&str],
    ) -> Result<()> {
        // count the amount of columns based on character count. The actual
        // character used is irrelevant.
        if splits.len() <= 1 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "assert must be followed by a column count (e.g. assert III)"
                )
            ));
        }

        line.clear();

        // read the actual query
        let mut found_break = false;
        let mut sql = String::new();
        let mut reader_result = reader.read_line(line);

        loop {
            match reader_result {
                Some(Ok(_)) => {
                    if line.trim() == "----" {
                        found_break = true;
                        break;
                    }
                    sql.push('\n');
                    sql.push_str(line);
                }
                Some(Err(e)) => return Err(e),
                None => break,
            }

            // Clear the line buffer for the next iteration.
            line.clear();
            reader_result = reader.read_line(line);
        }

        if !found_break {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "assert must be followed by a query and a result (separated by ----)"
                )
            ));
        }

        bench
            .assert_queries
            .push(read_query_from_reader(reader, &sql, splits[1])?);

        Ok(())
    }

    fn process_results(
        bench: &mut SqlBenchmark,
        reader: &BenchmarkFileReader,
        splits: &[&str],
    ) -> Result<()> {
        if splits.len() <= 1 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "result must be followed by a path to a result file"
                )
            ));
        }

        if !bench.result_queries.is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception("multiple results found")
            ));
        }

        let path = process_replacements(splits[1], &bench.replacement_mapping)?;

        bench.result_queries.push(BenchmarkQuery {
            path: Some(path),
            query: String::new(),
            column_count: 0,
            expected_result: vec![],
        });

        Ok(())
    }

    fn process_result_query(
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
        splits: &[&str],
    ) -> Result<()> {
        if splits.len() <= 1 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "result_query must be followed by a column count (e.g. result_query III)"
                )
            ));
        }

        line.clear();

        let mut sql = String::new();
        let mut found_break = false;
        let mut reader_result = reader.read_line(line);

        loop {
            match reader_result {
                Some(Ok(_)) => {
                    if line.trim() == "----" {
                        found_break = true;
                        break;
                    }
                    sql.push_str(line);
                    sql.push('\n');
                }
                Some(Err(e)) => return Err(e),
                None => break,
            }

            // Clear the line buffer for the next iteration.
            line.clear();
            reader_result = reader.read_line(line);
        }

        if !found_break {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception(
                    "result_query must be followed by a query and a result (separated by ----)"
                )
            ));
        }

        let result_check = read_query_from_reader(reader, &sql, splits[1])?;

        if !bench.result_queries.is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception("multiple results found")
            ));
        }
        bench.result_queries.push(result_check);

        Ok(())
    }

    async fn process_template(
        ctx: &SessionContext,
        bench: &mut SqlBenchmark,
        reader: &mut BenchmarkFileReader,
        line: &mut String,
        splits: &[&str],
    ) -> Result<()> {
        if splits.len() != 2 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception("template requires a single template path")
            ));
        }

        // template: update the path to read
        bench.benchmark_path = PathBuf::from(splits[1]);

        line.clear();

        // now read parameters
        let mut reader_result = reader.read_line(line);

        loop {
            match reader_result {
                Some(Ok(_)) => {
                    if is_comment_line(line) {
                        // Clear the line buffer for the next iteration.
                        line.clear();
                        reader_result = reader.read_line(line);
                        continue;
                    }
                    if is_blank_line(line) {
                        break;
                    }

                    let Some((key, value)) = line.trim_start().split_once('=') else {
                        return Err(exec_datafusion_err!(
                            "{}",
                            reader.format_exception(
                                "Expected a template parameter in the form of X=Y"
                            )
                        ));
                    };
                    insert_replacement(
                        &mut bench.replacement_mapping,
                        key.trim(),
                        value.trim().to_string(),
                    );
                }
                Some(Err(e)) => return Err(e),
                None => break,
            }

            // Clear the line buffer for the next iteration.
            line.clear();
            reader_result = reader.read_line(line);
        }

        // restart the load from the template file
        let path = resolve_benchmark_file_path(splits[1]);
        Box::pin(bench.process_file(ctx, &path)).await
    }

    async fn process_include(
        ctx: &SessionContext,
        bench: &mut SqlBenchmark,
        reader: &BenchmarkFileReader,
        splits: &[&str],
    ) -> Result<()> {
        if splits.len() != 2 || splits[1].is_empty() {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception("include requires a single argument")
            ));
        }

        let path = resolve_benchmark_file_path(splits[1]);
        Box::pin(bench.process_file(ctx, &path)).await
    }

    fn process_echo(
        bench: &mut SqlBenchmark,
        reader: &BenchmarkFileReader,
        splits: &[&str],
    ) -> Result<()> {
        if splits.len() < 2 {
            return Err(exec_datafusion_err!(
                "{}",
                reader.format_exception("Echo requires an argument")
            ));
        }

        bench.echo.push(splits[1..].join(" "));

        Ok(())
    }
}

struct BenchmarkFileReader {
    path: PathBuf,
    reader: BufReader<File>,
    line_nr: usize,
    replacements: HashMap<String, String>,
}

impl BenchmarkFileReader {
    fn new<P: Into<PathBuf>>(
        path: P,
        replacements: HashMap<String, String>,
    ) -> Result<Self> {
        let path = path.into();
        let file = OpenOptions::new().read(true).open(&path)?;

        Ok(Self {
            path,
            reader: BufReader::new(file),
            line_nr: 0,
            replacements,
        })
    }

    /// Read the next line, applying replacements and removing line terminators.
    fn read_line(&mut self, line: &mut String) -> Option<Result<()>> {
        match self.reader.read_line(line) {
            Ok(0) => None,
            Ok(_) => {
                self.line_nr += 1;

                // Trim newline and carriage return without changing other content.
                let trimmed_len = line.trim_end_matches(['\n', '\r']).len();
                line.truncate(trimmed_len);

                match process_replacements(line, &self.replacements) {
                    Ok(l) => {
                        *line = l;
                        Some(Ok(()))
                    }
                    Err(error) => Some(Err(error)),
                }
            }
            Err(e) => Some(Err(e.into())),
        }
    }

    fn format_exception(&self, msg: &str) -> String {
        format!("{}:{} - {}", self.path.display(), self.line_nr, msg)
    }
}

#[derive(Debug, Clone)]
pub struct BenchmarkQuery {
    path: Option<String>,
    query: String,
    column_count: usize,
    expected_result: Vec<Vec<String>>,
}

// ---- utility function below

fn directive_value<'a>(
    reader: &BenchmarkFileReader,
    line: &'a str,
    directive: &str,
    message: &str,
) -> Result<&'a str> {
    let value = line
        .get(..directive.len())
        .filter(|prefix| prefix.eq_ignore_ascii_case(directive))
        .and_then(|_| line.get(directive.len()..))
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| exec_datafusion_err!("{}", reader.format_exception(message)))?;

    Ok(value)
}

fn parse_group_from_path(path: &Path, benchmark_directory: &Path) -> String {
    let mut group_name = String::new();
    let mut parent = path.parent();

    while let Some(p) = parent {
        if path_ends_with_ignore_ascii_case(p, benchmark_directory) {
            break;
        }

        if let Some(dir_name) = p.file_name() {
            group_name = dir_name.to_string_lossy().into_owned();
        }

        parent = p.parent();
    }

    if group_name.is_empty() {
        warn!("Unable to find group name in path: {}", path.display());
    }

    group_name
}

fn path_ends_with_ignore_ascii_case(path: &Path, suffix: &Path) -> bool {
    let mut path_components = path.components().rev();

    for suffix_component in suffix.components().rev() {
        let Some(path_component) = path_components.next() else {
            return false;
        };

        if !path_component
            .as_os_str()
            .to_string_lossy()
            .eq_ignore_ascii_case(&suffix_component.as_os_str().to_string_lossy())
        {
            return false;
        }
    }

    true
}

fn starts_with_ignore_ascii_case(input: &str, prefix: &str) -> bool {
    input
        .get(..prefix.len())
        .is_some_and(|value| value.eq_ignore_ascii_case(prefix))
}

fn is_result_statement(statement: &str) -> bool {
    let statement = statement.trim_start();
    starts_with_ignore_ascii_case(statement, "select")
        || starts_with_ignore_ascii_case(statement, "with")
}

fn split_query_statements(sql: &str) -> impl Iterator<Item = &str> {
    sql.split("\n\n")
        .flat_map(|query| {
            query
                .split_inclusive(";\n")
                .map(|part| part.trim_end_matches('\n'))
        })
        .filter(|query| !query.trim().is_empty())
}

fn is_blank_line(line: &str) -> bool {
    line.trim().is_empty()
}

fn is_comment_line(line: &str) -> bool {
    let line = line.trim_start();
    line.starts_with('#') || line.starts_with("--")
}

fn is_blank_or_comment_line(line: &str) -> bool {
    is_blank_line(line) || is_comment_line(line)
}

fn insert_replacement(
    replacement_map: &mut HashMap<String, String>,
    key: &str,
    value: String,
) {
    replacement_map.insert(key.to_lowercase(), value);
}

fn replace_all<E>(
    re: &Regex,
    haystack: &str,
    replacement: impl Fn(&regex::Captures) -> Result<String, E>,
) -> Result<String, E> {
    let mut new = String::with_capacity(haystack.len());
    let mut last_match = 0;

    for caps in re.captures_iter(haystack) {
        let m = caps.get(0).unwrap();

        new.push_str(&haystack[last_match..m.start()]);
        new.push_str(&replacement(&caps)?);

        last_match = m.end();
    }

    new.push_str(&haystack[last_match..]);

    Ok(new)
}

static TRUE_FALSE_REPLACEMENT_RE: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r"\$\{(\w+)(?::-([^|}]+))?\|([^|]+)\|([^}]+)}")
        .expect("Regex failed to compile")
});

static VARIABLE_REPLACEMENT_RE: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r"\$\{(\w+)(?::-([^}]+))?}").expect("Regex failed to compile")
});

/// Replace all `${KEY}` or `${KEY:-default}` placeholders in a string according to the mapping.
/// Also handles `${KEY:-default|True value|false value}` syntax.
fn process_replacements(
    input: &str,
    replacement_map: &HashMap<String, String>,
) -> Result<String> {
    process_replacements_with_env(input, replacement_map, |key| std::env::var(key).ok())
}

fn process_replacements_with_env(
    input: &str,
    replacement_map: &HashMap<String, String>,
    get_env: impl Fn(&str) -> Option<String>,
) -> Result<String> {
    debug!("processing replacements for line '{input}'");

    // handle ${VAR:-default|true value|false value} syntax
    let replacement = |caps: &regex::Captures| -> Result<String> {
        let key = &caps[1];
        let default = caps.get(2).map(|m| m.as_str().to_string());
        let true_val = &caps[3];
        let false_val = &caps[4];

        let value = lookup_replacement_value(key, replacement_map, &get_env).or(default);

        match value {
            Some(v) if v.eq_ignore_ascii_case("true") => Ok(true_val.to_string()),
            Some(_) => Ok(false_val.to_string()),
            None => Err(exec_datafusion_err!("Missing value for key '{key}'")),
        }
    };
    let input = replace_all(&TRUE_FALSE_REPLACEMENT_RE, input, replacement)?;

    // handle ${KEY} and ${KEY:-default}`
    let replacement = |caps: &regex::Captures| -> Result<String> {
        let key = &caps[1];
        let default = caps.get(2);

        if let Some(v) = lookup_replacement_value(key, replacement_map, &get_env) {
            return Ok(v.to_string());
        }

        // use default if it was set
        if let Some(def) = default {
            Ok(def.as_str().to_string())
        } else {
            Err(exec_datafusion_err!("Missing value for key '{key}'"))
        }
    };

    replace_all(&VARIABLE_REPLACEMENT_RE, &input, replacement)
}

fn lookup_replacement_value(
    key: &str,
    replacement_map: &HashMap<String, String>,
    get_env: &impl Fn(&str) -> Option<String>,
) -> Option<String> {
    if let Some(v) = replacement_map.get(&key.to_lowercase()) {
        return Some(v.to_string());
    }

    // look in env variables
    get_env(&key.to_uppercase())
}

fn read_query_from_reader(
    reader: &mut BenchmarkFileReader,
    sql: &str,
    header: &str,
) -> Result<BenchmarkQuery> {
    let column_count = header.len();
    let mut expected_result = vec![];
    let mut line = String::new();
    let mut reader_result = reader.read_line(&mut line);

    loop {
        match reader_result {
            Some(Ok(_)) => {
                if is_comment_line(&line) {
                    // comment, ignore
                } else if is_blank_line(&line) {
                    break;
                } else {
                    let result_splits: Vec<&str> = line.split(['\t', '|']).collect();

                    if result_splits.len() != column_count {
                        return Err(exec_datafusion_err!(
                            "{} {line}",
                            reader.format_exception(&format!(
                                "expected {} values but got {}",
                                column_count,
                                result_splits.len(),
                            ))
                        ));
                    }

                    expected_result
                        .push(result_splits.into_iter().map(|s| s.to_string()).collect());
                }
            }
            Some(Err(e)) => return Err(e),
            None => break,
        }

        // Clear the line buffer for the next iteration.
        line.clear();
        reader_result = reader.read_line(&mut line);
    }

    Ok(BenchmarkQuery {
        path: None,
        query: sql.to_string(),
        column_count,
        expected_result,
    })
}

async fn read_query_from_file(
    ctx: &SessionContext,
    path: impl AsRef<Path>,
    replacement_mapping: &HashMap<String, String>,
) -> Result<BenchmarkQuery> {
    // Process replacements in file path
    let path = path.as_ref().to_string_lossy();
    let path = process_replacements(&path, replacement_mapping)?;
    let df: DataFrame = ctx
        .read_csv(
            path.clone(),
            CsvReadOptions::new()
                .has_header(true)
                .delimiter(b'|')
                .null_regex(Some("NULL".to_string()))
                // we only want string values, we do not want to infer the schema
                .schema_infer_max_records(0),
        )
        .await?;

    // Get schema to determine column count
    let schema = df.schema();
    let column_count = schema.fields().len();

    if column_count == 0 {
        return Err(exec_datafusion_err!(
            "Result file {path} did not contain any columns"
        ));
    }

    // Execute and collect results
    let batches = df.collect().await?;
    // Convert record batches to string vectors
    let expected_result = format_record_batches(&batches)?;

    Ok(BenchmarkQuery {
        path: Some(path),
        query: String::new(),
        column_count,
        expected_result,
    })
}

fn format_record_batches(
    batches: &[RecordBatch],
) -> Result<Vec<Vec<String>>, DataFusionError> {
    let mut expected_result = vec![];
    let arrow_format_options = FormatOptions::default()
        .with_null("NULL")
        .with_display_error(true);

    for batch in batches {
        let schema = batch.schema_ref();

        let formatters = batch
            .columns()
            .iter()
            .zip(schema.fields().iter())
            .map(|(c, field)| make_array_formatter(c, &arrow_format_options, Some(field)))
            .collect::<Result<Vec<_>, ArrowError>>()?;

        for row in 0..batch.num_rows() {
            let mut cells = vec![];
            for formatter in &formatters {
                cells.push(formatter.value(row).to_string());
            }
            expected_result.push(cells);
        }
    }

    Ok(expected_result)
}

fn make_array_formatter<'a>(
    array: &'a dyn Array,
    options: &FormatOptions<'a>,
    field: Option<&'a Field>,
) -> Result<ArrayFormatter<'a>, ArrowError> {
    match options.formatter_factory() {
        None => ArrayFormatter::try_new(array, options),
        Some(formatters) => formatters
            .create_array_formatter(array, options, field)
            .transpose()
            .unwrap_or_else(|| ArrayFormatter::try_new(array, options)),
    }
}

fn resolve_benchmark_file_path(path: &str) -> PathBuf {
    let path = PathBuf::from(path);
    if path.is_absolute() || path.exists() {
        path
    } else {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::prelude::SessionContext;
    use std::fs;
    use std::path::{Path, PathBuf};
    use tempfile::{TempDir, tempdir};

    fn write_test_file(temp_dir: &TempDir, name: &str, contents: &str) -> PathBuf {
        let path = temp_dir.path().join(name);
        fs::write(&path, contents).expect("failed to write benchmark test file");
        path
    }

    async fn parse_benchmark_file(path: &Path) -> Result<SqlBenchmark> {
        let ctx = SessionContext::new();
        let path_string = path.to_string_lossy().into_owned();
        SqlBenchmark::new(&ctx, &path_string, "/tmp").await
    }

    async fn parse_benchmark(contents: &str) -> Result<SqlBenchmark> {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let path = write_test_file(&temp_dir, "parser.benchmark", contents);

        parse_benchmark_file(&path).await
    }

    async fn assert_parse_error(contents: &str, expected_message: &str) {
        let error = parse_benchmark(contents)
            .await
            .expect_err("benchmark parsing should fail");

        let message = error.to_string();
        assert!(
            message.contains(expected_message),
            "expected error containing {expected_message:?}, got {message:?}"
        );
    }

    fn assert_result_error_contains<T: Debug>(result: Result<T>, expected_message: &str) {
        let error = result.expect_err("operation should fail");
        let message = error.to_string();
        assert!(
            message.contains(expected_message),
            "expected error containing {expected_message:?}, got {message:?}"
        );
    }

    fn formatted_last_results(benchmark: &SqlBenchmark) -> Vec<Vec<String>> {
        format_record_batches(
            benchmark
                .last_results
                .as_ref()
                .expect("last results should be set"),
        )
        .expect("results should format")
    }

    fn read_all_files_in_dir(path: &Path) -> String {
        let mut entries = fs::read_dir(path)
            .expect("directory should be readable")
            .filter_map(Result::ok)
            .map(|entry| entry.path())
            .filter(|path| path.is_file())
            .collect::<Vec<_>>();
        entries.sort();

        let mut contents = String::new();
        for path in entries {
            contents
                .push_str(&fs::read_to_string(path).expect("file should be readable"));
        }
        contents
    }

    fn replacement_map(entries: &[(&str, &str)]) -> HashMap<String, String> {
        let mut replacements = HashMap::new();
        for (key, value) in entries {
            insert_replacement(&mut replacements, key, value.to_string());
        }
        replacements
    }

    #[test]
    fn resolves_sql_benchmarks_paths_from_manifest_directory() {
        let path =
            resolve_benchmark_file_path("sql_benchmarks/clickbench/init/set_config.sql");

        assert!(path.exists(), "resolved path should exist: {path:?}");
    }

    #[tokio::test]
    async fn run_returns_row_count_when_not_saving_results() {
        let contents = "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2), (3)) AS t(v)\n";
        let mut benchmark = parse_benchmark(contents).await.unwrap();
        let ctx = SessionContext::new();

        benchmark.initialize(&ctx).await.unwrap();
        let row_count = benchmark.run(&ctx, false).await.unwrap();

        assert_eq!(row_count, 3);
    }

    #[tokio::test]
    async fn streaming_run_reports_last_select_row_count_after_ddl() {
        let contents = "name Q15\n\nrun\nCREATE VIEW v AS SELECT 1 AS value;\nSELECT * FROM v;\nDROP VIEW v;\n";
        let mut benchmark = parse_benchmark(contents).await.unwrap();
        let ctx = SessionContext::new();

        benchmark.initialize(&ctx).await.unwrap();
        let row_count = benchmark.run(&ctx, false).await.unwrap();

        assert_eq!(row_count, 1);
        assert!(ctx.table("v").await.is_err(), "trailing DDL should execute");
    }

    #[tokio::test]
    async fn run_returns_row_count_when_saving_results() {
        let contents = "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2)) AS t(v)\n";
        let mut benchmark = parse_benchmark(contents).await.unwrap();
        let ctx = SessionContext::new();

        benchmark.initialize(&ctx).await.unwrap();
        let row_count = benchmark.run(&ctx, true).await.unwrap();

        assert_eq!(row_count, 2);
        assert_eq!(
            formatted_last_results(&benchmark),
            vec![vec!["1"], vec!["2"]]
        );
    }

    fn env_map(entries: &[(&str, &str)]) -> HashMap<String, String> {
        entries
            .iter()
            .map(|(key, value)| (key.to_string(), value.to_string()))
            .collect()
    }

    // Replacement tests cover benchmark variable expansion syntax.

    #[test]
    fn process_replacements_replaces_map_values_case_insensitively() {
        let replacements = replacement_map(&[
            ("BENCH_NAME", "tpch"),
            ("QUERY_NUMBER_PADDED", "01"),
            ("format_1", "parquet"),
        ]);

        let actual = process_replacements_with_env(
            "${bench_name}/q${query_number_padded}.${FORMAT_1}",
            &replacements,
            |_| None,
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "tpch/q01.parquet");
    }

    #[test]
    fn process_replacements_uses_env_when_map_value_is_missing() {
        let replacements = HashMap::new();
        let env = env_map(&[("DATA_DIR", "/tmp/data")]);

        let actual = process_replacements_with_env(
            "${data_dir}/lineitem.parquet",
            &replacements,
            |key| env.get(key).cloned(),
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "/tmp/data/lineitem.parquet");
    }

    #[test]
    fn process_replacements_prefers_map_over_env() {
        let replacements = replacement_map(&[("BENCH_SIZE", "10")]);
        let env = env_map(&[("BENCH_SIZE", "100")]);

        let actual =
            process_replacements_with_env("sf${BENCH_SIZE}", &replacements, |key| {
                env.get(key).cloned()
            })
            .expect("replacement should succeed");

        assert_eq!(actual, "sf10");
    }

    #[test]
    #[expect(
        clippy::literal_string_with_formatting_args,
        reason = "The `${VAR:-default}` braces are shell-style placeholders, not format args"
    )]
    fn process_replacements_uses_default_for_missing_variable() {
        let replacements = HashMap::new();

        let actual = process_replacements_with_env(
            "load_${BENCH_SUBGROUP:-groupby}_${FILE_TYPE:-csv}.sql",
            &replacements,
            |_| None,
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_groupby_csv.sql");
    }

    #[test]
    fn process_replacements_reports_missing_variable_without_default() {
        let replacements = HashMap::new();

        let error = process_replacements_with_env("${MISSING}", &replacements, |_| None)
            .expect_err("replacement should fail");

        assert!(
            error
                .to_string()
                .contains("Missing value for key 'MISSING'"),
            "unexpected error: {error}"
        );
    }

    #[test]
    fn process_replacements_applies_true_false_true_branch() {
        let replacements = HashMap::new();
        let env = env_map(&[("USE_PARQUET", "TrUe")]);

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-false|parquet|csv}.sql",
            &replacements,
            |key| env.get(key).cloned(),
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_parquet.sql");
    }

    #[test]
    fn process_replacements_applies_true_false_false_branch() {
        let replacements = HashMap::new();
        let env = env_map(&[("USE_PARQUET", "false")]);

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-true|parquet|csv}.sql",
            &replacements,
            |key| env.get(key).cloned(),
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_csv.sql");
    }

    #[test]
    fn process_replacements_uses_map_for_true_false_branch() {
        let replacements = replacement_map(&[("USE_PARQUET", "true")]);

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-false|parquet|csv}.sql",
            &replacements,
            |_| None,
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_parquet.sql");
    }

    #[test]
    fn process_replacements_prefers_map_over_env_for_true_false_branch() {
        let replacements = replacement_map(&[("USE_PARQUET", "false")]);
        let env = env_map(&[("USE_PARQUET", "true")]);

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-true|parquet|csv}.sql",
            &replacements,
            |key| env.get(key).cloned(),
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_csv.sql");
    }

    #[test]
    fn process_replacements_uses_true_false_default_for_missing_true_value() {
        let replacements = HashMap::new();

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-true|parquet|csv}.sql",
            &replacements,
            |_| None,
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_parquet.sql");
    }

    #[test]
    fn process_replacements_uses_true_false_default_for_missing_false_value() {
        let replacements = HashMap::new();

        let actual = process_replacements_with_env(
            "load_${USE_PARQUET:-false|parquet|csv}.sql",
            &replacements,
            |_| None,
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "load_csv.sql");
    }

    #[test]
    fn process_replacements_reports_missing_true_false_variable_without_default() {
        let replacements = HashMap::new();

        let error = process_replacements_with_env(
            "load_${USE_PARQUET|parquet|csv}.sql",
            &replacements,
            |_| None,
        )
        .expect_err("replacement should fail");

        assert!(
            error
                .to_string()
                .contains("Missing value for key 'USE_PARQUET'"),
            "unexpected error: {error}"
        );
    }

    #[test]
    fn process_replacements_resolves_variables_after_true_false_replacement() {
        let replacements = replacement_map(&[("FILE_TYPE", "parquet")]);
        let env = env_map(&[("USE_TYPED_PATH", "true")]);

        let actual = process_replacements_with_env(
            "${USE_TYPED_PATH:-false|data.${FILE_TYPE}|data.csv}",
            &replacements,
            |key| env.get(key).cloned(),
        )
        .expect("replacement should succeed");

        assert_eq!(actual, "data.parquet");
    }

    #[test]
    fn process_replacements_leaves_unsupported_placeholder_syntax_unchanged() {
        let replacements = HashMap::new();

        let actual =
            process_replacements_with_env("${BAD-KEY:-fallback}", &replacements, |_| {
                None
            })
            .expect("unsupported placeholder should not match replacement regex");

        assert_eq!(actual, "${BAD-KEY:-fallback}");
    }

    // Parser tests cover benchmark directives and parse-time validation.

    #[tokio::test]
    async fn parser_accepts_metadata_expect_echo_and_sql_sections() {
        let benchmark = parse_benchmark(
            r#"
# top-level comments are ignored
name Parser Success
group Parser Group
subgroup Parser Subgroup
expect_plan ProjectionExec with details
echo hello from parser

load
-- query comments are ignored
CREATE TABLE t AS VALUES (1);

init
CREATE VIEW v AS SELECT * FROM t;

run
SELECT * FROM v;

cleanup
DROP VIEW v;
"#,
        )
        .await
        .expect("benchmark should parse");

        assert_eq!(benchmark.name(), "Parser Success");
        assert_eq!(benchmark.group(), "Parser Group");
        assert_eq!(benchmark.subgroup(), "Parser Subgroup");
        assert_eq!(benchmark.expect, vec!["ProjectionExec with details"]);
        assert_eq!(benchmark.echo, vec!["hello from parser"]);
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Load)
                .expect("load query"),
            &vec!["CREATE TABLE t AS VALUES (1);".to_string()]
        );
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Init)
                .expect("init query"),
            &vec!["CREATE VIEW v AS SELECT * FROM t;".to_string()]
        );
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["SELECT * FROM v;".to_string()]
        );
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Cleanup)
                .expect("cleanup query"),
            &vec!["DROP VIEW v;".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_splits_inline_run_block_on_semicolon_newline() {
        let benchmark = parse_benchmark(
            r#"
run
CREATE TABLE t AS SELECT 1 AS value;
SELECT value + 1 AS value FROM t;
DROP TABLE t;
"#,
        )
        .await
        .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec![
                "CREATE TABLE t AS SELECT 1 AS value;".to_string(),
                "SELECT value + 1 AS value FROM t;".to_string(),
                "DROP TABLE t;".to_string(),
            ]
        );
    }

    #[tokio::test]
    async fn parser_accepts_assert_with_expected_rows() {
        let benchmark = parse_benchmark(
            r#"
assert II
select 1, 'one'
----
1|one
2	two
"#,
        )
        .await
        .expect("benchmark should parse");

        let query = benchmark
            .assert_queries()
            .first()
            .expect("assert query should be parsed");

        assert_eq!(query.column_count, 2);
        assert!(query.query.contains("select 1, 'one'"));
        assert_eq!(
            query.expected_result,
            vec![
                vec!["1".to_string(), "one".to_string()],
                vec!["2".to_string(), "two".to_string()]
            ]
        );
    }

    #[tokio::test]
    async fn parser_accepts_result_query_with_expected_rows() {
        let benchmark = parse_benchmark(
            r#"
result_query II
select 1, 'one'
----
1|one
NULL|(empty)
"#,
        )
        .await
        .expect("benchmark should parse");

        let query = benchmark
            .result_queries()
            .first()
            .expect("result query should be parsed");

        assert_eq!(query.path, None);
        assert_eq!(query.column_count, 2);
        assert!(query.query.contains("select 1, 'one'"));
        assert_eq!(
            query.expected_result,
            vec![
                vec!["1".to_string(), "one".to_string()],
                vec!["NULL".to_string(), "(empty)".to_string()]
            ]
        );
    }

    #[tokio::test]
    async fn parser_records_result_file_without_parsing_contents() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path =
            write_test_file(&temp_dir, "result.csv", "col_a|col_b\n1|one\nNULL|two\n");
        let benchmark_path = write_test_file(
            &temp_dir,
            "result.benchmark",
            &format!("result {}\n", result_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        let query = benchmark
            .result_queries()
            .first()
            .expect("result file should be parsed");

        assert_eq!(query.path, Some(result_path.to_string_lossy().into_owned()));
        assert_eq!(query.column_count, 0);
        assert!(query.expected_result.is_empty());
    }

    #[tokio::test]
    async fn parser_accepts_include_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let include_path =
            write_test_file(&temp_dir, "include.benchmark", "run\nselect 1\n");

        let benchmark_path = write_test_file(
            &temp_dir,
            "include_driver.benchmark",
            &format!("include {}\n", include_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        let benchmark = result.expect("benchmark should parse");
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["select 1".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_accepts_template_file_with_parameters() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let template_path = write_test_file(
            &temp_dir,
            "template_success.benchmark",
            "# template comments are ignored\nrun\n-- query comments are ignored\nselect '${TABLE_NAME}', '${BENCHMARK_DIR}'\n",
        );

        let benchmark_path = write_test_file(
            &temp_dir,
            "template_success_driver.benchmark",
            &format!(
                "template {}\n# parameter comments are ignored\nTABLE_NAME=orders\n",
                template_path.display()
            ),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        let benchmark = result.expect("benchmark should parse");
        assert_eq!(benchmark.benchmark_path(), template_path.as_path());
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["select 'orders', '/tmp'".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_trims_template_parameter_keys_and_values() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let template_path = write_test_file(
            &temp_dir,
            "template_trim.benchmark",
            "run\nselect '${TABLE_NAME}'\n",
        );

        let benchmark_path = write_test_file(
            &temp_dir,
            "template_trim_driver.benchmark",
            &format!(
                "template {}\n  TABLE_NAME = orders  \n",
                template_path.display()
            ),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["select 'orders'".to_string()]
        );
        assert_eq!(
            benchmark.replacement_mapping().get("table_name"),
            Some(&"orders".to_string())
        );
    }

    #[tokio::test]
    async fn parser_preserves_expected_result_cell_whitespace() {
        let benchmark = parse_benchmark("assert I\nselect '  x  '\n----\n  x  \n")
            .await
            .expect("benchmark should parse");

        let query = benchmark
            .assert_queries()
            .first()
            .expect("assert query should be parsed");

        assert_eq!(query.expected_result, vec![vec!["  x  ".to_string()]]);
    }

    #[tokio::test]
    async fn parser_accepts_indented_comments_and_blank_lines() {
        let benchmark =
            parse_benchmark("  # comment\n  -- comment\n  run\n  select 1\n   \n")
                .await
                .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["select 1".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_accepts_case_insensitive_query_directives() {
        let benchmark = parse_benchmark("RUN\nselect 1\n")
            .await
            .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["select 1".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_accepts_query_file_and_splits_statements() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "queries.sql",
            "-- leading comment\nSELECT 1 AS value;\nSELECT 2 AS value;\n\n# another comment\nWITH t AS (SELECT 3 AS value) SELECT * FROM t;\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "query_file.benchmark",
            &format!("run {}\n", query_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run queries"),
            &vec![
                "SELECT 1 AS value;".to_string(),
                "SELECT 2 AS value;".to_string(),
                "WITH t AS (SELECT 3 AS value) SELECT * FROM t;".to_string(),
            ]
        );
    }

    #[tokio::test]
    async fn parser_accepts_replacements_in_query_file_path() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path =
            write_test_file(&temp_dir, "queries.sql", "SELECT 5 AS value;\n");
        let template_path = write_test_file(
            &temp_dir,
            "query_file_path_template.benchmark",
            "run ${QUERY_PATH}\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "query_file_path_driver.benchmark",
            &format!(
                "template {}\nQUERY_PATH={}\n",
                template_path.display(),
                query_path.display()
            ),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["SELECT 5 AS value;".to_string()]
        );
    }

    #[tokio::test]
    #[expect(
        clippy::literal_string_with_formatting_args,
        reason = "The `${VAR:-default}` braces are shell-style placeholders, not format args"
    )]
    async fn parser_applies_data_dir_replacement_in_load_query_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let data_dir = temp_dir.path().join("non_default_data");
        let csv_dir = data_dir.join("tpch_sf1/csv/generated");
        fs::create_dir_all(&csv_dir).expect("failed to create generated data directory");
        fs::write(csv_dir.join("generated.1.csv"), "value\n42\n")
            .expect("failed to write generated csv file");

        let load_path = write_test_file(
            &temp_dir,
            "load_generated_csv.sql",
            "CREATE EXTERNAL TABLE generated(value INT) STORED AS CSV LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/csv/generated/generated.1.csv' OPTIONS ('format.has_header' 'true');\n",
        );
        let template_path = write_test_file(
            &temp_dir,
            "load_file_template.benchmark",
            &format!(
                "load {}\n\nrun\nSELECT value FROM generated;\n",
                load_path.display()
            ),
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "load_file_driver.benchmark",
            &format!(
                "template {}\nDATA_DIR={}\n",
                template_path.display(),
                data_dir.display()
            ),
        );

        let ctx = SessionContext::new();
        let path_string = benchmark_path.to_string_lossy().into_owned();
        let mut benchmark = SqlBenchmark::new(&ctx, &path_string, "/tmp")
            .await
            .expect("benchmark should parse");

        let load_queries = benchmark
            .queries()
            .get(&QueryDirective::Load)
            .expect("load queries");
        assert_eq!(load_queries.len(), 1);
        assert!(
            load_queries.iter().all(|query| !query.contains("${")),
            "all placeholders should be replaced: {load_queries:?}"
        );
        let expected_location = format!(
            "LOCATION '{}/tpch_sf1/csv/generated/generated.1.csv'",
            data_dir.display()
        );
        assert!(
            load_queries[0].contains(&expected_location),
            "all load locations should use the non-default DATA_DIR: {load_queries:?}"
        );

        benchmark
            .initialize(&ctx)
            .await
            .expect("benchmark should load generated csv file");
        benchmark
            .run(&ctx, true)
            .await
            .expect("benchmark should read generated csv file");

        assert_eq!(formatted_last_results(&benchmark), vec![vec!["42"]]);
    }

    #[tokio::test]
    async fn parser_rejects_inline_sql_when_query_file_is_provided() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path =
            write_test_file(&temp_dir, "queries.sql", "SELECT 1 AS value;\n");
        let benchmark_path = write_test_file(
            &temp_dir,
            "query_file_with_inline_body.benchmark",
            &format!("run {}\nSELECT 999 AS value;\n", query_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(
            result,
            "run directive must use either a query file or inline SQL, not both",
        );
    }

    #[tokio::test]
    async fn parser_rejects_inline_sql_when_load_file_is_provided() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "load.sql",
            "CREATE TABLE t AS SELECT 1 AS value;\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "load_file_with_inline_body.benchmark",
            &format!(
                "load {}\nCREATE TABLE u AS SELECT 2 AS value;\n",
                query_path.display()
            ),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(
            result,
            "load directive must use either a query file or inline SQL, not both",
        );
    }

    #[tokio::test]
    async fn parser_rejects_inline_sql_when_init_file_is_provided() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "init.sql",
            "CREATE VIEW v AS SELECT 1 AS value;\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "init_file_with_inline_body.benchmark",
            &format!(
                "init {}\nCREATE VIEW w AS SELECT 2 AS value;\n",
                query_path.display()
            ),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(
            result,
            "init directive must use either a query file or inline SQL, not both",
        );
    }

    #[tokio::test]
    async fn parser_rejects_inline_sql_when_cleanup_file_is_provided() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(&temp_dir, "cleanup.sql", "DROP TABLE t;\n");
        let benchmark_path = write_test_file(
            &temp_dir,
            "cleanup_file_with_inline_body.benchmark",
            &format!("cleanup {}\nDROP TABLE u;\n", query_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(
            result,
            "cleanup directive must use either a query file or inline SQL, not both",
        );
    }

    #[tokio::test]
    async fn parser_ignores_query_file_with_only_comments_and_blank_lines() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "queries.sql",
            "# comment\n\n-- another comment\n\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "empty_query_file.benchmark",
            &format!("run {}\n", query_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        assert!(!benchmark.queries().contains_key(&QueryDirective::Run));
    }

    #[tokio::test]
    async fn parser_splits_query_file_with_windows_line_endings() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "queries.sql",
            "SELECT 1 AS value;\r\nSELECT 2 AS value;\r\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "windows_query_file.benchmark",
            &format!("run {}\n", query_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run queries"),
            &vec![
                "SELECT 1 AS value;".to_string(),
                "SELECT 2 AS value;".to_string()
            ]
        );
    }

    #[tokio::test]
    async fn parser_rejects_unknown_command() {
        assert_parse_error("wat\n", "Unrecognized command: wat").await;
    }

    #[tokio::test]
    async fn parser_rejects_assert_without_column_count() {
        assert_parse_error(
            "assert\nselect 1\n----\n1\n",
            "assert must be followed by a column count",
        )
        .await;
    }

    #[tokio::test]
    async fn parser_rejects_assert_without_result_separator() {
        assert_parse_error(
            "assert I\nselect 1\n1\n",
            "assert must be followed by a query and a result (separated by ----)",
        )
        .await;
    }

    #[tokio::test]
    async fn parser_rejects_result_query_without_separator() {
        assert_parse_error(
            "result_query I\nselect 1\n1\n",
            "result_query must be followed by a query and a result (separated by ----)",
        )
        .await;
    }

    #[tokio::test]
    async fn parser_rejects_result_query_with_wrong_column_count() {
        assert_parse_error(
            "result_query II\nselect 1\n----\n1\n",
            "expected 2 values but got 1",
        )
        .await;
    }

    #[tokio::test]
    async fn parser_rejects_multiple_result_queries() {
        assert_parse_error(
            "result_query I\nselect 1\n----\n1\n\nresult_query I\nselect 2\n----\n2\n",
            "multiple results found",
        )
        .await;
    }

    #[tokio::test]
    async fn parser_rejects_duplicate_run_directives() {
        assert_parse_error("run\nselect 1\n\nrun\nselect 2\n", "Multiple calls to run")
            .await;
    }

    #[tokio::test]
    async fn parser_accepts_multiple_load_directives() {
        let benchmark = parse_benchmark(
            "load\nCREATE TABLE t AS SELECT 1;\n\nload\nCREATE TABLE u AS SELECT 2;\n",
        )
        .await
        .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Load)
                .expect("load queries"),
            &vec![
                "CREATE TABLE t AS SELECT 1;".to_string(),
                "CREATE TABLE u AS SELECT 2;".to_string(),
            ]
        );
    }

    #[tokio::test]
    async fn parser_accepts_multiple_init_directives() {
        let benchmark = parse_benchmark(
            "init\nCREATE VIEW v AS SELECT 1;\n\ninit\nCREATE VIEW w AS SELECT 2;\n",
        )
        .await;

        let benchmark = benchmark.expect("benchmark should parse");
        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Init)
                .expect("init queries"),
            &vec![
                "CREATE VIEW v AS SELECT 1;".to_string(),
                "CREATE VIEW w AS SELECT 2;".to_string(),
            ]
        );
    }

    #[tokio::test]
    async fn parser_accepts_multiple_cleanup_directives() {
        let benchmark =
            parse_benchmark("cleanup\nDROP TABLE t;\n\ncleanup\nDROP TABLE u;\n")
                .await
                .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Cleanup)
                .expect("cleanup queries"),
            &vec!["DROP TABLE t;".to_string(), "DROP TABLE u;".to_string(),]
        );
    }

    #[tokio::test]
    async fn parser_rejects_missing_query_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let missing_path = temp_dir.path().join("missing.sql");
        let benchmark_path = write_test_file(
            &temp_dir,
            "missing_query_file.benchmark",
            &format!("run {}\n", missing_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(result, "Failed to read query file");
    }

    #[tokio::test]
    async fn parser_rejects_template_with_invalid_parameter_assignment() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let template_path =
            write_test_file(&temp_dir, "template.benchmark", "run\nselect 1\n");

        let benchmark_path = write_test_file(
            &temp_dir,
            "template_driver.benchmark",
            &format!("template {}\nINVALID\n", template_path.display()),
        );

        let ctx = SessionContext::new();
        let benchmark_path_string = benchmark_path.to_string_lossy().into_owned();
        let result = SqlBenchmark::new(&ctx, &benchmark_path_string, "/tmp").await;

        let error = result.expect_err("benchmark parsing should fail");
        let message = error.to_string();
        assert!(
            message.contains("Expected a template parameter in the form of X=Y"),
            "expected template parameter error, got {message:?}"
        );
    }

    #[tokio::test]
    async fn parser_rejects_metadata_and_result_directives_without_values() {
        assert_parse_error("name\n", "name must be followed by a value").await;
        assert_parse_error("group\n", "group must be followed by a value").await;
        assert_parse_error("subgroup\n", "subgroup must be followed by a value").await;
        assert_parse_error(
            "expect_plan\n",
            "expect_plan must be followed by a string to search in the physical plan",
        )
        .await;
        assert_parse_error("echo\n", "Echo requires an argument").await;
        assert_parse_error(
            "result\n",
            "result must be followed by a path to a result file",
        )
        .await;
        assert_parse_error("include\n", "include requires a single argument").await;
        assert_parse_error("template\n", "template requires a single template path")
            .await;
    }

    #[tokio::test]
    async fn parser_rejects_include_and_template_with_too_many_arguments() {
        assert_parse_error("include a b\n", "include requires a single argument").await;
        assert_parse_error("template a b\n", "template requires a single template path")
            .await;
    }

    #[tokio::test]
    async fn parser_rejects_missing_include_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let missing_path = temp_dir.path().join("missing_include.benchmark");
        let benchmark_path = write_test_file(
            &temp_dir,
            "missing_include_driver.benchmark",
            &format!("include {}\n", missing_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(result, "No such file");
    }

    #[tokio::test]
    async fn parser_rejects_missing_template_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let missing_path = temp_dir.path().join("missing_template.benchmark");
        let benchmark_path = write_test_file(
            &temp_dir,
            "missing_template_driver.benchmark",
            &format!("template {}\n", missing_path.display()),
        );

        let result = parse_benchmark_file(&benchmark_path).await;

        assert_result_error_contains(result, "No such file");
    }

    #[tokio::test]
    async fn parser_uses_metadata_values_as_replacements() {
        let benchmark = parse_benchmark(
            r#"
name Q01
group tpch
subgroup sf1

run
SELECT '${BENCH_NAME}', '${BENCH_GROUP}', '${BENCH_SUBGROUP}'
"#,
        )
        .await
        .expect("benchmark should parse");

        assert_eq!(
            benchmark
                .queries()
                .get(&QueryDirective::Run)
                .expect("run query"),
            &vec!["SELECT 'Q01', 'tpch', 'sf1'".to_string()]
        );
    }

    #[tokio::test]
    async fn parser_accepts_replacement_in_result_file_path() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path = write_test_file(&temp_dir, "result.csv", "value\n1\n");
        let template_path = write_test_file(
            &temp_dir,
            "result_path_template.benchmark",
            "result ${RESULT_PATH}\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "result_path_driver.benchmark",
            &format!(
                "template {}\nRESULT_PATH={}\n",
                template_path.display(),
                result_path.display()
            ),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        let query = benchmark
            .result_queries()
            .first()
            .expect("result query should be parsed");
        assert_eq!(query.path, Some(result_path.to_string_lossy().into_owned()));
        assert_eq!(query.column_count, 0);
        assert!(query.expected_result.is_empty());
    }

    #[tokio::test]
    async fn parser_rejects_missing_replacement_in_result_file_path() {
        assert_parse_error("result ${MISSING_RESULT_PATH}\n", "Missing value for key")
            .await;
    }

    #[tokio::test]
    async fn parser_accepts_missing_result_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let missing_path = temp_dir.path().join("missing_result.csv");
        let benchmark_path = write_test_file(
            &temp_dir,
            "missing_result_file.benchmark",
            &format!("result {}\n", missing_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        let query = benchmark
            .result_queries()
            .first()
            .expect("result file should be parsed");
        assert_eq!(
            query.path,
            Some(missing_path.to_string_lossy().into_owned())
        );
        assert!(query.expected_result.is_empty());
    }

    #[tokio::test]
    async fn parser_accepts_malformed_result_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path = temp_dir.path().join("malformed_result.csv");
        fs::write(&result_path, [0xff]).expect("failed to write malformed result file");
        let benchmark_path = write_test_file(
            &temp_dir,
            "malformed_result_file.benchmark",
            &format!("result {}\n", result_path.display()),
        );

        let benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");

        let query = benchmark
            .result_queries()
            .first()
            .expect("result file should be parsed");
        assert_eq!(query.path, Some(result_path.to_string_lossy().into_owned()));
        assert!(query.expected_result.is_empty());
    }

    // Lifecycle tests cover initialization, assertions, and cleanup execution.

    #[tokio::test]
    async fn initialize_executes_load_before_init_and_is_idempotent() {
        let mut benchmark = parse_benchmark(
            r#"
load
CREATE TABLE t AS SELECT 1 AS value;

load
CREATE TABLE u AS SELECT value + 1 AS value FROM t;

init
CREATE TABLE v AS SELECT value + 1 AS value FROM u;

init
CREATE TABLE initialized AS SELECT value + 1 AS value FROM v;

run
SELECT value FROM initialized;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark
            .initialize(&ctx)
            .await
            .expect("initialize should succeed");
        benchmark
            .initialize(&ctx)
            .await
            .expect("second initialize should be a no-op");

        assert!(benchmark.is_loaded());

        let rows = ctx
            .sql("SELECT value FROM initialized")
            .await
            .expect("query should plan")
            .collect()
            .await
            .expect("query should run");

        assert_eq!(format_record_batches(&rows).unwrap(), vec![vec!["4"]]);
    }

    #[tokio::test]
    async fn initialize_rejects_benchmark_without_run_query() {
        let mut benchmark = parse_benchmark(
            r#"
load
CREATE TABLE t AS SELECT 1 AS value;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.initialize(&ctx).await,
            "Invalid benchmark file: no \"run\" query specified",
        );
    }

    #[tokio::test]
    async fn initialize_propagates_load_query_failures() {
        let mut benchmark = parse_benchmark(
            r#"
load
CREATE TABLE t AS SELECT * FROM missing_load_table;

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.initialize(&ctx).await,
            "missing_load_table",
        );
    }

    #[tokio::test]
    async fn initialize_propagates_init_query_failures() {
        let mut benchmark = parse_benchmark(
            r#"
init
CREATE TABLE t AS SELECT * FROM missing_init_table;

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.initialize(&ctx).await,
            "missing_init_table",
        );
    }

    #[tokio::test]
    async fn cleanup_executes_cleanup_queries() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 1;

cleanup
CREATE TABLE cleanup_marker_a AS SELECT 7 AS value;

cleanup
CREATE TABLE cleanup_marker_b AS SELECT value + 1 AS value FROM cleanup_marker_a;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.cleanup(&ctx).await.expect("cleanup should run");

        let rows = ctx
            .sql("SELECT value FROM cleanup_marker_b")
            .await
            .expect("query should plan")
            .collect()
            .await
            .expect("query should run");
        assert_eq!(format_record_batches(&rows).unwrap(), vec![vec!["8"]]);
    }

    #[tokio::test]
    async fn cleanup_propagates_query_failures() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 1;

cleanup
SELECT * FROM missing_cleanup_table;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.cleanup(&ctx).await,
            "missing_cleanup_table",
        );
    }

    #[tokio::test]
    async fn assert_executes_assert_queries_successfully() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT 1 AS value
----
1

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.assert(&ctx).await.expect("assert should pass");
    }

    #[tokio::test]
    async fn assert_accepts_null_expected_for_empty_actual() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT '' AS value
----
NULL

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.assert(&ctx).await.expect("assert should pass");
    }

    #[tokio::test]
    async fn assert_accepts_empty_marker_for_empty_actual() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT '' AS value
----
(empty)

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.assert(&ctx).await.expect("assert should pass");
    }

    #[tokio::test]
    async fn assert_accepts_empty_marker_for_null_actual() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT CAST(NULL AS VARCHAR) AS value
----
(empty)

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.assert(&ctx).await.expect("assert should pass");
    }

    #[tokio::test]
    async fn assert_succeeds_with_zero_actual_and_expected_rows() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT 1 AS value WHERE false
----

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.assert(&ctx).await.expect("assert should pass");
    }

    #[tokio::test]
    async fn assert_propagates_query_failures() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT * FROM missing_assert_table
----
1

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.assert(&ctx).await,
            "missing_assert_table",
        );
    }

    #[tokio::test]
    async fn assert_reports_row_count_mismatch() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT 1 AS value
----
1
2

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.assert(&ctx).await,
            "expected 2 rows but got 1",
        );
    }

    #[tokio::test]
    async fn assert_reports_column_count_mismatch() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT 1 AS a, 2 AS b
----
1

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.assert(&ctx).await,
            "expected 1 columns but got 2",
        );
    }

    #[tokio::test]
    async fn assert_reports_value_mismatch() {
        let mut benchmark = parse_benchmark(
            r#"
assert I
SELECT 1 AS value
----
2

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.assert(&ctx).await,
            "expected value \"2\" but got value \"1\"",
        );
    }

    // Run tests cover result buffering and physical-plan expectations.

    #[tokio::test]
    async fn run_saves_uppercase_select_results() {
        let mut benchmark = parse_benchmark("run\nSELECT 1 AS value\n")
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_eq!(formatted_last_results(&benchmark), vec![vec!["1"]]);
    }

    #[tokio::test]
    async fn run_saves_with_query_results() {
        let mut benchmark =
            parse_benchmark("run\nWITH t AS (SELECT 3 AS value) SELECT value FROM t\n")
                .await
                .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_eq!(formatted_last_results(&benchmark), vec![vec!["3"]]);
    }

    #[tokio::test]
    async fn run_only_keeps_last_select_or_with_result() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let query_path = write_test_file(
            &temp_dir,
            "queries.sql",
            "SELECT 1 AS value;\nSELECT 2 AS value;\nWITH t AS (SELECT 3 AS value) SELECT value FROM t;\n",
        );
        let benchmark_path = write_test_file(
            &temp_dir,
            "run_file.benchmark",
            &format!("run {}\n", query_path.display()),
        );
        let mut benchmark = parse_benchmark_file(&benchmark_path)
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_eq!(formatted_last_results(&benchmark), vec![vec!["3"]]);
    }

    #[tokio::test]
    async fn run_inline_multi_statement_only_keeps_last_select_or_with_result() {
        let mut benchmark = parse_benchmark(
            "run\nCREATE TABLE t AS SELECT 1 AS value;\nSELECT 2 AS value;\nWITH u AS (SELECT 3 AS value) SELECT value FROM u;\n",
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_eq!(formatted_last_results(&benchmark), vec![vec!["3"]]);
    }

    #[tokio::test]
    async fn run_does_not_save_results_for_non_select_statement() {
        let mut benchmark =
            parse_benchmark("run\nCREATE TABLE run_created AS SELECT 1 AS value;\n")
                .await
                .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert!(
            benchmark
                .last_results
                .as_ref()
                .expect("last results should be set")
                .is_empty()
        );
    }

    #[tokio::test]
    async fn run_propagates_query_failures_when_buffering_results() {
        let mut benchmark = parse_benchmark("run\nSELECT * FROM missing_run_table\n")
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.run(&ctx, true).await,
            "missing_run_table",
        );
    }

    #[tokio::test]
    async fn run_propagates_query_failures_when_streaming_results() {
        let mut benchmark = parse_benchmark("run\nSELECT * FROM missing_stream_table\n")
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.run(&ctx, false).await,
            "missing_stream_table",
        );
    }

    #[tokio::test]
    async fn run_rejects_missing_expect_plan_for_buffered_and_streaming_modes() {
        let ctx = SessionContext::new();
        let benchmark_text = "expect_plan definitely_not_in_plan\nrun\nSELECT 1\n";

        let mut buffered = parse_benchmark(benchmark_text)
            .await
            .expect("benchmark should parse");
        assert_result_error_contains(
            buffered.run(&ctx, true).await,
            "does not contain the expected string 'definitely_not_in_plan'",
        );

        let mut streaming = parse_benchmark(benchmark_text)
            .await
            .expect("benchmark should parse");
        assert_result_error_contains(
            streaming.run(&ctx, false).await,
            "does not contain the expected string 'definitely_not_in_plan'",
        );
    }

    #[tokio::test]
    async fn run_checks_expect_plan_once_per_benchmark() {
        let ctx = SessionContext::new();
        let benchmark_text = "expect_plan PlaceholderRowExec\nrun\nSELECT 1\n";

        let mut benchmark = parse_benchmark(benchmark_text)
            .await
            .expect("benchmark should parse");
        assert!(!benchmark.plans_validated);

        benchmark
            .run(&ctx, false)
            .await
            .expect("first run should accept the matching plan");
        assert!(
            benchmark.plans_validated,
            "the first run should mark the plans as checked"
        );

        benchmark
            .run(&ctx, false)
            .await
            .expect("later runs should not repeat the check");
    }

    #[tokio::test]
    async fn run_accepts_matching_expect_plan_for_buffered_and_streaming_modes() {
        let ctx = SessionContext::new();
        let benchmark_text = "expect_plan PlaceholderRowExec\nrun\nSELECT 1\n";

        let mut buffered = parse_benchmark(benchmark_text)
            .await
            .expect("benchmark should parse");
        buffered
            .run(&ctx, true)
            .await
            .expect("buffered run should accept matching plan");
        assert_eq!(formatted_last_results(&buffered), vec![vec!["1"]]);

        let mut streaming = parse_benchmark(benchmark_text)
            .await
            .expect("benchmark should parse");
        streaming
            .run(&ctx, false)
            .await
            .expect("streaming run should accept matching plan");
    }

    // Verification tests cover result_query and persisted-result comparison paths.

    #[tokio::test]
    async fn verify_without_result_query_returns_ok() {
        let mut benchmark = parse_benchmark("run\nSELECT 1 AS value\n")
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.verify(&ctx).await.expect("verify should pass");
    }

    #[tokio::test]
    async fn verify_errors_when_benchmark_has_not_run() {
        let mut benchmark = parse_benchmark(
            r#"
result_query I
SELECT 1 AS value
----
1

run
SELECT 1;
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.verify(&ctx).await,
            "No results available for verification. Run the benchmark first.",
        );
    }

    #[tokio::test]
    async fn verify_uses_last_results_for_result_file_entries() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path = write_test_file(&temp_dir, "result.csv", "value\n1\n");
        let mut benchmark = parse_benchmark(&format!(
            "result {}\n\nrun\nSELECT 1 AS value\n",
            result_path.display()
        ))
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");
        benchmark.verify(&ctx).await.expect("verify should pass");
    }

    #[tokio::test]
    async fn verify_uses_last_results_for_zero_row_result_file_entries() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path = write_test_file(&temp_dir, "result.csv", "value\n");
        let mut benchmark = parse_benchmark(&format!(
            "result {}\n\nrun\nSELECT 1 AS value WHERE false\n",
            result_path.display()
        ))
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");
        benchmark.verify(&ctx).await.expect("verify should pass");
    }

    #[tokio::test]
    async fn verify_rejects_missing_result_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let missing_path = temp_dir.path().join("missing_result.csv");
        let mut benchmark = parse_benchmark(&format!(
            "result {}\n\nrun\nSELECT 1 AS value\n",
            missing_path.display()
        ))
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_result_error_contains(benchmark.verify(&ctx).await, "missing_result.csv");
    }

    #[tokio::test]
    async fn verify_rejects_malformed_result_file() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let result_path = temp_dir.path().join("malformed_result.csv");
        fs::write(&result_path, [0xff]).expect("failed to write malformed result file");
        let mut benchmark = parse_benchmark(&format!(
            "result {}\n\nrun\nSELECT 1 AS value\n",
            result_path.display()
        ))
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_result_error_contains(benchmark.verify(&ctx).await, "CSV");
    }

    #[tokio::test]
    async fn verify_executes_result_query_instead_of_last_results() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 100 AS value

result_query I
SELECT 1 AS value
----
1
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");
        benchmark.verify(&ctx).await.expect("verify should pass");
    }

    #[tokio::test]
    async fn verify_propagates_result_query_failures() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 1 AS value

result_query I
SELECT * FROM missing_verify_table
----
1
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        assert_result_error_contains(
            benchmark.verify(&ctx).await,
            "missing_verify_table",
        );
    }

    #[tokio::test]
    async fn verify_reports_result_mismatch_context() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 1 AS value

result_query I
SELECT 1 AS value
----
2
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.run(&ctx, true).await.expect("run should succeed");

        let error = benchmark
            .verify(&ctx)
            .await
            .expect_err("verify should fail");
        let message = error.to_string();
        assert!(
            message.contains("row 1, column 1")
                && message.contains("expected value \"2\"")
                && message.contains("got value \"1\""),
            "unexpected error: {message}"
        );
    }

    // Persistence tests cover CSV writing and persist-time error paths.

    #[tokio::test]
    async fn persist_without_result_query_returns_ok() {
        let mut benchmark = parse_benchmark("run\nSELECT 1 AS value\n")
            .await
            .expect("benchmark should parse");
        let ctx = SessionContext::new();

        benchmark.persist(&ctx).await.expect("persist should pass");
    }

    #[tokio::test]
    async fn persist_rejects_result_query_without_file_path() {
        let mut benchmark = parse_benchmark(
            r#"
run
SELECT 1 AS value

result_query I
SELECT 1 AS value
----
1
"#,
        )
        .await
        .expect("benchmark should parse");
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.persist(&ctx).await,
            "Unable to persist results from query",
        );
    }

    #[tokio::test]
    async fn persist_rejects_run_without_saved_result_batches() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let output_path = temp_dir.path().join("persisted");
        let mut benchmark =
            parse_benchmark("run\nCREATE TABLE persist_source AS SELECT 1 AS value;\n")
                .await
                .expect("benchmark should parse");
        benchmark.result_queries.push(BenchmarkQuery {
            path: Some(output_path.to_string_lossy().into_owned()),
            query: String::new(),
            column_count: 1,
            expected_result: vec![],
        });
        let ctx = SessionContext::new();

        assert_result_error_contains(
            benchmark.persist(&ctx).await,
            "Results should be loaded",
        );
    }

    #[tokio::test]
    async fn persist_writes_header_and_pipe_delimited_rows() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let output_path = temp_dir.path().join("persisted");
        let mut benchmark = parse_benchmark("run\nSELECT 1 AS a, 'one' AS b\n")
            .await
            .expect("benchmark should parse");
        benchmark.result_queries.push(BenchmarkQuery {
            path: Some(output_path.to_string_lossy().into_owned()),
            query: String::new(),
            column_count: 2,
            expected_result: vec![],
        });
        let ctx = SessionContext::new();

        benchmark.persist(&ctx).await.expect("persist should pass");

        let contents = read_all_files_in_dir(&output_path);
        assert!(
            contents.contains("a|b\n") && contents.contains("1|one\n"),
            "unexpected persisted contents: {contents:?}"
        );
    }

    #[tokio::test]
    async fn persist_writes_header_for_zero_row_select_results() {
        let temp_dir = tempdir().expect("failed to create benchmark test directory");
        let output_path = temp_dir.path().join("persisted_empty");
        let mut benchmark = parse_benchmark("run\nSELECT 1 AS value WHERE false\n")
            .await
            .expect("benchmark should parse");
        benchmark.result_queries.push(BenchmarkQuery {
            path: Some(output_path.to_string_lossy().into_owned()),
            query: String::new(),
            column_count: 1,
            expected_result: vec![],
        });
        let ctx = SessionContext::new();

        benchmark.persist(&ctx).await.expect("persist should pass");

        let contents = read_all_files_in_dir(&output_path);
        assert!(
            contents.contains("value\n"),
            "unexpected persisted contents: {contents:?}"
        );
    }

    // Path helper tests cover group derivation from benchmark file paths.

    #[test]
    fn parse_group_from_path_returns_group_under_benchmark_directory() {
        let group = parse_group_from_path(
            Path::new("sql_benchmarks/tpch/benchmarks/q01.benchmark"),
            Path::new("sql_benchmarks"),
        );

        assert_eq!(group, "tpch");
    }

    #[test]
    fn parse_group_from_path_matches_benchmark_directory_case_insensitively() {
        let group = parse_group_from_path(
            Path::new("/tmp/SQL_BENCHMARKS/Tpch/benchmarks/q01.benchmark"),
            Path::new("sql_benchmarks"),
        );

        assert_eq!(group, "Tpch");
    }

    #[test]
    fn parse_group_from_path_handles_relative_and_absolute_paths() {
        let relative = parse_group_from_path(
            Path::new("sql_benchmarks/h2o/q01.benchmark"),
            Path::new("sql_benchmarks"),
        );
        let absolute = parse_group_from_path(
            Path::new("/tmp/sql_benchmarks/imdb/q01.benchmark"),
            Path::new("sql_benchmarks"),
        );

        assert_eq!(relative, "h2o");
        assert_eq!(absolute, "imdb");
    }

    #[test]
    fn parse_group_from_path_pins_fallback_for_paths_outside_benchmark_directory() {
        let group = parse_group_from_path(
            Path::new("outside/group/q01.benchmark"),
            Path::new("sql_benchmarks"),
        );

        assert_eq!(group, "outside");
    }

    #[test]
    fn path_ends_with_ignore_ascii_case_matches_component_suffixes() {
        assert!(path_ends_with_ignore_ascii_case(
            Path::new("/tmp/SQL_BENCHMARKS"),
            Path::new("sql_benchmarks")
        ));
        assert!(!path_ends_with_ignore_ascii_case(
            Path::new("/tmp/sql_benchmarks_extra"),
            Path::new("sql_benchmarks")
        ));
    }
}