benchmarks/bench.sh

#!/usr/bin/env bash
# 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.

# This script is meant for developers of DataFusion -- it is runnable
# from the standard DataFusion development environment and uses cargo,
# etc and orchestrates gathering data and run the benchmark binary in
# different configurations.


# Exit on error
set -e

# https://stackoverflow.com/questions/59895/how-do-i-get-the-directory-where-a-bash-script-is-located-from-within-the-script
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

# Execute command and also print it, for debugging purposes
debug_run() {
    set -x
    "$@"
    set +x
}

# Set Defaults
COMMAND=
BENCHMARK=all
DATAFUSION_DIR=${DATAFUSION_DIR:-$SCRIPT_DIR/..}
DATA_DIR=${DATA_DIR:-$SCRIPT_DIR/data}
CARGO_COMMAND=${CARGO_COMMAND:-"cargo run --release"}
SQL_CARGO_COMMAND=${SQL_CARGO_COMMAND:-"cargo bench --bench sql"}
PREFER_HASH_JOIN=${PREFER_HASH_JOIN:-true}
SIMULATE_LATENCY=${SIMULATE_LATENCY:-false}

# Build latency arg based on SIMULATE_LATENCY setting
LATENCY_ARG=""
if [ "$SIMULATE_LATENCY" = "true" ]; then
    LATENCY_ARG="--simulate-latency"
fi

usage() {
    echo "
Orchestrates running benchmarks against DataFusion checkouts

Usage:
$0 data [benchmark]
$0 run [benchmark] [query]
$0 compare <branch1> <branch2>
$0 compare_detail <branch1> <branch2>

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Examples:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Create the datasets for all benchmarks in $DATA_DIR
./bench.sh data

# Run the 'tpch' benchmark on the datafusion checkout in /source/datafusion
DATAFUSION_DIR=/source/datafusion ./bench.sh run tpch

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Commands
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
data:            Generates or downloads data needed for benchmarking
run:             Runs the named benchmark
compare:         Compares fastest results from benchmark runs
compare_detail:  Compares minimum, average (±stddev), and maximum results from benchmark runs

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Benchmarks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# Run all of the following benchmarks
all(default): Data/Run/Compare for all benchmarks

# TPC-H Benchmarks
tpch:                   TPCH inspired benchmark on Scale Factor (SF) 1 (~1GB), single parquet file per table, hash join
tpch_csv:               TPCH inspired benchmark on Scale Factor (SF) 1 (~1GB), single csv file per table, hash join
tpch_mem:               TPCH inspired benchmark on Scale Factor (SF) 1 (~1GB), query from memory
tpch10:                 TPCH inspired benchmark on Scale Factor (SF) 10 (~10GB), single parquet file per table, hash join
tpch_csv10:             TPCH inspired benchmark on Scale Factor (SF) 10 (~10GB), single csv file per table, hash join
tpch_mem10:             TPCH inspired benchmark on Scale Factor (SF) 10 (~10GB), query from memory

# TPC-DS Benchmarks
tpcds:                  TPCDS inspired benchmark on Scale Factor (SF) 1 (~1GB), single parquet file per table, hash join

# Extended TPC-H Benchmarks
sort_tpch:              Benchmark of sorting speed for end-to-end sort queries on TPC-H dataset (SF=1)
sort_tpch10:            Benchmark of sorting speed for end-to-end sort queries on TPC-H dataset (SF=10)
topk_tpch:              Benchmark of top-k (sorting with limit) queries on TPC-H dataset (SF=1)
topk_sorted_tpch:       Benchmark of top-k queries on TPC-H lineitem ordered by l_orderkey (SF=1)
push_down_topk:         Benchmark of ORDER BY ... LIMIT over outer joins on TPC-H dataset (SF=1) — exercises pushing TopK through a join
external_aggr:          External aggregation benchmark on TPC-H dataset (SF=1)
wide_schema:            Small-projection queries on a wide synthetic dataset (1024 cols × 256 files) — measures per-file metadata overhead
                          (runs both 'wide' and 'narrow' subgroups: narrow is an internal baseline; the wide-vs-narrow ratio is the signal)
predicate_eval:         Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an
                          adaptive predicate-ordering system behaves across them (see https://github.com/apache/datafusion/issues/11262)
                          (subgroups via BENCH_SUBGROUP: costsel, cost, selectivity, cardinality, width, scale, neutral, correlation, drift, nulls)
                          (toggle a system under test with its native DATAFUSION_* env var; size data with PRED_ROWS, string width with PRED_FILL)
parquet_row_filter_skip: Per-RG fully-matched RowFilter skip on Parquet (apache/datafusion#23696); clustered string key + low-selectivity
                          range filter + pushdown, so most row groups are fully matched and the per-row RowFilter is skipped on them
                          (subgroups via BENCH_SUBGROUP: skip = clustered key so the skip fires, control = scrambled key so it never fires)
                          (data generated inline by the suite's load SQL; knobs: PRED_ROWS, RG_SIZE)
spill_views:            Sort and GROUP BY queries that spill StringView/BinaryView columns read from Parquet
                          (https://github.com/apache/datafusion/issues/23564); each query sets its own memory limit
                          (subgroups via BENCH_SUBGROUP: repeated = low-cardinality payload, distinct = all-distinct strings)
                          (data generated inline by the suite's load SQL)
null_aware_join:        Null-aware (NOT IN) hash join micro-benchmarks: uncorrelated, non-equality-correlated and equality-correlated
                          NOT IN across NULL fractions, to measure the per-pair join-filter work the correlated cases do
                          (data generated inline by the suite's load SQL from range(); knobs: NAJ_ROWS, NAJ_LARGE_ROWS)
projection_subquery:    IN / NOT IN / EXISTS subqueries in the SELECT list (see https://github.com/apache/datafusion/issues/25341); each query projects the
                          boolean subquery result and aggregates it, so the cost is the decorrelation plan and not the output size
                          (q07 correlates on '<' instead of '=', so it keeps the nested-loop plan and acts as the control)
                          (data generated inline by the suite's load SQL; knob: PSQ_ROWS)

# ClickBench Benchmarks
clickbench_1:           ClickBench queries against a single parquet file
clickbench_partitioned: ClickBench queries against partitioned (100 files) parquet
clickbench_pushdown:    ClickBench queries against partitioned (100 files) parquet w/ filter_pushdown enabled
clickbench_extended:    ClickBench \"inspired\" queries against a single parquet (DataFusion specific)

# Sort Pushdown Benchmarks
sort_pushdown:          Sort pushdown baseline (no WITH ORDER) on TPC-H data (SF=1)
sort_pushdown_sorted:   Sort pushdown with WITH ORDER — tests sort elimination on non-overlapping files
sort_pushdown_inexact:  Sort pushdown Inexact path (--sorted DESC) — multi-file with scrambled RGs, tests reverse scan + RG reorder
sort_pushdown_inexact_unsorted: Sort pushdown Inexact path (no WITH ORDER) — same data, tests Unsupported path + RG reorder
sort_pushdown_inexact_overlap: Sort pushdown Inexact path — multi-file scrambled RGs (streaming data scenario)

# Sorted Data Benchmarks (ORDER BY Optimization)
clickbench_sorted:     ClickBench queries on pre-sorted data using prefer_existing_sort (tests sort elimination optimization)

# H2O.ai Benchmarks (Group By, Join, Window)
h2o_small:                      h2oai benchmark with small dataset (1e7 rows) for groupby,  default file format is csv
h2o_medium:                     h2oai benchmark with medium dataset (1e8 rows) for groupby, default file format is csv
h2o_big:                        h2oai benchmark with large dataset (1e9 rows) for groupby,  default file format is csv
h2o_small_join:                 h2oai benchmark with small dataset (1e7 rows) for join,  default file format is csv
h2o_medium_join:                h2oai benchmark with medium dataset (1e8 rows) for join, default file format is csv
h2o_big_join:                   h2oai benchmark with large dataset (1e9 rows) for join,  default file format is csv
h2o_small_window:               Extended h2oai benchmark with small dataset (1e7 rows) for window,  default file format is csv
h2o_medium_window:              Extended h2oai benchmark with medium dataset (1e8 rows) for window, default file format is csv
h2o_big_window:                 Extended h2oai benchmark with large dataset (1e9 rows) for window,  default file format is csv
h2o_small_window_sorted:        Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows),  default file format is csv
h2o_medium_window_sorted:       Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), default file format is csv
h2o_big_window_sorted:          Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows),  default file format is csv
h2o_small_window_sorted_parquet:  Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows),  source file format is parquet
h2o_medium_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), source file format is parquet
h2o_big_window_sorted_parquet:    Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows),  source file format is parquet
h2o_small_parquet:              h2oai benchmark with small dataset (1e7 rows) for groupby,  file format is parquet
h2o_medium_parquet:             h2oai benchmark with medium dataset (1e8 rows) for groupby, file format is parquet
h2o_big_parquet:                h2oai benchmark with large dataset (1e9 rows) for groupby,  file format is parquet
h2o_small_join_parquet:         h2oai benchmark with small dataset (1e7 rows) for join,  file format is parquet
h2o_medium_join_parquet:        h2oai benchmark with medium dataset (1e8 rows) for join, file format is parquet
h2o_big_join_parquet:           h2oai benchmark with large dataset (1e9 rows) for join,  file format is parquet
h2o_small_window_parquet:       Extended h2oai benchmark with small dataset (1e7 rows) for window,  file format is parquet
h2o_medium_window_parquet:      Extended h2oai benchmark with medium dataset (1e8 rows) for window, file format is parquet
h2o_big_window_parquet:         Extended h2oai benchmark with large dataset (1e9 rows) for window,  file format is parquet

# Join Order Benchmark (IMDB)
imdb:                   Join Order Benchmark (JOB) using the IMDB dataset converted to parquet

# Micro-Benchmarks (specific operators and features)
cancellation:           How long cancelling a query takes
asof_join:              ASOF join workloads varying size, ordering, grouping, match direction, and payload width
nlj:                    Benchmark for simple nested loop joins, testing various join scenarios
hj:                     Benchmark for simple hash joins, testing various join scenarios
smj:                    Benchmark for simple sort merge joins, testing various join scenarios
dict:                   Benchmark for dictionary-encoded group-by scenarios
array_agg_distinct:     1000K-group, two-row-per-group array_agg(DISTINCT) benchmark
compile_profile:        Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size


━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Supported Configuration (Environment Variables)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
DATA_DIR            directory to store datasets
CARGO_COMMAND       command that runs the benchmark binary
DATAFUSION_DIR      directory to use (default $DATAFUSION_DIR)
RESULTS_NAME        folder where the benchmark files are stored
PREFER_HASH_JOIN    Prefer hash join algorithm (default true)
SIMULATE_LATENCY    Serve local files as if they were remote object storage: per-request
                    latency, per-connection bandwidth and paginated LIST (default false)
DATAFUSION_*        Set the given datafusion configuration
"
    exit 1
}

# https://stackoverflow.com/questions/192249/how-do-i-parse-command-line-arguments-in-bash
POSITIONAL_ARGS=()

while [[ $# -gt 0 ]]; do
    case $1 in
        # -e|--extension)
        #   EXTENSION="$2"
        #   shift # past argument
        #   shift # past value
        #   ;;
        -h|--help)
            shift # past argument
            usage
            ;;
        -*)
            echo "Unknown option $1"
            exit 1
            ;;
        *)
            POSITIONAL_ARGS+=("$1") # save positional arg
            shift # past argument
            ;;
    esac
done

set -- "${POSITIONAL_ARGS[@]}" # restore positional parameters
COMMAND=${1:-"${COMMAND}"}
ARG2=$2
ARG3=$3

# Do what is requested
main() {
    # Command Dispatch
    case "$COMMAND" in
        data)
            BENCHMARK=${ARG2:-"${BENCHMARK}"}
            echo "***************************"
            echo "DataFusion Benchmark Runner and Data Generator"
            echo "COMMAND: ${COMMAND}"
            echo "BENCHMARK: ${BENCHMARK}"
            echo "DATA_DIR: ${DATA_DIR}"
            echo "CARGO_COMMAND: ${CARGO_COMMAND}"
            echo "PREFER_HASH_JOIN: ${PREFER_HASH_JOIN}"
            echo "***************************"
            case "$BENCHMARK" in
                all)
                    data_tpch "1" "parquet"
                    data_tpch "10" "parquet"
                    data_h2o "SMALL"
                    data_h2o "MEDIUM"
                    data_h2o "BIG"
                    data_h2o_join "SMALL"
                    data_h2o_join "MEDIUM"
                    data_h2o_join "BIG"
                    data_clickbench_1
                    data_clickbench_partitioned
                    data_imdb
                    data_asof_join
                    # nlj uses range() function, no data generation needed
                    ;;
                tpch)
                    data_tpch "1" "parquet"
                    ;;
                tpch_mem)
                    data_tpch "1" "parquet"
                    ;;
                tpch_csv)
                    data_tpch "1" "csv"
                    ;;
                tpch10)
                    data_tpch "10" "parquet"
                    ;;
                tpch_mem10)
                    data_tpch "10" "parquet"
                    ;;
                tpch_csv10)
                    data_tpch "10" "csv"
                    ;;
                wide_schema)
                    data_wide_schema
                    ;;
                predicate_eval)
                    # Data is generated inline by the suite's load SQL.
                    echo "predicate_eval: no external data to generate"
                    ;;
                parquet_row_filter_skip)
                    # Data is generated inline by the suite's load SQL (COPY).
                    echo "parquet_row_filter_skip: no external data to generate"
                    ;;
                spill_views)
                    # Data is generated inline by the suite's load SQL (COPY).
                    echo "spill_views: no external data to generate"
                    ;;
                null_aware_join)
                    # Data is generated inline by the suite's load SQL from range().
                    echo "null_aware_join: no external data to generate"
                    ;;
                projection_subquery)
                    # Data is generated inline by the suite's load SQL.
                    echo "projection_subquery: no external data to generate"
                    ;;
                asof_join)
                    data_asof_join
                    ;;
                tpcds)
                    data_tpcds
                    ;;
                clickbench_1)
                    data_clickbench_1
                    ;;
                clickbench_partitioned)
                    data_clickbench_partitioned
                    ;;
                clickbench_pushdown)
                    data_clickbench_partitioned # same data as clickbench_partitioned
                    ;;
                clickbench_extended)
                    data_clickbench_1
                    ;;
                imdb)
                    data_imdb
                    ;;
                h2o_small)
                    data_h2o "SMALL" "CSV"
                    ;;
                h2o_medium)
                    data_h2o "MEDIUM" "CSV"
                    ;;
                h2o_big)
                    data_h2o "BIG" "CSV"
                    ;;
                h2o_small_join)
                    data_h2o_join "SMALL" "CSV"
                    ;;
                h2o_medium_join)
                    data_h2o_join "MEDIUM" "CSV"
                    ;;
                h2o_big_join)
                    data_h2o_join "BIG" "CSV"
                    ;;
                # h2o window benchmark uses the same data as the h2o join
                h2o_small_window)
                    data_h2o_join "SMALL" "CSV"
                    ;;
                h2o_medium_window)
                    data_h2o_join "MEDIUM" "CSV"
                    ;;
                h2o_big_window)
                    data_h2o_join "BIG" "CSV"
                    ;;
                # the sorted window subgroup derives its data from the same
                # source, then sorts it inside its load SQL
                h2o_small_window_sorted)
                    data_h2o_join "SMALL" "CSV"
                    ;;
                h2o_medium_window_sorted)
                    data_h2o_join "MEDIUM" "CSV"
                    ;;
                h2o_big_window_sorted)
                    data_h2o_join "BIG" "CSV"
                    ;;
                h2o_small_window_sorted_parquet)
                    data_h2o_join "SMALL" "PARQUET"
                    ;;
                h2o_medium_window_sorted_parquet)
                    data_h2o_join "MEDIUM" "PARQUET"
                    ;;
                h2o_big_window_sorted_parquet)
                    data_h2o_join "BIG" "PARQUET"
                    ;;
                h2o_small_parquet)
                    data_h2o "SMALL" "PARQUET"
                    ;;
                h2o_medium_parquet)
                    data_h2o "MEDIUM" "PARQUET"
                    ;;
                h2o_big_parquet)
                    data_h2o "BIG" "PARQUET"
                    ;;
                h2o_small_join_parquet)
                    data_h2o_join "SMALL" "PARQUET"
                    ;;
                h2o_medium_join_parquet)
                    data_h2o_join "MEDIUM" "PARQUET"
                    ;;
                h2o_big_join_parquet)
                    data_h2o_join "BIG" "PARQUET"
                    ;;
                # h2o window benchmark uses the same data as the h2o join
                h2o_small_window_parquet)
                    data_h2o_join "SMALL" "PARQUET"
                    ;;
                h2o_medium_window_parquet)
                    data_h2o_join "MEDIUM" "PARQUET"
                    ;;
                h2o_big_window_parquet)
                    data_h2o_join "BIG" "PARQUET"
                    ;;
                external_aggr)
                    # same data as for tpch
                    data_tpch "1" "parquet"
                    ;;
                sort_pushdown|sort_pushdown_sorted)
                    data_sort_pushdown
                    ;;
                sort_pushdown_inexact|sort_pushdown_inexact_unsorted|sort_pushdown_inexact_overlap)
                    data_sort_pushdown_inexact
                    ;;
                sort_tpch)
                    # same data as for tpch
                    data_tpch "1" "parquet"
                    ;;
                sort_tpch10)
                    # same data as for tpch10
                    data_tpch "10" "parquet"
                    ;;
                topk_tpch|topk_sorted_tpch)
                    # same data as for tpch
                    data_tpch "1" "parquet"
                    ;;
                push_down_topk)
                    # same data as for tpch
                    data_tpch "1" "parquet"
                    ;;
                nlj)
                    # nlj uses range() function, no data generation needed
                    echo "NLJ benchmark does not require data generation"
                    ;;
                hj)
                    data_tpch "10" "parquet"
                    ;;
                smj)
                    # smj uses range() function, no data generation needed
                    echo "SMJ benchmark does not require data generation"
                    ;;
                dict)
                    # dict generates in-memory data, no data generation needed
                    echo "DICT benchmark does not require data generation"
                    ;;
                compile_profile)
                    data_tpch "1" "parquet"
                    ;;
                clickbench_sorted)
                    clickbench_sorted
                    ;;
                *)
                    echo "Error: unknown benchmark '$BENCHMARK' for data generation"
                    usage
                    ;;
            esac
            ;;
        run)
            # Parse positional parameters
            BENCHMARK=${ARG2:-"${BENCHMARK}"}
            EXTRA_ARGS=("${POSITIONAL_ARGS[@]:2}")
            PROFILE_ARGS=()
            QUERY=""
            QUERY_ARG=""
            if [ "$BENCHMARK" = "compile_profile" ]; then
                PROFILE_ARGS=("${EXTRA_ARGS[@]}")
            else
                QUERY=${EXTRA_ARGS[0]}
                if [ -n "$QUERY" ]; then
                    QUERY_ARG="--query ${QUERY}"
                fi
            fi
            BRANCH_NAME=$(cd "${DATAFUSION_DIR}" && git rev-parse --abbrev-ref HEAD)
            BRANCH_NAME=${BRANCH_NAME//\//_} # mind blowing syntax to replace / with _
            RESULTS_NAME=${RESULTS_NAME:-"${BRANCH_NAME}"}
            RESULTS_DIR=${RESULTS_DIR:-"$SCRIPT_DIR/results/$RESULTS_NAME"}

            echo "***************************"
            echo "DataFusion Benchmark Script"
            echo "COMMAND: ${COMMAND}"
            echo "BENCHMARK: ${BENCHMARK}"
            if [ "$BENCHMARK" = "compile_profile" ]; then
                echo "PROFILES: ${PROFILE_ARGS[*]:-All}"
            else
                echo "QUERY: ${QUERY:-All}"
            fi
            echo "DATAFUSION_DIR: ${DATAFUSION_DIR}"
            echo "BRANCH_NAME: ${BRANCH_NAME}"
            echo "DATA_DIR: ${DATA_DIR}"
            echo "RESULTS_DIR: ${RESULTS_DIR}"
            echo "CARGO_COMMAND: ${CARGO_COMMAND}"
            echo "PREFER_HASH_JOIN: ${PREFER_HASH_JOIN}"
            echo "SIMULATE_LATENCY: ${SIMULATE_LATENCY}"
            echo "***************************"

            # navigate to the appropriate directory
            pushd "${DATAFUSION_DIR}/benchmarks" > /dev/null
            mkdir -p "${RESULTS_DIR}"
            mkdir -p "${DATA_DIR}"
            case "$BENCHMARK" in
                all)
                    run_tpch "1" "parquet"
                    run_tpch "1" "csv"
                    run_tpch_mem "1"
                    run_tpch "10" "parquet"
                    run_tpch "10" "csv"
                    run_tpch_mem "10"
                    run_cancellation
                    run_clickbench_1
                    run_clickbench_partitioned
                    run_clickbench_pushdown
                    run_clickbench_extended
                    run_h2o "SMALL" "PARQUET" "groupby"
                    run_h2o "MEDIUM" "PARQUET" "groupby"
                    run_h2o "BIG" "PARQUET" "groupby"
                    run_h2o_join "SMALL" "PARQUET" "join"
                    run_h2o_join "MEDIUM" "PARQUET" "join"
                    run_h2o_join "BIG" "PARQUET" "join"
                    run_imdb
                    run_external_aggr
                    run_asof_join
                    run_nlj
                    run_hj
                    run_tpcds
                    run_smj
                    run_dict 
                    run_null_aware_join
                    ;;
                tpch)
                    run_tpch "1" "parquet"
                    ;;
                tpch_csv)
                    run_tpch "1" "csv"
                    ;;
                tpch_mem)
                    run_tpch_mem "1"
                    ;;
                tpch10)
                    run_tpch "10" "parquet"
                    ;;
                tpch_csv10)
                    run_tpch "10" "csv"
                    ;;
                tpch_mem10)
                    run_tpch_mem "10"
                    ;;
                wide_schema)
                    run_wide_schema
                    ;;
                predicate_eval)
                    run_predicate_eval
                    ;;
                parquet_row_filter_skip)
                    run_parquet_row_filter_skip
                    ;;
                spill_views)
                    run_spill_views
                    ;;
                null_aware_join)
                    run_null_aware_join
                    ;;
                projection_subquery)
                    run_projection_subquery
                    ;;
                asof_join)
                    run_asof_join
                    ;;
                tpcds)
                    run_tpcds
                    ;;
                cancellation)
                    run_cancellation
                    ;;
                clickbench_1)
                    run_clickbench_1
                    ;;
                clickbench_partitioned)
                    run_clickbench_partitioned
                    ;;
                clickbench_pushdown)
                    run_clickbench_pushdown
                    ;;
                clickbench_extended)
                    run_clickbench_extended
                    ;;
                imdb)
                    run_imdb
                    ;;
                h2o_small)
                    run_h2o "SMALL" "CSV" "groupby"
                    ;;
                h2o_medium)
                    run_h2o "MEDIUM" "CSV" "groupby"
                    ;;
                h2o_big)
                    run_h2o "BIG" "CSV" "groupby"
                    ;;
                h2o_small_join)
                    run_h2o_join "SMALL" "CSV" "join"
                    ;;
                h2o_medium_join)
                    run_h2o_join "MEDIUM" "CSV" "join"
                    ;;
                h2o_big_join)
                    run_h2o_join "BIG" "CSV" "join"
                    ;;
                h2o_small_window)
                    run_h2o_window "SMALL" "CSV" "window"
                    ;;
                h2o_medium_window)
                    run_h2o_window "MEDIUM" "CSV" "window"
                    ;;
                h2o_big_window)
                    run_h2o_window "BIG" "CSV" "window"
                    ;;
                h2o_small_window_sorted)
                    run_h2o_window_sorted "small" "csv"
                    ;;
                h2o_medium_window_sorted)
                    run_h2o_window_sorted "medium" "csv"
                    ;;
                h2o_big_window_sorted)
                    run_h2o_window_sorted "big" "csv"
                    ;;
                h2o_small_window_sorted_parquet)
                    run_h2o_window_sorted "small" "parquet"
                    ;;
                h2o_medium_window_sorted_parquet)
                    run_h2o_window_sorted "medium" "parquet"
                    ;;
                h2o_big_window_sorted_parquet)
                    run_h2o_window_sorted "big" "parquet"
                    ;;
                h2o_small_parquet)
                    run_h2o "SMALL" "PARQUET"
                    ;;
                h2o_medium_parquet)
                    run_h2o "MEDIUM" "PARQUET"
                    ;;
                h2o_big_parquet)
                    run_h2o "BIG" "PARQUET"
                    ;;
                h2o_small_join_parquet)
                    run_h2o_join "SMALL" "PARQUET"
                    ;;
                h2o_medium_join_parquet)
                    run_h2o_join "MEDIUM" "PARQUET"
                    ;;
                h2o_big_join_parquet)
                    run_h2o_join "BIG" "PARQUET"
                    ;;
                # h2o window benchmark uses the same data as the h2o join
                h2o_small_window_parquet)
                    run_h2o_window "SMALL" "PARQUET"
                    ;;
                h2o_medium_window_parquet)
                    run_h2o_window "MEDIUM" "PARQUET"
                    ;;
                h2o_big_window_parquet)
                    run_h2o_window "BIG" "PARQUET"
                    ;;
                external_aggr)
                    run_external_aggr
                    ;;
                sort_pushdown)
                    run_sort_pushdown
                    ;;
                sort_pushdown_sorted)
                    run_sort_pushdown_sorted
                    ;;
                sort_pushdown_inexact)
                    run_sort_pushdown_inexact
                    ;;
                sort_pushdown_inexact_unsorted)
                    run_sort_pushdown_inexact_unsorted
                    ;;
                sort_pushdown_inexact_overlap)
                    run_sort_pushdown_inexact_overlap
                    ;;
                sort_tpch)
                    run_sort_tpch "1"
                    ;;
                sort_tpch10)
                    run_sort_tpch "10"
                    ;;
                topk_tpch)
                    run_topk_tpch
                    ;;
                topk_sorted_tpch)
                    run_topk_sorted_tpch
                    ;;
                push_down_topk)
                    run_push_down_topk
                    ;;
                nlj)
                    run_nlj
                    ;;
                hj)
                    run_hj
                    ;;
                smj)
                    run_smj
                    ;;
                dict)
                    run_dict
                    ;;
                array_agg_distinct)
                    run_array_agg_distinct
                    ;;
                compile_profile)
                    run_compile_profile "${PROFILE_ARGS[@]}"
                    ;;
                clickbench_sorted)
                    run_clickbench_sorted
                    ;;
                *)
                    echo "Error: unknown benchmark '$BENCHMARK' for run"
                    usage
                    ;;
            esac
            popd > /dev/null
            echo "Done"
            ;;
        compare)
            compare_benchmarks "$ARG2" "$ARG3"
            ;;
        compare_detail)
            compare_benchmarks "$ARG2" "$ARG3" "--detailed"
            ;;
        "")
            usage
            ;;
        *)
            echo "Error: unknown command: $COMMAND"
            usage
            ;;
    esac
}



# Creates TPCH data at a certain scale factor, if it doesn't already
# exist
#
# call like: data_tpch($scale_factor, format)
#
# Creates data in $DATA_DIR/tpch_sf1 for scale factor 1
# Creates data in $DATA_DIR/tpch_sf10 for scale factor 10
# etc
data_tpch() {
    SCALE_FACTOR=$1
    if [ -z "$SCALE_FACTOR" ] ; then
        echo "Internal error: Scale factor not specified"
        exit 1
    fi
    FORMAT=$2
    if [ -z "$FORMAT" ] ; then
        echo "Internal error: Format not specified"
        exit 1
    fi

    TPCH_DIR="${DATA_DIR}/tpch_sf${SCALE_FACTOR}"
    echo "Creating tpch $FORMAT dataset at Scale Factor ${SCALE_FACTOR} in ${TPCH_DIR}..."

    # Ensure the target data directory exists
    mkdir -p "${TPCH_DIR}"

    # check if tpchgen-cli is installed
    if ! command -v tpchgen-cli &> /dev/null
    then
        echo "tpchgen-cli could not be found, please install it via 'cargo install tpchgen-cli'"
        exit 1
    fi

    # Copy expected answers into the ./data/answers directory if it does not already exist
    FILE="${TPCH_DIR}/answers/q1.out"
    if test -f "${FILE}"; then
        echo " Expected answers exist (${FILE} exists)."
    else
        echo " Copying answers to ${TPCH_DIR}/answers"
        mkdir -p "${TPCH_DIR}/answers"
        docker run -v "${TPCH_DIR}":/data -it --entrypoint /bin/bash --rm ghcr.io/scalytics/tpch-docker:main  -c "cp -f /opt/tpch/2.18.0_rc2/dbgen/answers/* /data/answers/"
    fi

    if [ "$FORMAT" = "parquet" ]; then
      # Create 'parquet' files, one directory per file
      FILE="${TPCH_DIR}/supplier"
      if test -d "${FILE}"; then
          echo " parquet files exist ($FILE exists)."
      else
          echo " creating parquet files using tpchgen-cli ..."
          tpchgen-cli --scale-factor "${SCALE_FACTOR}" --format parquet --parquet-compression='ZSTD(1)' --parts=1 --output-dir "${TPCH_DIR}"
      fi
      return
    fi

    # Create 'csv' files, one directory per file
    if [ "$FORMAT" = "csv" ]; then
      FILE="${TPCH_DIR}/csv/supplier"
      if test -d "${FILE}"; then
          echo " csv files exist ($FILE exists)."
      else
          echo " creating csv files using tpchgen-cli binary ..."
          tpchgen-cli --scale-factor "${SCALE_FACTOR}" --format csv --parts=1 --output-dir "${TPCH_DIR}/csv"
      fi
      return
    fi

    echo "Error: unknown format '$FORMAT' for tpch data generation, expected 'parquet' or 'csv'"
    exit 1
}

# Downloads TPC-DS data
data_tpcds() {
    TPCDS_DIR="${DATA_DIR}/tpcds_sf1"

    # Check if `web_site.parquet` exists in the TPCDS data directory to verify data presence
    echo "Checking TPC-DS data directory: ${TPCDS_DIR}"
    if [ ! -f "${TPCDS_DIR}/web_site.parquet" ]; then
        mkdir -p "${TPCDS_DIR}"
        # Download the DataFusion benchmarks repository zip if it is not already downloaded
        if [ ! -f "${DATA_DIR}/datafusion-benchmarks.zip" ]; then
          echo "Downloading DataFusion benchmarks repository zip to: ${DATA_DIR}/datafusion-benchmarks.zip"
          wget --timeout=30 --tries=3 -O "${DATA_DIR}/datafusion-benchmarks.zip" https://github.com/apache/datafusion-benchmarks/archive/refs/heads/main.zip
        fi
        echo "Extracting TPC-DS parquet data to ${TPCDS_DIR}..."
        unzip -o -j -d "${TPCDS_DIR}" "${DATA_DIR}/datafusion-benchmarks.zip" datafusion-benchmarks-main/tpcds/data/sf1/*
        echo "TPC-DS data extracted."
    fi
    echo "Done."
}

# Path of the results JSON the Criterion SQL harness ($SQL_CARGO_COMMAND)
# writes, passed to it as BENCH_RESULTS_FILE. It has the same format as the
# dfbench results files and holds each query's peak memory pool reservation
# (recorded only under a memory limit). It has no timings: Criterion keeps
# those under target/criterion. It goes in a subdirectory so that `compare`,
# which reads ${RESULTS_DIR}/*.json as timing results, does not read it.
sql_results_file() {
    mkdir -p "${RESULTS_DIR}/criterion"
    echo "${RESULTS_DIR}/criterion/$1.json"
}

# Runs the tpch benchmark
run_tpch() {
    SCALE_FACTOR=$1
    if [ -z "$SCALE_FACTOR" ] ; then
        echo "Internal error: Scale factor not specified"
        exit 1
    fi
    FORMAT=$2
    echo "Running tpch benchmark..."

    debug_run env BENCH_NAME=tpch \
      BENCH_RESULTS_FILE="$(sql_results_file tpch_sf${SCALE_FACTOR}_${FORMAT})" \
      BENCH_SIZE="${SCALE_FACTOR}" \
      DATA_DIR="${DATA_DIR}" \
      PREFER_HASH_JOIN="${PREFER_HASH_JOIN}" \
      TPCH_FILE_TYPE="${FORMAT}" \
      SIMULATE_LATENCY="${SIMULATE_LATENCY}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Synthesizes two parquet datasets used to measure per-file metadata
# overhead of a wide schema:
#
#   - data/wide_schema/wide/    1024-col events × 256 files (~225 MB)
#   - data/wide_schema/narrow/    8-col events × 256 files (~110 MB)
#
# Both share row count, file count, and per-file row-group shape; only
# schema width differs. No external data source required — gen_wide_data
# synthesizes everything from scratch in ~60 s.
data_wide_schema() {
    NUM_FILES=256
    ROWS_PER_FILE=50000
    WIDTH_FACTOR=128

    DST_DIR="${DATA_DIR}/wide_schema"
    WIDE_DIR="${DST_DIR}/wide"
    NARROW_DIR="${DST_DIR}/narrow"

    if [ -d "${WIDE_DIR}" ] && [ "$(ls -A "${WIDE_DIR}" 2>/dev/null | wc -l)" -ge ${NUM_FILES} ]; then
        echo " wide parquet exists (${WIDE_DIR})."
    else
        mkdir -p "${WIDE_DIR}"
        echo " synthesizing wide -> ${WIDE_DIR} (factor ${WIDTH_FACTOR}, ${NUM_FILES} files × ${ROWS_PER_FILE} rows) ..."
        debug_run $CARGO_COMMAND --bin gen_wide_data -- \
            --dst-dir "${WIDE_DIR}" \
            --width-factor ${WIDTH_FACTOR} \
            --num-files ${NUM_FILES} \
            --rows-per-file ${ROWS_PER_FILE}
    fi

    if [ -d "${NARROW_DIR}" ] && [ "$(ls -A "${NARROW_DIR}" 2>/dev/null | wc -l)" -ge ${NUM_FILES} ]; then
        echo " narrow parquet exists (${NARROW_DIR})."
    else
        mkdir -p "${NARROW_DIR}"
        echo " synthesizing narrow -> ${NARROW_DIR} (8 base cols, ${NUM_FILES} files × ${ROWS_PER_FILE} rows) ..."
        debug_run $CARGO_COMMAND --bin gen_wide_data -- \
            --dst-dir "${NARROW_DIR}" \
            --width-factor 1 \
            --num-files ${NUM_FILES} \
            --rows-per-file ${ROWS_PER_FILE}
    fi
}

# Runs the wide_schema benchmark. Each query has a `subgroup`
# directive that picks up BENCH_SUBGROUP, so we invoke the framework
# twice — once with subgroup=wide (the actual workload) and once with
# subgroup=narrow (the baseline). The wide-only queries (Q02/Q08/Q11/Q12)
# hardcode `subgroup wide`, so they're skipped on the narrow pass.
run_wide_schema() {
    echo "Running wide_schema benchmark (wide subgroup)..."
    debug_run env BENCH_NAME=wide_schema BENCH_SUBGROUP=wide \
      BENCH_RESULTS_FILE="$(sql_results_file wide_schema_wide)" \
      DATA_DIR="${DATA_DIR}" \
      SIMULATE_LATENCY="${SIMULATE_LATENCY}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"

    echo "Running wide_schema benchmark (narrow baseline subgroup)..."
    debug_run env BENCH_NAME=wide_schema BENCH_SUBGROUP=narrow \
      BENCH_RESULTS_FILE="$(sql_results_file wide_schema_narrow)" \
      DATA_DIR="${DATA_DIR}" \
      SIMULATE_LATENCY="${SIMULATE_LATENCY}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the push_down_topk benchmark (ORDER BY ... LIMIT over outer joins).
# Reuses the TPC-H parquet data, so it needs `./bench.sh data tpch` (or
# `data push_down_topk`) first.
run_push_down_topk() {
    echo "Running push_down_topk benchmark..."

    debug_run env BENCH_NAME=push_down_topk \
      BENCH_RESULTS_FILE="$(sql_results_file push_down_topk)" \
      BENCH_SIZE="1" \
      DATA_DIR="${DATA_DIR}" \
      SIMULATE_LATENCY="${SIMULATE_LATENCY}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the predicate_eval benchmark suite: conjunctive (AND) filter-evaluation
# micro-benchmarks where each subgroup is a different predicate pattern, used to
# test how an adaptive predicate-ordering system behaves across them (see
# https://github.com/apache/datafusion/issues/11262). Data is generated inline
# by the suite's load SQL, so there is no data step (drift q82 and q83 share 16
# small Parquet files written into sql_benchmarks/predicate_eval/scratch/, which
# is gitignored; q82 reads them with target_partitions=1, so the selectivity flip
# lands halfway through one stream, and q83 with 16, so each stream gets one whole
# file and a fixed profile).
#
# By default the suite measures DataFusion's built-in left-deep AND short-circuit
# and sets no engine config of its own. To evaluate a system under test, export
# its native DATAFUSION_* config before invoking bench.sh -- the harness reads
# SessionConfig::from_env, and that environment is inherited here, e.g.
#   DATAFUSION_EXECUTION_ADAPTIVE_FILTER_REORDERING=true ./bench.sh run predicate_eval
# Suite-specific knobs (string-substituted into the load SQL, not engine config):
#   BENCH_SUBGROUP   run one subgroup (costsel, cost, selectivity, cardinality,
#                    width, scale, neutral, correlation, drift, nulls)
#   PRED_ROWS        synthetic row count (default 1_000_000; the scale subgroup
#                    overrides this per query)
#   PRED_FILL        filler chars per marker = string-column width knob
run_predicate_eval() {
    echo "Running predicate_eval benchmark (subgroup=${BENCH_SUBGROUP:-all}, rows=${PRED_ROWS:-1000000})..."
    debug_run env BENCH_NAME=predicate_eval \
      BENCH_RESULTS_FILE="$(sql_results_file predicate_eval)" \
      ${BENCH_SUBGROUP:+BENCH_SUBGROUP="${BENCH_SUBGROUP}"} \
      PRED_ROWS="${PRED_ROWS:-1000000}" \
      ${PRED_FILL:+PRED_FILL="${PRED_FILL}"} \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the parquet_row_filter_skip suite: the load SQL COPYs a Parquet file
# inline (fixed-width string key, ordered so each row group holds a disjoint
# sorted range) and the query applies a low-selectivity range filter, so all
# but the first row group is fully matched by statistics and the per-row
# RowFilter is skipped on them (apache/datafusion#23696). The control subgroup
# scrambles the key so no row group is ever fully matched, measuring the
# overhead of the check when it cannot fire. Data is inline, so no data step.
# Knobs (string-substituted into the load SQL, not engine config):
#   BENCH_SUBGROUP  run one subgroup (skip, control)
#   PRED_ROWS       synthetic row count      (default 10_000_000)
#   RG_SIZE         parquet row-group size   (default 1_000_000)
run_parquet_row_filter_skip() {
    echo "Running parquet_row_filter_skip benchmark (subgroup=${BENCH_SUBGROUP:-all}, rows=${PRED_ROWS:-10000000}, rg_size=${RG_SIZE:-1000000})..."
    debug_run env BENCH_NAME=parquet_row_filter_skip \
      BENCH_RESULTS_FILE="$(sql_results_file parquet_row_filter_skip)" \
      ${BENCH_SUBGROUP:+BENCH_SUBGROUP="${BENCH_SUBGROUP}"} \
      PRED_ROWS="${PRED_ROWS:-10000000}" \
      RG_SIZE="${RG_SIZE:-1000000}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the null_aware_join suite: NOT IN (null-aware) hash joins. The load SQL
# builds every table inline from range(), so there is no data step.
#
# Q01-Q03 are uncorrelated NOT IN and are linear in the table size; they are the
# regression guard for the plain null-aware path. Q04-Q08 are correlated, where
# the correlation predicate stays behind as a join filter that the join applies
# per candidate (build row x probe row) pair while deciding which rows are
# UNKNOWN; with no equality correlation there are no scope keys to narrow those
# pairs, so Q05-Q07 scale with the NULL count times the opposite table's size.
#
#   NAJ_ROWS         rows per table for the correlated queries (default 10_000)
#   NAJ_LARGE_ROWS   rows per table for the uncorrelated queries (default 1_000_000)
run_null_aware_join() {
    echo "Running null_aware_join benchmark (rows=${NAJ_ROWS:-10000}, large_rows=${NAJ_LARGE_ROWS:-1000000})..."
    debug_run env BENCH_NAME=null_aware_join \
      BENCH_RESULTS_FILE="$(sql_results_file null_aware_join)" \
      NAJ_ROWS="${NAJ_ROWS:-10000}" \
      NAJ_LARGE_ROWS="${NAJ_LARGE_ROWS:-1000000}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the projection_subquery suite: IN / NOT IN / EXISTS subqueries that sit
# in the SELECT list instead of a filter (see
# https://github.com/apache/datafusion/issues/25341). The load SQL builds the
# two tables inline, so there is no data step. Each query projects the boolean
# subquery result and aggregates it, so the measured cost is the decorrelation
# plan and not the size of the output. Query 07 correlates on '<' instead of
# '=', so it keeps the nested-loop plan and acts as the control.
# Knob (string-substituted into the load SQL, not engine config):
#   PSQ_ROWS  rows in each of the two tables (default 30_000; the checked-in
#             result files hold the counts for that value)
run_projection_subquery() {
    echo "Running projection_subquery benchmark (rows=${PSQ_ROWS:-30000})..."
    debug_run env BENCH_NAME=projection_subquery \
      BENCH_RESULTS_FILE="$(sql_results_file projection_subquery)" \
      PSQ_ROWS="${PSQ_ROWS:-30000}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the spill_views suite: sort and GROUP BY queries that spill StringView
# and BinaryView columns read from Parquet
# (https://github.com/apache/datafusion/issues/23564). The load SQL COPYs 1M-row
# Parquet files inline, so there is no data step. Each query sets its own memory
# limit and target_partitions, so it spills regardless of the environment.
#   BENCH_SUBGROUP  run one subgroup (repeated, distinct)
run_spill_views() {
    echo "Running spill_views benchmark (subgroup=${BENCH_SUBGROUP:-all})..."
    debug_run env BENCH_NAME=spill_views \
      BENCH_RESULTS_FILE="$(sql_results_file spill_views)" \
      ${BENCH_SUBGROUP:+BENCH_SUBGROUP="${BENCH_SUBGROUP}"} \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the tpch in memory (needs tpch parquet data)
run_tpch_mem() {
    SCALE_FACTOR=$1
    if [ -z "$SCALE_FACTOR" ] ; then
        echo "Internal error: Scale factor not specified"
        exit 1
    fi
    echo "Running tpch_mem benchmark..."

    debug_run env BENCH_NAME=tpch \
      BENCH_RESULTS_FILE="$(sql_results_file tpch_mem_sf${SCALE_FACTOR})" \
      BENCH_SIZE="${SCALE_FACTOR}" \
      DATA_DIR="${DATA_DIR}" \
      TPCH_FILE_TYPE="mem" \
      PREFER_HASH_JOIN="${PREFER_HASH_JOIN}" \
      SIMULATE_LATENCY="${SIMULATE_LATENCY}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"}  \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the tpcds benchmark
run_tpcds() {
    TPCDS_DIR="${DATA_DIR}/tpcds_sf1"

    # Check if TPCDS data directory and representative file exists
    if [ ! -f "${TPCDS_DIR}/web_site.parquet" ]; then
        echo "" >&2
        echo "Please prepare TPC-DS data first by following instructions:" >&2
        echo "  ./bench.sh data tpcds" >&2
        echo "" >&2
        exit 1
    fi

    RESULTS_FILE="${RESULTS_DIR}/tpcds_sf1.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running tpcds benchmark..."

    debug_run $CARGO_COMMAND --bin dfbench -- tpcds --iterations 5 --path "${TPCDS_DIR}" --query_path "../datafusion/core/tests/tpc-ds" --prefer_hash_join "${PREFER_HASH_JOIN}" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the compile profile benchmark helper
run_compile_profile() {
    local profiles=("$@")
    local runner="${SCRIPT_DIR}/compile_profile.py"
    local data_path="${DATA_DIR}/tpch_sf1"

    echo "Running compile profile benchmark..."
    local cmd=(uv run python3 "${runner}" --data "${data_path}")
    if [ ${#profiles[@]} -gt 0 ]; then
        cmd+=(--profiles "${profiles[@]}")
    fi
    debug_run "${cmd[@]}"
}

# Runs the cancellation benchmark
run_cancellation() {
    RESULTS_FILE="${RESULTS_DIR}/cancellation.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running cancellation benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- cancellation --iterations 5 --path "${DATA_DIR}/cancellation" -o "${RESULTS_FILE}" ${LATENCY_ARG}
}


# Downloads the single file hits.parquet ClickBench datasets from
# https://github.com/ClickHouse/ClickBench/tree/main#data-loading
#
# Creates data in $DATA_DIR/hits.parquet
data_clickbench_1() {
    pushd "${DATA_DIR}" > /dev/null

    # Avoid downloading if it already exists and is the right size
    OUTPUT_SIZE=$(wc -c hits.parquet  2>/dev/null  | awk '{print $1}' || true)
    echo -n "Checking hits.parquet..."
    if test "${OUTPUT_SIZE}" = "14779976446"; then
        echo -n "... found ${OUTPUT_SIZE} bytes ..."
    else
        URL="https://datasets.clickhouse.com/hits_compatible/hits.parquet"
        echo -n "... downloading ${URL} (14GB) ... "
        wget --continue ${URL}
    fi
    echo " Done"
    popd > /dev/null
}

# Downloads the 100 file partitioned ClickBench datasets from
# https://github.com/ClickHouse/ClickBench/tree/main#data-loading
#
# Creates data in $DATA_DIR/hits_partitioned
data_clickbench_partitioned() {
    MAX_CONCURRENT_DOWNLOADS=10

    mkdir -p "${DATA_DIR}/hits_partitioned"
    pushd "${DATA_DIR}/hits_partitioned" > /dev/null

    echo -n "Checking hits_partitioned..."
    OUTPUT_SIZE=$(wc -c -- * 2>/dev/null | tail -n 1  | awk '{print $1}' || true)
    if test "${OUTPUT_SIZE}" = "14737666736"; then
        echo -n "... found ${OUTPUT_SIZE} bytes ..."
    else
        echo -n " downloading with ${MAX_CONCURRENT_DOWNLOADS} parallel workers"
        seq 0 99 | xargs -P${MAX_CONCURRENT_DOWNLOADS} -I{} bash -c 'wget -q --continue https://datasets.clickhouse.com/hits_compatible/athena_partitioned/hits_{}.parquet && echo -n "."'
    fi

    echo " Done"
    popd > /dev/null
}


# Runs the clickbench benchmark with a single large parquet file
run_clickbench_1() {
    RESULTS_FILE="${RESULTS_DIR}/clickbench_1.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running clickbench (1 file) benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- clickbench  --iterations 5 --path "${DATA_DIR}/hits.parquet"  --queries-path "${SCRIPT_DIR}/queries/clickbench/queries" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

 # Runs the clickbench benchmark with the partitioned parquet dataset (100 files)
run_clickbench_partitioned() {
    RESULTS_FILE="${RESULTS_DIR}/clickbench_partitioned.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running clickbench (partitioned, 100 files) benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- clickbench  --iterations 5 --path "${DATA_DIR}/hits_partitioned" --queries-path "${SCRIPT_DIR}/queries/clickbench/queries" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}


 # Runs the clickbench benchmark with the partitioned parquet files and filter_pushdown enabled
run_clickbench_pushdown() {
    RESULTS_FILE="${RESULTS_DIR}/clickbench_pushdown.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running clickbench (partitioned, 100 files) benchmark with pushdown_filters=true, reorder_filters=true..."
    debug_run $CARGO_COMMAND --bin dfbench -- clickbench --pushdown --iterations 5 --path "${DATA_DIR}/hits_partitioned" --queries-path "${SCRIPT_DIR}/queries/clickbench/queries" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}


# Runs the clickbench "extended" benchmark with a single large parquet file
run_clickbench_extended() {
    RESULTS_FILE="${RESULTS_DIR}/clickbench_extended.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running clickbench (1 file) extended benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- clickbench  --iterations 5 --path "${DATA_DIR}/hits.parquet" --queries-path "${SCRIPT_DIR}/queries/clickbench/extended" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Downloads the csv.gz files IMDB datasets from Peter Boncz's homepage(one of the JOB paper authors)
# https://event.cwi.nl/da/job/imdb.tgz
data_imdb() {
    local imdb_dir="${DATA_DIR}/imdb"
    local imdb_temp_gz="${imdb_dir}/imdb.tgz"
    local imdb_url="https://event.cwi.nl/da/job/imdb.tgz"

   # imdb has 21 files, we just separate them into 3 groups for better readability
    local first_required_files=(
        "aka_name.parquet"
        "aka_title.parquet"
        "cast_info.parquet"
        "char_name.parquet"
        "comp_cast_type.parquet"
        "company_name.parquet"
        "company_type.parquet"
    )

    local second_required_files=(
        "complete_cast.parquet"
        "info_type.parquet"
        "keyword.parquet"
        "kind_type.parquet"
        "link_type.parquet"
        "movie_companies.parquet"
        "movie_info.parquet"
    )

    local third_required_files=(
        "movie_info_idx.parquet"
        "movie_keyword.parquet"
        "movie_link.parquet"
        "name.parquet"
        "person_info.parquet"
        "role_type.parquet"
        "title.parquet"
    )

    # Combine the three arrays into one
    local required_files=("${first_required_files[@]}" "${second_required_files[@]}" "${third_required_files[@]}")
    local convert_needed=false

    # Create directory if it doesn't exist
    mkdir -p "${imdb_dir}"

    # Check if required files exist
    for file in "${required_files[@]}"; do
        if [ ! -f "${imdb_dir}/${file}" ]; then
            convert_needed=true
            break
        fi
    done

    if [ "$convert_needed" = true ]; then
        # Expected size of the dataset
        expected_size="1263193115" # 1.18 GB

        echo -n "Looking for imdb.tgz... "
        if [ -f "${imdb_temp_gz}" ]; then
            echo "found"
            echo -n "Checking size... "
            OUTPUT_SIZE=$(wc -c "${imdb_temp_gz}" 2>/dev/null | awk '{print $1}' || true)

            #Checking the size of the existing file
            if [ "${OUTPUT_SIZE}" = "${expected_size}" ]; then
                # Existing file is of the expected size, no need for download
                echo "OK ${OUTPUT_SIZE}"
            else
                # Existing file is partially installed, remove it and initiate a new download
                echo "MISMATCH"
                echo "Size less than expected: ${OUTPUT_SIZE} found, ${expected_size} required"
                echo "Downloading IMDB dataset..."
                rm -f "${imdb_temp_gz}"

                # Download the dataset
                curl -o "${imdb_temp_gz}" "${imdb_url}"

                # Size check of the installed file
                DOWNLOADED_SIZE=$(wc -c "${imdb_temp_gz}" | awk '{print $1}')
                if [ "${DOWNLOADED_SIZE}" != "${expected_size}" ]; then
                    echo "Error: Download size mismatch"
                    echo "Expected: ${expected_size}"
                    echo "Got: ${DOWNLOADED_SIZE}"
                    echo "Please re-initiate the download"
                    return 1
                fi
            fi
        else
            # No existing file found, initiate a new download
            echo "not found"
            echo "Downloading IMDB dataset ${expected_size} expected)..."
            # Download the dataset
            curl -o "${imdb_temp_gz}" "${imdb_url}"
        fi

        # Extract the dataset
        tar -xzvf "${imdb_temp_gz}" -C "${imdb_dir}"
        $CARGO_COMMAND --bin imdb -- convert --input ${imdb_dir} --output ${imdb_dir} --format parquet
        echo "IMDB dataset downloaded and extracted."

    else
        echo "IMDB dataset already exists and contains required parquet files."
    fi
}

# Runs the imdb benchmark
run_imdb() {
    IMDB_DIR="${DATA_DIR}/imdb"

    RESULTS_FILE="${RESULTS_DIR}/imdb.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running imdb benchmark..."
    debug_run $CARGO_COMMAND --bin imdb -- benchmark datafusion --iterations 5 --path "${IMDB_DIR}" --prefer_hash_join "${PREFER_HASH_JOIN}" --format parquet -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

data_h2o() {
    # Default values for size and data format
    SIZE=${1:-"SMALL"}
    DATA_FORMAT=${2:-"CSV"}

    # Create directory if it doesn't exist
    H2O_DIR="${DATA_DIR}/h2o"
    mkdir -p "${H2O_DIR}"

    # Generate h2o test data
    echo "Generating h2o test data in ${H2O_DIR} with size=${SIZE} and format=${DATA_FORMAT}"
    uv run falsa groupby --path-prefix="${H2O_DIR}" --size "${SIZE}" --data-format "${DATA_FORMAT}"
}

data_h2o_join() {
    # Default values for size and data format
    SIZE=${1:-"SMALL"}
    DATA_FORMAT=${2:-"CSV"}

    # Create directory if it doesn't exist
    H2O_DIR="${DATA_DIR}/h2o"
    mkdir -p "${H2O_DIR}"

    # Generate h2o test data
    echo "Generating h2o test data in ${H2O_DIR} with size=${SIZE} and format=${DATA_FORMAT}"
    uv run falsa join --path-prefix="${H2O_DIR}" --size "${SIZE}" --data-format "${DATA_FORMAT}"
}

# Runner for h2o groupby benchmark
run_h2o() {
    # Default values for size and data format
    SIZE=${1:-"SMALL"}
    DATA_FORMAT=${2:-"CSV"}
    DATA_FORMAT=$(echo "$DATA_FORMAT" | tr '[:upper:]' '[:lower:]')
    RUN_Type=${3:-"groupby"}

    # Data directory and results file path
    H2O_DIR="${DATA_DIR}/h2o"
    RESULTS_FILE="${RESULTS_DIR}/h2o.json"

    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running h2o groupby benchmark..."

    # Set the file name based on the size
    case "$SIZE" in
        "SMALL")
            FILE_NAME="G1_1e7_1e7_100_0.${DATA_FORMAT}"  # For small dataset
            ;;
        "MEDIUM")
            FILE_NAME="G1_1e8_1e8_100_0.${DATA_FORMAT}"  # For medium dataset
            ;;
        "BIG")
            FILE_NAME="G1_1e9_1e9_100_0.${DATA_FORMAT}"  # For big dataset
            ;;
        *)
            echo "Invalid size. Valid options are SMALL, MEDIUM, or BIG."
            return 1
            ;;
    esac

     # Set the query file name based on the RUN_Type
    QUERY_FILE="${SCRIPT_DIR}/queries/h2o/${RUN_Type}.sql"

    # Run the benchmark using the dynamically constructed file path and query file
    debug_run $CARGO_COMMAND --bin dfbench -- h2o \
        --iterations 3 \
        --path "${H2O_DIR}/${FILE_NAME}" \
        --queries-path "${QUERY_FILE}" \
        -o "${RESULTS_FILE}" \
         ${QUERY_ARG} ${LATENCY_ARG}
}

# Utility function to run h2o join/window benchmark
h2o_runner() {
    # Default values for size and data format
    SIZE=${1:-"SMALL"}
    DATA_FORMAT=${2:-"CSV"}
    DATA_FORMAT=$(echo "$DATA_FORMAT" | tr '[:upper:]' '[:lower:]')
    RUN_Type=${3:-"join"}

    # Data directory and results file path
    H2O_DIR="${DATA_DIR}/h2o"
    RESULTS_FILE="${RESULTS_DIR}/h2o_${RUN_Type}.json"

    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running h2o ${RUN_Type} benchmark..."

    # Set the file name based on the size
    case "$SIZE" in
        "SMALL")
            X_TABLE_FILE_NAME="J1_1e7_NA_0.${DATA_FORMAT}"
            SMALL_TABLE_FILE_NAME="J1_1e7_1e1_0.${DATA_FORMAT}"
            MEDIUM_TABLE_FILE_NAME="J1_1e7_1e4_0.${DATA_FORMAT}"
            LARGE_TABLE_FILE_NAME="J1_1e7_1e7_NA.${DATA_FORMAT}"
            ;;
        "MEDIUM")
            X_TABLE_FILE_NAME="J1_1e8_NA_0.${DATA_FORMAT}"
            SMALL_TABLE_FILE_NAME="J1_1e8_1e2_0.${DATA_FORMAT}"
            MEDIUM_TABLE_FILE_NAME="J1_1e8_1e5_0.${DATA_FORMAT}"
            LARGE_TABLE_FILE_NAME="J1_1e8_1e8_NA.${DATA_FORMAT}"
            ;;
        "BIG")
            X_TABLE_FILE_NAME="J1_1e9_NA_0.${DATA_FORMAT}"
            SMALL_TABLE_FILE_NAME="J1_1e9_1e3_0.${DATA_FORMAT}"
            MEDIUM_TABLE_FILE_NAME="J1_1e9_1e6_0.${DATA_FORMAT}"
            LARGE_TABLE_FILE_NAME="J1_1e9_1e9_NA.${DATA_FORMAT}"
            ;;
        *)
            echo "Invalid size. Valid options are SMALL, MEDIUM, or BIG."
            return 1
            ;;
    esac

    # Set the query file name based on the RUN_Type
    QUERY_FILE="${SCRIPT_DIR}/queries/h2o/${RUN_Type}.sql"

    debug_run $CARGO_COMMAND --bin dfbench -- h2o \
        --iterations 3 \
        --join-paths "${H2O_DIR}/${X_TABLE_FILE_NAME},${H2O_DIR}/${SMALL_TABLE_FILE_NAME},${H2O_DIR}/${MEDIUM_TABLE_FILE_NAME},${H2O_DIR}/${LARGE_TABLE_FILE_NAME}" \
        --queries-path "${QUERY_FILE}" \
        -o "${RESULTS_FILE}" \
         ${QUERY_ARG} ${LATENCY_ARG}
}

# Runners for h2o join benchmark
run_h2o_join() {
    h2o_runner "$1" "$2" "join"
}

# Runners for h2o join benchmark
run_h2o_window() {
    h2o_runner "$1" "$2" "window"
}

# Runs the h2o window_sorted subgroup: window Top-N over an input that declares
# the ordering the window requires.
#
# The `window` subgroup registers `x` with no declared ordering, so
# `output_ordering()` is None and any plan that depends on a declared ordering is
# unreachable from it, however the data happens to sit on disk. This subgroup's
# load SQL writes a sorted copy and registers it `WITH ORDER`, then asserts a
# per-partition top-K operator is actually in the plan so a silent fallback to
# window-plus-filter cannot masquerade as a result.
#
# The sort happens in the untimed `load` step, so it stays out of the
# measurement. Data comes from the same source as the window subgroup, so the
# data step is data_h2o_join.
#
# H2O_FILE_TYPE selects the *source* format and so only affects that untimed
# load: the measured query always reads the sorted Parquet copy, whatever the
# source was. A Parquet source still makes the load markedly cheaper than
# re-parsing a multi-GB CSV, which is why the `_parquet` entries exist; they are
# expected to produce the same measured numbers as their CSV counterparts, not
# different ones. Each entry pairs with the data step that generates its format.
run_h2o_window_sorted() {
    SIZE=${1:-"small"}
    FILE_TYPE=${2:-"csv"}
    echo "Running h2o window_sorted benchmark (size=${SIZE}, source format=${FILE_TYPE})..."
    debug_run env BENCH_NAME=h2o \
      BENCH_RESULTS_FILE="$(sql_results_file h2o_window_sorted_${SIZE}_${FILE_TYPE})" \
      BENCH_SUBGROUP=window_sorted \
      H2O_BENCH_SIZE="${SIZE}" \
      H2O_FILE_TYPE="${FILE_TYPE}" \
      ${QUERY:+BENCH_QUERY="${QUERY}"} \
      bash -c "$SQL_CARGO_COMMAND"
}

# Runs the external aggregation benchmark
run_external_aggr() {
    # Use TPC-H SF1 dataset
    TPCH_DIR="${DATA_DIR}/tpch_sf1"
    RESULTS_FILE="${RESULTS_DIR}/external_aggr.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running external aggregation benchmark..."

    # Only parquet is supported.
    # Since per-operator memory limit is calculated as (total-memory-limit /
    # number-of-partitions), and by default `--partitions` is set to number of
    # CPU cores, we set a constant number of partitions to prevent this
    # benchmark to fail on some machines.
    debug_run $CARGO_COMMAND --bin external_aggr -- benchmark --partitions 4 --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" ${QUERY_ARG}
}

# Runs the sort pushdown benchmark (without WITH ORDER)
# Generates sort pushdown benchmark data: TPC-H lineitem with 3 parts,
# renamed so alphabetical order does NOT match sort key order.
# This forces the sort pushdown optimizer to reorder files by statistics.
#
# tpchgen produces 3 sorted, non-overlapping parquet files:
#   lineitem.1.parquet: l_orderkey 1 ~ 2M        (lowest keys)
#   lineitem.2.parquet: l_orderkey 2M ~ 4M
#   lineitem.3.parquet: l_orderkey 4M ~ 6M       (highest keys)
#
# We rename them so alphabetical order is reversed:
#   a_part3.parquet (highest keys, sorts first alphabetically)
#   b_part2.parquet
#   c_part1.parquet (lowest keys, sorts last alphabetically)
data_sort_pushdown() {
    SORT_PUSHDOWN_DIR="${DATA_DIR}/sort_pushdown/lineitem"
    if [ -d "${SORT_PUSHDOWN_DIR}" ] && [ "$(ls -A ${SORT_PUSHDOWN_DIR}/*.parquet 2>/dev/null)" ]; then
        echo "Sort pushdown data already exists at ${SORT_PUSHDOWN_DIR}"
        return
    fi

    echo "Generating sort pushdown benchmark data (3 parts with reversed naming)..."

    TEMP_DIR="${DATA_DIR}/sort_pushdown_temp"
    mkdir -p "${TEMP_DIR}" "${SORT_PUSHDOWN_DIR}"

    tpchgen-cli --scale-factor 1 --format parquet --parquet-compression='ZSTD(1)' --parts=3 --output-dir "${TEMP_DIR}"

    # Rename: reverse alphabetical order vs key order
    mv "${TEMP_DIR}/lineitem/lineitem.3.parquet" "${SORT_PUSHDOWN_DIR}/a_part3.parquet"
    mv "${TEMP_DIR}/lineitem/lineitem.2.parquet" "${SORT_PUSHDOWN_DIR}/b_part2.parquet"
    mv "${TEMP_DIR}/lineitem/lineitem.1.parquet" "${SORT_PUSHDOWN_DIR}/c_part1.parquet"

    rm -rf "${TEMP_DIR}"

    echo "Sort pushdown data generated at ${SORT_PUSHDOWN_DIR}"
    ls -la "${SORT_PUSHDOWN_DIR}"
}

run_sort_pushdown() {
    SORT_PUSHDOWN_DIR="${DATA_DIR}/sort_pushdown"
    RESULTS_FILE="${RESULTS_DIR}/sort_pushdown.json"
    echo "Running sort pushdown benchmark (no WITH ORDER)..."
    debug_run $CARGO_COMMAND --bin dfbench -- sort-pushdown --iterations 5 --path "${SORT_PUSHDOWN_DIR}" --queries-path "${SCRIPT_DIR}/queries/sort_pushdown" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sort pushdown benchmark with WITH ORDER (enables sort elimination)
run_sort_pushdown_sorted() {
    SORT_PUSHDOWN_DIR="${DATA_DIR}/sort_pushdown"
    RESULTS_FILE="${RESULTS_DIR}/sort_pushdown_sorted.json"
    echo "Running sort pushdown benchmark (with WITH ORDER)..."
    debug_run $CARGO_COMMAND --bin dfbench -- sort-pushdown --sorted --iterations 5 --path "${SORT_PUSHDOWN_DIR}" --queries-path "${SCRIPT_DIR}/queries/sort_pushdown" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Generates data for sort pushdown Inexact benchmark.
#
# Produces multiple parquet files where each file has MULTIPLE row groups
# with scrambled RG order. This tests both:
#   - Row-group-level reorder within each file (reorder_by_statistics)
#   - TopK threshold initialization from RG statistics
#
# Strategy:
# 1. Write a single sorted file with small (100K-row) RGs (~61 RGs total).
# 2. Use pyarrow to redistribute RGs into N_FILES files, scrambling the
#    RG order within each file using a deterministic permutation.
#    Each file gets ~61/N_FILES RGs with narrow, non-overlapping ranges
#    but in scrambled order.
#
# Writing a single file with ORDER BY scramble does NOT work: the parquet
# writer merges rows from adjacent chunks at RG boundaries, widening
# ranges and defeating reorder_by_statistics.
#
# Requires pyarrow (pip install pyarrow).
data_sort_pushdown_inexact() {
    INEXACT_DIR="${DATA_DIR}/sort_pushdown_inexact/lineitem"
    if [ -d "${INEXACT_DIR}" ] && [ "$(ls -A ${INEXACT_DIR}/*.parquet 2>/dev/null)" ]; then
        echo "Sort pushdown Inexact data already exists at ${INEXACT_DIR}"
        return
    fi

    # Check pyarrow dependency (needed to split/scramble RGs)
    if ! python3 -c "import pyarrow" 2>/dev/null; then
        echo "Error: pyarrow is required for sort pushdown Inexact data generation."
        echo "Install with: pip install pyarrow"
        return 1
    fi

    echo "Generating sort pushdown Inexact benchmark data (multi-file, scrambled RGs)..."

    # Re-use the sort_pushdown data as the source (generate if missing)
    data_sort_pushdown

    mkdir -p "${INEXACT_DIR}"
    SRC_DIR="${DATA_DIR}/sort_pushdown/lineitem"

    # Step 1: Write a single sorted file with small (100K-row) RGs
    TMPFILE="${INEXACT_DIR}/_sorted_small_rgs.parquet"
    (cd "${SCRIPT_DIR}/.." && cargo run --release -p datafusion-cli -- -c "
        CREATE EXTERNAL TABLE src
        STORED AS PARQUET
        LOCATION '${SRC_DIR}';

        COPY (SELECT * FROM src ORDER BY l_orderkey)
        TO '${TMPFILE}'
        STORED AS PARQUET
        OPTIONS ('format.max_row_group_size' '100000');
    ")

    # Step 2: Redistribute RGs into 3 files with scrambled RG order.
    # Each file gets ~20 RGs. RG assignment: rg_idx % 3 determines file,
    # permutation (rg_idx * 41 + 7) % n scrambles the order within file.
    python3 -c "
import pyarrow.parquet as pq

pf = pq.ParquetFile('${TMPFILE}')
n = pf.metadata.num_row_groups
n_files = 3

# Assign each RG to a file, scramble order within each file
file_rgs = [[] for _ in range(n_files)]
for rg_idx in range(n):
    slot = (rg_idx * 41 + 7) % n  # scrambled index
    file_id = slot % n_files
    file_rgs[file_id].append(rg_idx)

# Write each file with its assigned RGs (in scrambled order)
for file_id in range(n_files):
    rgs = file_rgs[file_id]
    if not rgs:
        continue
    tables = [pf.read_row_group(rg) for rg in rgs]
    writer = pq.ParquetWriter(
        '${INEXACT_DIR}/part_%03d.parquet' % file_id,
        pf.schema_arrow)
    for t in tables:
        writer.write_table(t)
    writer.close()
    print(f'File part_{file_id:03d}.parquet: {len(rgs)} RGs')
"

    rm -f "${TMPFILE}"
    echo "Sort pushdown Inexact data generated at ${INEXACT_DIR}"
    ls -la "${INEXACT_DIR}"

    # Also generate overlap data: same strategy but with different file count
    # and permutation. Simulates streaming data with network delays where
    # chunks arrive out of sequence.
    #
    # Requires pyarrow (pip install pyarrow).
    OVERLAP_DIR="${DATA_DIR}/sort_pushdown_inexact_overlap/lineitem"
    if [ -d "${OVERLAP_DIR}" ] && [ "$(ls -A ${OVERLAP_DIR}/*.parquet 2>/dev/null)" ]; then
        echo "Sort pushdown Inexact overlap data already exists at ${OVERLAP_DIR}"
        return
    fi

    echo "Generating sort pushdown Inexact overlap data (multi-file, scrambled RGs)..."
    mkdir -p "${OVERLAP_DIR}"

    # Step 1: Write a single sorted file with small (100K-row) RGs
    TMPFILE="${OVERLAP_DIR}/_sorted_small_rgs.parquet"
    (cd "${SCRIPT_DIR}/.." && cargo run --release -p datafusion-cli -- -c "
        CREATE EXTERNAL TABLE src
        STORED AS PARQUET
        LOCATION '${SRC_DIR}';

        COPY (SELECT * FROM src ORDER BY l_orderkey)
        TO '${TMPFILE}'
        STORED AS PARQUET
        OPTIONS ('format.max_row_group_size' '100000');
    ")

    # Step 2: Redistribute into 5 files with scrambled RG order.
    python3 -c "
import pyarrow.parquet as pq

pf = pq.ParquetFile('${TMPFILE}')
n = pf.metadata.num_row_groups
n_files = 5

file_rgs = [[] for _ in range(n_files)]
for rg_idx in range(n):
    slot = (rg_idx * 37 + 13) % n
    file_id = slot % n_files
    file_rgs[file_id].append(rg_idx)

for file_id in range(n_files):
    rgs = file_rgs[file_id]
    if not rgs:
        continue
    tables = [pf.read_row_group(rg) for rg in rgs]
    writer = pq.ParquetWriter(
        '${OVERLAP_DIR}/part_%03d.parquet' % file_id,
        pf.schema_arrow)
    for t in tables:
        writer.write_table(t)
    writer.close()
    print(f'File part_{file_id:03d}.parquet: {len(rgs)} RGs')
"

    rm -f "${TMPFILE}"
}

# Runs the sort pushdown Inexact benchmark (tests RG reorder by statistics).
# Enables pushdown_filters so TopK's dynamic filter is pushed to the parquet
# reader for late materialization (only needed for Inexact path).
run_sort_pushdown_inexact() {
    INEXACT_DIR="${DATA_DIR}/sort_pushdown_inexact"
    RESULTS_FILE="${RESULTS_DIR}/sort_pushdown_inexact.json"
    echo "Running sort pushdown Inexact benchmark (multi-file scrambled RGs, --sorted DESC)..."
    DATAFUSION_EXECUTION_PARQUET_PUSHDOWN_FILTERS=true \
    debug_run $CARGO_COMMAND --bin dfbench -- sort-pushdown --sorted --iterations 5 --path "${INEXACT_DIR}" --queries-path "${SCRIPT_DIR}/queries/sort_pushdown_inexact" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sort pushdown Inexact benchmark WITHOUT declared ordering.
# Tests the Unsupported path in try_pushdown_sort where RG reorder by
# statistics can still help TopK queries without any file ordering guarantee.
run_sort_pushdown_inexact_unsorted() {
    INEXACT_DIR="${DATA_DIR}/sort_pushdown_inexact"
    RESULTS_FILE="${RESULTS_DIR}/sort_pushdown_inexact_unsorted.json"
    echo "Running sort pushdown Inexact benchmark (no WITH ORDER, Unsupported path)..."
    DATAFUSION_EXECUTION_PARQUET_PUSHDOWN_FILTERS=true \
    debug_run $CARGO_COMMAND --bin dfbench -- sort-pushdown --iterations 5 --path "${INEXACT_DIR}" --queries-path "${SCRIPT_DIR}/queries/sort_pushdown_inexact_unsorted" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sort pushdown benchmark with multi-file scrambled RG order.
# Simulates streaming data with network delays — multiple files, each with
# scrambled RGs. Tests both RG-level reorder and TopK stats initialization.
run_sort_pushdown_inexact_overlap() {
    OVERLAP_DIR="${DATA_DIR}/sort_pushdown_inexact_overlap"
    RESULTS_FILE="${RESULTS_DIR}/sort_pushdown_inexact_overlap.json"
    echo "Running sort pushdown Inexact benchmark (multi-file scrambled RGs, streaming data pattern)..."
    DATAFUSION_EXECUTION_PARQUET_PUSHDOWN_FILTERS=true \
    debug_run $CARGO_COMMAND --bin dfbench -- sort-pushdown --sorted --iterations 5 --path "${OVERLAP_DIR}" --queries-path "${SCRIPT_DIR}/queries/sort_pushdown_inexact_overlap" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sort integration benchmark
run_sort_tpch() {
    SCALE_FACTOR=$1
    if [ -z "$SCALE_FACTOR" ] ; then
        echo "Internal error: Scale factor not specified"
        exit 1
    fi
    TPCH_DIR="${DATA_DIR}/tpch_sf${SCALE_FACTOR}"
    RESULTS_FILE="${RESULTS_DIR}/sort_tpch${SCALE_FACTOR}.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running sort tpch benchmark..."

    debug_run $CARGO_COMMAND --bin dfbench -- sort-tpch --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sort tpch integration benchmark with limit 100 (topk)
run_topk_tpch() {
    TPCH_DIR="${DATA_DIR}/tpch_sf1"
    RESULTS_FILE="${RESULTS_DIR}/run_topk_tpch.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running topk tpch benchmark..."

    $CARGO_COMMAND --bin dfbench -- sort-tpch --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" --limit 100 ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the sorted sort tpch integration benchmark with limit 100 (topk)
run_topk_sorted_tpch() {
    TPCH_DIR="${DATA_DIR}/tpch_sf1"
    RESULTS_FILE="${RESULTS_DIR}/run_topk_sorted_tpch.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running sorted topk tpch benchmark..."

    $CARGO_COMMAND --bin dfbench -- sort-tpch --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" --sorted --limit 100 ${QUERY_ARG} ${LATENCY_ARG}
}

# Generates the pre-sorted Parquet inputs for the ASOF join benchmark.
data_asof_join() {
    ASOF_DIR="${DATA_DIR}/asof_join"
    LEFT_FILE="${ASOF_DIR}/q07_left.parquet"
    RIGHT_FILE="${ASOF_DIR}/q07_right.parquet"

    if [ -f "${LEFT_FILE}" ] && [ -f "${RIGHT_FILE}" ]; then
        echo "ASOF join benchmark data already exists at ${ASOF_DIR}"
        return
    fi

    mkdir -p "${ASOF_DIR}"
    echo "Generating ASOF join benchmark data at ${ASOF_DIR}..."
    (
        cd "${DATAFUSION_DIR}"
        debug_run $CARGO_COMMAND -p datafusion-cli -- -c "
            COPY (
                SELECT value / 10 AS group_key,
                       value % 10 + 1 AS ts,
                       value AS payload
                FROM range(100000)
                ORDER BY group_key, ts
            )
            TO '${LEFT_FILE}' STORED AS PARQUET;

            COPY (
                SELECT value / 10 AS group_key,
                       value % 10 AS ts,
                       value AS payload
                FROM range(100000)
                ORDER BY group_key, ts
            )
            TO '${RIGHT_FILE}' STORED AS PARQUET;
        "
    )
}

# Runs the ASOF join benchmark
run_asof_join() {
    RESULTS_FILE="${RESULTS_DIR}/asof_join.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running ASOF join benchmark..."
    debug_run $CARGO_COMMAND --bin benchmark_runner -- asof_join --iterations 5 --path "${DATA_DIR}" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the nlj benchmark
run_nlj() {
    RESULTS_FILE="${RESULTS_DIR}/nlj.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running nlj benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- nlj --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the hj benchmark
run_hj() {
    TPCH_DIR="${DATA_DIR}/tpch_sf10"
    RESULTS_FILE="${RESULTS_DIR}/hj.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running hj benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- hj --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the smj benchmark
run_smj() {
    RESULTS_FILE="${RESULTS_DIR}/smj.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running smj benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- smj --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the dict benchmark
run_dict() {
    RESULTS_FILE="${RESULTS_DIR}/dict.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running dict benchmark..."
    debug_run $CARGO_COMMAND --bin dfbench -- dict --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
}

# Runs the data-free high-cardinality array_agg(DISTINCT) SQL benchmark.
run_array_agg_distinct() {
    echo "Running array_agg_distinct benchmark..."
    debug_run env BENCH_NAME=array_agg_distinct \
      BENCH_RESULTS_FILE="$(sql_results_file array_agg_distinct)" \
      ${QUERY:+BENCH_QUERY="${QUERY}"} \
      bash -c "$SQL_CARGO_COMMAND"
}


compare_benchmarks() {
    BASE_RESULTS_DIR="${SCRIPT_DIR}/results"
    BRANCH1="$1"
    BRANCH2="$2"
    OPTS="$3"

    if [ -z "$BRANCH1" ] ; then
        echo "<branch1> not specified. Available branches:"
        ls -1 "${BASE_RESULTS_DIR}"
        exit 1
    fi

    if [ -z "$BRANCH2" ] ; then
        echo "<branch2> not specified"
        ls -1 "${BASE_RESULTS_DIR}"
        exit 1
    fi

    echo "Comparing ${BRANCH1} and ${BRANCH2}"
    for RESULTS_FILE1 in "${BASE_RESULTS_DIR}/${BRANCH1}"/*.json ; do
	BENCH=$(basename "${RESULTS_FILE1}")
        RESULTS_FILE2="${BASE_RESULTS_DIR}/${BRANCH2}/${BENCH}"
        if test -f "${RESULTS_FILE2}" ; then
            echo "--------------------"
            echo "Benchmark ${BENCH}"
            echo "--------------------"
            uv run python3 "${SCRIPT_DIR}"/compare.py $OPTS "${RESULTS_FILE1}" "${RESULTS_FILE2}"
        else
            echo "Note: Skipping ${RESULTS_FILE1} as ${RESULTS_FILE2} does not exist"
        fi
    done

}

# Creates sorted ClickBench data from hits.parquet (full dataset)
# The data is sorted by EventTime in ascending order
# Uses datafusion-cli to reduce dependencies
clickbench_sorted() {
    SORTED_FILE="${DATA_DIR}/hits_sorted.parquet"
    ORIGINAL_FILE="${DATA_DIR}/hits.parquet"

    # Default memory limit is 12GB, can be overridden with DATAFUSION_MEMORY_GB env var
    MEMORY_LIMIT_GB=${DATAFUSION_MEMORY_GB:-12}

    echo "Creating sorted ClickBench dataset from hits.parquet..."
    echo "Configuration:"
    echo "  Memory limit: ${MEMORY_LIMIT_GB}G"
    echo "  Row group size: 64K rows"
    echo "  Compression: uncompressed"

    if [ ! -f "${ORIGINAL_FILE}" ]; then
        echo "hits.parquet not found. Running data_clickbench_1 first..."
        data_clickbench_1
    fi

    if [ -f "${SORTED_FILE}" ]; then
        echo "Sorted hits.parquet already exists at ${SORTED_FILE}"
        return 0
    fi

    echo "Sorting hits.parquet by EventTime (this may take several minutes)..."

    pushd "${DATAFUSION_DIR}" > /dev/null
    echo "Building datafusion-cli..."
    cargo build --release --bin datafusion-cli
    DATAFUSION_CLI="${DATAFUSION_DIR}/target/release/datafusion-cli"
    popd > /dev/null


    START_TIME=$(date +%s)
    echo "Start time: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Using datafusion-cli to create sorted parquet file..."
    "${DATAFUSION_CLI}" << EOF
-- Memory and performance configuration
SET datafusion.runtime.memory_limit = '${MEMORY_LIMIT_GB}G';
SET datafusion.execution.spill_compression = 'uncompressed';
SET datafusion.execution.sort_spill_reservation_bytes = 10485760; -- 10MB
SET datafusion.execution.batch_size = 8192;
SET datafusion.execution.target_partitions = 1;

-- Parquet output configuration
SET datafusion.execution.parquet.max_row_group_size = 65536;
SET datafusion.execution.parquet.compression = 'uncompressed';

-- Execute sort and write
COPY (SELECT * FROM '${ORIGINAL_FILE}' ORDER BY "EventTime")
TO '${SORTED_FILE}'
STORED AS PARQUET;
EOF

    local result=$?

    END_TIME=$(date +%s)
    DURATION=$((END_TIME - START_TIME))
    echo "End time: $(date '+%Y-%m-%d %H:%M:%S')"

    if [ $result -eq 0 ]; then
        echo "✓ Successfully created sorted ClickBench dataset"

        INPUT_SIZE=$(stat -f%z "${ORIGINAL_FILE}" 2>/dev/null || stat -c%s "${ORIGINAL_FILE}" 2>/dev/null)
        OUTPUT_SIZE=$(stat -f%z "${SORTED_FILE}" 2>/dev/null || stat -c%s "${SORTED_FILE}" 2>/dev/null)
        INPUT_MB=$((INPUT_SIZE / 1024 / 1024))
        OUTPUT_MB=$((OUTPUT_SIZE / 1024 / 1024))

        echo "  Input:  ${INPUT_MB} MB"
        echo "  Output: ${OUTPUT_MB} MB"

        echo ""
        echo "Time Statistics:"
        echo "  Total duration: ${DURATION} seconds ($(printf '%02d:%02d:%02d' $((DURATION/3600)) $((DURATION%3600/60)) $((DURATION%60))))"
        echo "  Throughput: $((INPUT_MB / DURATION)) MB/s"

        return 0
    else
        echo "✗ Error: Failed to create sorted dataset"
        echo "💡 Tip: Try increasing memory with: DATAFUSION_MEMORY_GB=16 ./bench.sh data clickbench_sorted"
        return 1
    fi
}

# Runs the sorted data benchmark with prefer_existing_sort configuration
run_clickbench_sorted() {
    RESULTS_FILE="${RESULTS_DIR}/clickbench_sorted.json"
    echo "RESULTS_FILE: ${RESULTS_FILE}"
    echo "Running sorted data benchmark with prefer_existing_sort optimization..."

    # Ensure sorted data exists
    clickbench_sorted

    # Run benchmark with prefer_existing_sort configuration
    # This allows DataFusion to optimize away redundant sorts while maintaining parallelism
    debug_run $CARGO_COMMAND --bin dfbench -- clickbench \
        --iterations 5 \
        --path "${DATA_DIR}/hits_sorted.parquet" \
        --queries-path "${SCRIPT_DIR}/queries/clickbench/queries/sorted_data" \
        --sorted-by "EventTime" \
        -c datafusion.optimizer.prefer_existing_sort=true \
        -o "${RESULTS_FILE}" \
        ${QUERY_ARG} ${LATENCY_ARG}
}


# And start the process up
main