This is an automated email from the ASF dual-hosted git repository.
avantgardnerio pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 28c26eabe ci: add h2o window suite CI job (#2230)
28c26eabe is described below
commit 28c26eabe921a6e64b3c97c6ce11c6f45f605fe2
Author: Brent Gardner <[email protected]>
AuthorDate: Wed Aug 5 18:41:42 2026 -0600
ci: add h2o window suite CI job (#2230)
* add h2o benchmark suite (groupby / join / window)
Ports the h2o `db-benchmark` runner from arrow-datafusion into
`benchmarks/src/bin/h2o.rs`, with `datafusion` and `ballista`
subcommands sharing the same query loader and table registration.
Suite is inferred from the queries file name:
- groupby.sql → registers `x`
- join.sql → registers `x`, `small`, `medium`, `large`
- window.sql → registers `large` (the 4th `--join-paths` entry)
Query files are copied verbatim from
`arrow-datafusion/benchmarks/queries/h2o/` so query indices align 1:1
between the two runners.
Data generation is unchanged from upstream — point `--path` /
`--join-paths` at h2o CSV or Parquet files produced by
`arrow-datafusion/benchmarks/bench.sh data h2o_small_window` (or the
equivalent for `h2o_small` / `h2o_small_join`).
* feat(benchmarks): add --explain flag to the h2o binary
Adds `--explain` to both `datafusion` and `ballista` subcommands: prints
the physical plan for each query in `--query`/`--queries-path` and exits
without executing. Convenient for iterating on planner rules against the
h2o window suite without spinning up an executor loop.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* ci(h2o): run window suite on SMALL data with AQE + MPT
Mirrors the TPC-H CI job (single job, one scheduler + one executor,
`tpch-ci` build profile) but scoped to the h2o window suite, which
adds coverage beyond what TPC-H already exercises. Data is generated
via `uv run falsa join --size SMALL` — same generator arrow-datafusion
uses.
No --verify yet; this first pass catches crashes and planner
regressions only. Verification against a local DataFusion oracle
(matching TPC-H's --verify) is a follow-up in this PR once we know
how long the run takes.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* ci(h2o): pin setup-uv to ASF-allowlisted SHA (v8.2.0)
`astral-sh/setup-uv@v6` is not on the ASF allowlist, so the workflow
never started (0s startup_failure). Switch to the SHA that's on the
allowlist in apache/infrastructure-actions#actions.yml.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* fix(benchmarks): make h2o ballista --path optional
The join and window suites read data via --join-paths, not --path, so
requiring --path forced callers to pass a dummy value. Make it
Option<String>; the groupby suite errors clearly if unset. Fixes the
h2o CI job.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* chore(benchmarks): add ASF license header to pyproject.toml
Fixes RAT.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* ci(h2o): shorten workflow + job name
GitHub renders `<workflow> / <job>` in the check list; the previous
setting rendered as `h2o window SMALL / h2o window SMALL (AQE on,
multi-partition tasks)` which pushed the runtime off-screen. Now it
renders as `h2o / window SMALL`.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* feat(benchmarks): --verify h2o against local DataFusion oracle
Ports TPC-H's --verify pattern to h2o.rs: when set, builds a
single-process DataFusion SessionContext over the same parquet files,
runs each query on both engines, and diffs the row sets.
h2o queries have no ORDER BY, so distributed and single-process
execution can produce the same rows in different orders. Adds a
canonicalize() helper that lex-sorts both sides by every column before
handing off to compare_results (which handles float tolerance).
Wired into the CI workflow at 100% (no sampling); local run at 1e7
finishes in ~2m45s (release build, wide CPU). CI numbers TBD.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* fix(chaos-testing): gate executor spawns on scheduler readiness
TestClusterBuilder::start spawned the scheduler and immediately looped
spawning executors. Each executor's first act is dialing the scheduler
over gRPC — if that happens before the scheduler has bound its port,
the executor gets GrpcConnectionError("Could not connect to scheduler")
and exits. The scheduler never learns about that executor, and
await_executors then burns its full 120s deadline waiting for a ghost
registration.
Add await_scheduler_ready between spawn and the executor loop: poll
GET /api/executors on the scheduler's REST/gRPC port until it answers,
or fail fast with the log tail if the scheduler already exited.
Reproducer was the cluster tests' 120s timeout in CI on PR #2230.
---------
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
---
.github/workflows/h2o.yml | 187 +++++++++++++
benchmarks/pyproject.toml | 23 ++
benchmarks/queries/h2o/groupby.sql | 19 ++
benchmarks/queries/h2o/join.sql | 9 +
benchmarks/queries/h2o/window.sql | 150 +++++++++++
benchmarks/src/bin/h2o.rs | 540 +++++++++++++++++++++++++++++++++++++
chaos-testing/src/cluster.rs | 38 +++
7 files changed, 966 insertions(+)
diff --git a/.github/workflows/h2o.yml b/.github/workflows/h2o.yml
new file mode 100644
index 000000000..ec9562e77
--- /dev/null
+++ b/.github/workflows/h2o.yml
@@ -0,0 +1,187 @@
+# 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.
+
+# Runs the h2o `window.sql` suite against a Ballista cluster with AQE + MPT on.
+# groupby/join are already exercised by the TPC-H job; window is where h2o adds
+# incremental coverage.
+
+name: h2o
+
+permissions:
+ contents: read
+
+concurrency:
+ group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{
github.workflow }}
+ cancel-in-progress: true
+
+on:
+ push:
+ branches: [ "main" ]
+ paths:
+ - "ballista/**"
+ - "benchmarks/**"
+ - "Cargo.toml"
+ - "Cargo.lock"
+ - "rust-toolchain.toml"
+ - ".github/workflows/h2o.yml"
+ - ".github/actions/setup-builder/**"
+ # docs-only changes cannot affect this job; exclusions are applied last
+ - "!**/*.md"
+ pull_request:
+ paths:
+ - "ballista/**"
+ - "benchmarks/**"
+ - "Cargo.toml"
+ - "Cargo.lock"
+ - "rust-toolchain.toml"
+ - ".github/workflows/h2o.yml"
+ - ".github/actions/setup-builder/**"
+ # docs-only changes cannot affect this job; exclusions are applied last
+ - "!**/*.md"
+ workflow_dispatch:
+
+jobs:
+ h2o-window-small:
+ name: window SMALL
+ runs-on: ubuntu-latest
+ container:
+ image: amd64/rust
+ steps:
+ - name: Install dependencies
+ run: |
+ apt-get update
+ apt-get install -y netcat-traditional
+
+ - uses: actions/[email protected]
+ with:
+ fetch-depth: 1
+
+ - name: Setup Rust toolchain
+ uses: ./.github/actions/setup-builder
+ with:
+ rust-version: stable
+
+ - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4
#2.9.1
+ with:
+ shared-key: h2o-window-small
+
+ - name: Build Ballista binaries
+ run: |
+ cargo build --profile tpch-ci --locked \
+ -p ballista-scheduler \
+ -p ballista-executor \
+ -p ballista-benchmarks
+
+ - name: Install uv
+ uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+
+ - name: Generate h2o SMALL data (join dataset — window reuses it)
+ env:
+ H2O_DIR: ${{ runner.temp }}/h2o-data
+ run: |
+ mkdir -p "$H2O_DIR"
+ # `falsa join` produces J1_1e7_{NA_0,1e1_0,1e4_0,1e7_NA}.parquet.
+ # The window suite only needs the `large` file
(J1_1e7_1e7_NA.parquet),
+ # but the h2o binary's --join-paths CLI expects all four, so we
+ # generate the full set here.
+ cd benchmarks
+ uv run falsa join --path-prefix "$H2O_DIR" --size SMALL
--data-format PARQUET
+
+ - name: Run h2o window suite against Ballista cluster
+ env:
+ H2O_DIR: ${{ runner.temp }}/h2o-data
+ WORK_DIR: ${{ runner.temp }}/work
+ SCHEDULER_LOG: ${{ runner.temp }}/scheduler.log
+ EXECUTOR_LOG: ${{ runner.temp }}/executor.log
+ run: |
+ set -euo pipefail
+
+ mkdir -p "$WORK_DIR"
+
+ ./target/tpch-ci/ballista-scheduler \
+ --bind-host 127.0.0.1 \
+ > "$SCHEDULER_LOG" 2>&1 &
+ SCHEDULER_PID=$!
+
+ ./target/tpch-ci/ballista-executor \
+ --bind-host 127.0.0.1 \
+ --bind-port 50051 \
+ --scheduler-host 127.0.0.1 \
+ --scheduler-connect-timeout-seconds 10 \
+ --vcores 4 \
+ --memory-pool-size 2GB \
+ --work-dir "$WORK_DIR" \
+ > "$EXECUTOR_LOG" 2>&1 &
+ EXECUTOR_PID=$!
+
+ cleanup() {
+ echo "::group::scheduler log (tail)"
+ tail -n 200 "$SCHEDULER_LOG" || true
+ echo "::endgroup::"
+ echo "::group::executor log (tail)"
+ tail -n 200 "$EXECUTOR_LOG" || true
+ echo "::endgroup::"
+ kill "$SCHEDULER_PID" "$EXECUTOR_PID" 2>/dev/null || true
+ wait "$SCHEDULER_PID" "$EXECUTOR_PID" 2>/dev/null || true
+ }
+ trap cleanup EXIT
+
+ echo "Waiting for scheduler on 127.0.0.1:50050..."
+ for _ in $(seq 1 30); do
+ if nc -z 127.0.0.1 50050; then
+ break
+ fi
+ sleep 1
+ done
+ nc -z 127.0.0.1 50050 || { echo "scheduler did not start"; exit 1; }
+
+ echo "Waiting for executor on 127.0.0.1:50051..."
+ for _ in $(seq 1 30); do
+ if nc -z 127.0.0.1 50051; then
+ break
+ fi
+ sleep 1
+ done
+ nc -z 127.0.0.1 50051 || { echo "executor did not start"; exit 1; }
+
+
JOIN_PATHS="$H2O_DIR/J1_1e7_NA_0.parquet,$H2O_DIR/J1_1e7_1e1_0.parquet,$H2O_DIR/J1_1e7_1e4_0.parquet,$H2O_DIR/J1_1e7_1e7_NA.parquet"
+
+ for q in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17; do
+ echo "::group::window Q$q"
+ ./target/tpch-ci/h2o ballista \
+ --host 127.0.0.1 --port 50050 \
+ --queries-path benchmarks/queries/h2o/window.sql \
+ --join-paths "$JOIN_PATHS" \
+ --query "$q" \
+ --iterations 1 \
+ --partitions 4 \
+ --verify \
+ -c ballista.planner.adaptive.enabled=true \
+ -c ballista.scheduler.max_partitions_per_task=0
+ echo "::endgroup::"
+ done
+
+ - name: Upload cluster logs on failure
+ if: failure()
+ uses: actions/upload-artifact@v7
+ with:
+ name: h2o-window-small-cluster-logs
+ retention-days: 14
+ path: |
+ ${{ runner.temp }}/scheduler.log
+ ${{ runner.temp }}/executor.log
+ if-no-files-found: ignore
diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml
new file mode 100644
index 000000000..34d5536ca
--- /dev/null
+++ b/benchmarks/pyproject.toml
@@ -0,0 +1,23 @@
+# 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.
+
+[project]
+name = "ballista-benchmarks"
+version = "0.1.0"
+requires-python = ">=3.11"
+# typing_extensions is an undeclared dependency of falsa
+dependencies = ["rich", "falsa", "typing_extensions"]
diff --git a/benchmarks/queries/h2o/groupby.sql
b/benchmarks/queries/h2o/groupby.sql
new file mode 100644
index 000000000..4fae7a138
--- /dev/null
+++ b/benchmarks/queries/h2o/groupby.sql
@@ -0,0 +1,19 @@
+SELECT id1, SUM(v1) AS v1 FROM x GROUP BY id1;
+
+SELECT id1, id2, SUM(v1) AS v1 FROM x GROUP BY id1, id2;
+
+SELECT id3, SUM(v1) AS v1, AVG(v3) AS v3 FROM x GROUP BY id3;
+
+SELECT id4, AVG(v1) AS v1, AVG(v2) AS v2, AVG(v3) AS v3 FROM x GROUP BY id4;
+
+SELECT id6, SUM(v1) AS v1, SUM(v2) AS v2, SUM(v3) AS v3 FROM x GROUP BY id6;
+
+SELECT id4, id5, MEDIAN(v3) AS median_v3, STDDEV(v3) AS sd_v3 FROM x GROUP BY
id4, id5;
+
+SELECT id3, MAX(v1) - MIN(v2) AS range_v1_v2 FROM x GROUP BY id3;
+
+SELECT id6, largest2_v3 FROM (SELECT id6, v3 AS largest2_v3, ROW_NUMBER() OVER
(PARTITION BY id6 ORDER BY v3 DESC) AS order_v3 FROM x WHERE v3 IS NOT NULL)
sub_query WHERE order_v3 <= 2;
+
+SELECT id2, id4, POWER(CORR(v1, v2), 2) AS r2 FROM x GROUP BY id2, id4;
+
+SELECT id1, id2, id3, id4, id5, id6, SUM(v3) AS v3, COUNT(*) AS count FROM x
GROUP BY id1, id2, id3, id4, id5, id6;
\ No newline at end of file
diff --git a/benchmarks/queries/h2o/join.sql b/benchmarks/queries/h2o/join.sql
new file mode 100644
index 000000000..84cd661fd
--- /dev/null
+++ b/benchmarks/queries/h2o/join.sql
@@ -0,0 +1,9 @@
+SELECT x.id1, x.id2, x.id3, x.id4 as xid4, small.id4 as smallid4, x.id5,
x.id6, x.v1, small.v2 FROM x INNER JOIN small ON x.id1 = small.id1;
+
+SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4,
medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1,
medium.v2 FROM x INNER JOIN medium ON x.id2 = medium.id2;
+
+SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4,
medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1,
medium.v2 FROM x LEFT JOIN medium ON x.id2 = medium.id2;
+
+SELECT x.id1 as xid1, medium.id1 as mediumid1, x.id2, x.id3, x.id4 as xid4,
medium.id4 as mediumid4, x.id5 as xid5, medium.id5 as mediumid5, x.id6, x.v1,
medium.v2 FROM x JOIN medium ON x.id5 = medium.id5;
+
+SELECT x.id1 as xid1, large.id1 as largeid1, x.id2 as xid2, large.id2 as
largeid2, x.id3, x.id4 as xid4, large.id4 as largeid4, x.id5 as xid5, large.id5
as largeid5, x.id6 as xid6, large.id6 as largeid6, x.v1, large.v2 FROM x JOIN
large ON x.id3 = large.id3;
\ No newline at end of file
diff --git a/benchmarks/queries/h2o/window.sql
b/benchmarks/queries/h2o/window.sql
new file mode 100644
index 000000000..346a8e471
--- /dev/null
+++ b/benchmarks/queries/h2o/window.sql
@@ -0,0 +1,150 @@
+-- Basic Window
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER () AS window_basic
+FROM large;
+
+-- Sorted Window
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ first_value(v2) OVER (ORDER BY id3) AS first_order_by,
+ row_number() OVER (ORDER BY id3) AS row_number_order_by
+FROM large;
+
+-- PARTITION BY
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER (PARTITION BY id1) AS sum_by_id1,
+ sum(v2) OVER (PARTITION BY id2) AS sum_by_id2,
+ sum(v2) OVER (PARTITION BY id3) AS sum_by_id3
+FROM large;
+
+-- PARTITION BY ORDER BY
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ first_value(v2) OVER (PARTITION BY id2 ORDER BY id3) AS
first_by_id2_ordered_by_id3
+FROM large;
+
+-- Lead and Lag
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 PRECEDING AND 1
PRECEDING) AS my_lag,
+ first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 FOLLOWING AND 1
FOLLOWING) AS my_lead
+FROM large;
+
+-- Moving Averages
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ avg(v2) OVER (ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) AS
my_moving_average
+FROM large;
+
+-- Rolling Sum
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER (ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW) AS my_rolling_sum
+FROM large;
+
+-- RANGE BETWEEN
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER (ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) AS
my_range_between
+FROM large;
+
+-- First PARTITION BY ROWS BETWEEN
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1
PRECEDING AND 1 PRECEDING) AS my_lag_by_id2,
+ first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1
FOLLOWING AND 1 FOLLOWING) AS my_lead_by_id2
+FROM large;
+
+-- Moving Averages PARTITION BY
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ avg(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND
CURRENT ROW) AS my_moving_average_by_id2
+FROM large;
+
+-- Rolling Sum PARTITION BY
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN UNBOUNDED
PRECEDING AND CURRENT ROW) AS my_rolling_sum_by_id2
+FROM large;
+
+-- RANGE BETWEEN PARTITION BY
+SELECT
+ id1,
+ id2,
+ id3,
+ v2,
+ sum(v2) OVER (PARTITION BY id2 ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND
CURRENT ROW) AS my_range_between_by_id2
+FROM large;
+
+-- Window Top-N (ROW_NUMBER top-2 per partition)
+SELECT id2, largest2_v2 FROM (
+ SELECT id2, v2 AS largest2_v2,
+ ROW_NUMBER() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS order_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE order_v2 <= 2;
+
+-- Window Top-N partition cardinality sweep (id3 % N gives N distinct
partitions).
+-- These exercise PartitionedTopKExec across cardinalities to validate it stays
+-- competitive with the SortExec+Filter baseline as partition count grows.
+-- Window Top-N: 100 partitions
+SELECT pk, largest2_v2 FROM (
+ SELECT id3 % 100 AS pk, v2 AS largest2_v2,
+ ROW_NUMBER() OVER (PARTITION BY id3 % 100 ORDER BY v2 DESC) AS
order_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE order_v2 <= 2;
+
+-- Window Top-N: 1,000 partitions
+SELECT pk, largest2_v2 FROM (
+ SELECT id3 % 1000 AS pk, v2 AS largest2_v2,
+ ROW_NUMBER() OVER (PARTITION BY id3 % 1000 ORDER BY v2 DESC) AS
order_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE order_v2 <= 2;
+
+-- Window Top-N: 10,000 partitions
+SELECT pk, largest2_v2 FROM (
+ SELECT id3 % 10000 AS pk, v2 AS largest2_v2,
+ ROW_NUMBER() OVER (PARTITION BY id3 % 10000 ORDER BY v2 DESC) AS
order_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE order_v2 <= 2;
+
+-- Window Top-N: 100,000 partitions
+SELECT pk, largest2_v2 FROM (
+ SELECT id3 % 100000 AS pk, v2 AS largest2_v2,
+ ROW_NUMBER() OVER (PARTITION BY id3 % 100000 ORDER BY v2 DESC) AS
order_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE order_v2 <= 2;
diff --git a/benchmarks/src/bin/h2o.rs b/benchmarks/src/bin/h2o.rs
new file mode 100644
index 000000000..209e82c39
--- /dev/null
+++ b/benchmarks/src/bin/h2o.rs
@@ -0,0 +1,540 @@
+// 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.
+
+//! H2O `db-benchmark` runner for Ballista.
+//!
+//! Ports `arrow-datafusion/benchmarks/src/h2o.rs` to Ballista. Supports both
+//! single-process DataFusion and distributed Ballista execution over the h2o
+//! groupby, join, and window query files.
+
+use ballista::extension::SessionConfigExt;
+use ballista::prelude::SessionContextExt;
+use ballista_benchmarks::compare_results;
+use ballista_core::object_store::{
+ session_config_with_s3_support, session_state_with_s3_support,
+};
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::compute::{SortColumn, lexsort_to_indices, take};
+use datafusion::error::{DataFusionError, Result};
+use datafusion::prelude::*;
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+use structopt::StructOpt;
+
+#[cfg(feature = "mimalloc")]
+#[global_allocator]
+static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+#[derive(Debug, StructOpt, Clone)]
+struct DataFusionBenchmarkOpt {
+ /// Query number (1..=max). If not specified, runs every query in the file.
+ #[structopt(short, long)]
+ query: Option<usize>,
+
+ /// Path to the queries SQL file. Suite is inferred from the file name
+ /// (`groupby.sql`, `join.sql`, `window.sql`).
+ #[structopt(
+ parse(from_os_str),
+ short = "r",
+ long = "queries-path",
+ default_value = "benchmarks/queries/h2o/groupby.sql"
+ )]
+ queries_path: PathBuf,
+
+ /// Path to the primary data file (used for `groupby.sql` / `window.sql`).
+ #[structopt(
+ parse(from_os_str),
+ short = "p",
+ long = "path",
+ default_value = "benchmarks/data/h2o/G1_1e7_1e7_100_0.csv"
+ )]
+ path: PathBuf,
+
+ /// Comma-separated join data files, in order: x, small, medium, large.
+ /// Used for `join.sql` and `window.sql` (window uses the `large` file).
+ #[structopt(
+ short = "j",
+ long = "join-paths",
+ default_value =
"benchmarks/data/h2o/J1_1e7_NA_0.csv,benchmarks/data/h2o/J1_1e7_1e1_0.csv,benchmarks/data/h2o/J1_1e7_1e4_0.csv,benchmarks/data/h2o/J1_1e7_1e7_NA.csv"
+ )]
+ join_paths: String,
+
+ /// Iterations per query.
+ #[structopt(short = "i", long = "iterations", default_value = "1")]
+ iterations: usize,
+
+ /// Number of partitions to process in parallel.
+ #[structopt(short = "n", long = "partitions", default_value = "2")]
+ partitions: usize,
+
+ /// Batch size when reading CSV or Parquet files.
+ #[structopt(short = "s", long = "batch-size", default_value = "8192")]
+ batch_size: usize,
+
+ /// Print plans and query text.
+ #[structopt(short, long)]
+ debug: bool,
+
+ /// Print the physical plan and exit without running the query.
+ #[structopt(long)]
+ explain: bool,
+}
+
+#[derive(Debug, StructOpt, Clone)]
+struct BallistaBenchmarkOpt {
+ /// Query number (1..=max). If not specified, runs every query in the file.
+ #[structopt(short, long)]
+ query: Option<usize>,
+
+ /// Path to the queries SQL file. Suite is inferred from the file name
+ /// (`groupby.sql`, `join.sql`, `window.sql`).
+ #[structopt(
+ parse(from_os_str),
+ short = "r",
+ long = "queries-path",
+ default_value = "benchmarks/queries/h2o/groupby.sql"
+ )]
+ queries_path: PathBuf,
+
+ /// Path to the primary data file (registered as `x` for `groupby.sql`).
+ /// Not required for the join or window suites — they use `--join-paths`.
+ /// May be a local path or an object-store URL such as
`s3://bucket/prefix`.
+ #[structopt(short = "p", long = "path")]
+ path: Option<String>,
+
+ /// Comma-separated join data files, in order: x, small, medium, large.
+ /// Used for `join.sql` and `window.sql` (window uses the `large` file).
+ /// Paths may be local or object-store URLs.
+ #[structopt(short = "j", long = "join-paths", default_value = "")]
+ join_paths: String,
+
+ /// Iterations per query.
+ #[structopt(short = "i", long = "iterations", default_value = "1")]
+ iterations: usize,
+
+ /// Number of partitions to process in parallel.
+ #[structopt(short = "n", long = "partitions", default_value = "2")]
+ partitions: usize,
+
+ /// Batch size when reading CSV or Parquet files.
+ #[structopt(short = "s", long = "batch-size", default_value = "8192")]
+ batch_size: usize,
+
+ /// Ballista scheduler host.
+ #[structopt(long = "host")]
+ host: String,
+
+ /// Ballista scheduler port.
+ #[structopt(long = "port")]
+ port: u16,
+
+ /// Configuration overrides in `key=value` format. Repeatable.
+ #[structopt(short = "c", long = "config", number_of_values = 1)]
+ config_overrides: Vec<String>,
+
+ /// Print plans and query text.
+ #[structopt(short, long)]
+ debug: bool,
+
+ /// Print the physical plan and exit without running the query.
+ #[structopt(long)]
+ explain: bool,
+
+ /// Verify Ballista's answer against a local DataFusion oracle running the
+ /// same SQL over the same files. Adds one full local execution per query.
+ #[structopt(long)]
+ verify: bool,
+}
+
+#[derive(Debug, StructOpt)]
+#[structopt(name = "h2o", about = "H2O db-benchmark for Ballista")]
+enum H2oOpt {
+ #[structopt(name = "datafusion")]
+ DataFusion(DataFusionBenchmarkOpt),
+ #[structopt(name = "ballista")]
+ Ballista(BallistaBenchmarkOpt),
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+ env_logger::init();
+ match H2oOpt::from_args() {
+ H2oOpt::DataFusion(opt) => run_datafusion(opt).await,
+ H2oOpt::Ballista(opt) => run_ballista(opt).await,
+ }
+}
+
+async fn run_datafusion(opt: DataFusionBenchmarkOpt) -> Result<()> {
+ println!("Running h2o (datafusion) with: {opt:?}");
+ let queries = AllQueries::from_file(&opt.queries_path)?;
+ let query_range = query_range(&queries, opt.query);
+
+ let config = SessionConfig::new()
+ .with_target_partitions(opt.partitions)
+ .with_batch_size(opt.batch_size);
+ let ctx = SessionContext::new_with_config(config);
+
+ let suite = Suite::from_queries_path(&opt.queries_path)?;
+ let paths = SuitePaths {
+ primary: Some(opt.path.to_string_lossy().into_owned()),
+ join_csv: opt.join_paths.clone(),
+ };
+ register_h2o_tables(&ctx, suite, &paths).await?;
+
+ if opt.explain {
+ return explain_queries(&ctx, &queries, query_range).await;
+ }
+
+ run_queries(
+ &ctx,
+ None,
+ &queries,
+ query_range,
+ opt.iterations,
+ opt.debug,
+ "datafusion",
+ )
+ .await
+}
+
+async fn run_ballista(opt: BallistaBenchmarkOpt) -> Result<()> {
+ println!("Running h2o (ballista) with: {opt:?}");
+ let queries = AllQueries::from_file(&opt.queries_path)?;
+ let query_range = query_range(&queries, opt.query);
+
+ let mut config = session_config_with_s3_support()
+ .with_target_partitions(opt.partitions)
+ .with_batch_size(opt.batch_size)
+ .with_ballista_job_name("h2o benchmark");
+
+ for kv in &opt.config_overrides {
+ match kv.split_once('=') {
+ Some((key, value)) => {
+ if let Err(err) = config.options_mut().set(key.trim(),
value.trim()) {
+ println!("Warning: could not set config '{kv}': {err}");
+ }
+ }
+ None => println!(
+ "Warning: ignoring invalid config override '{kv}'. Expected
key=value"
+ ),
+ }
+ }
+
+ let state = session_state_with_s3_support(config)?;
+ let address = format!("df://{}:{}", opt.host, opt.port);
+ let ctx = SessionContext::remote_with_state(&address, state).await?;
+
+ let suite = Suite::from_queries_path(&opt.queries_path)?;
+ let paths = SuitePaths {
+ primary: opt.path.clone(),
+ join_csv: opt.join_paths.clone(),
+ };
+ register_h2o_tables(&ctx, suite, &paths).await?;
+
+ if opt.explain {
+ return explain_queries(&ctx, &queries, query_range).await;
+ }
+
+ let oracle_ctx = if opt.verify {
+ let oracle_config = SessionConfig::new()
+ .with_target_partitions(opt.partitions)
+ .with_batch_size(opt.batch_size);
+ let oracle = SessionContext::new_with_config(oracle_config);
+ register_h2o_tables(&oracle, suite, &paths).await?;
+ Some(oracle)
+ } else {
+ None
+ };
+
+ run_queries(
+ &ctx,
+ oracle_ctx.as_ref(),
+ &queries,
+ query_range,
+ opt.iterations,
+ opt.debug,
+ "ballista",
+ )
+ .await
+}
+
+async fn explain_queries(
+ ctx: &SessionContext,
+ queries: &AllQueries,
+ query_range: std::ops::RangeInclusive<usize>,
+) -> Result<()> {
+ use datafusion::physical_plan::displayable;
+ for query_id in query_range {
+ let sql = queries.get(query_id)?;
+ println!("Q{query_id}: {sql}");
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ println!(
+ "=== Physical plan (Q{query_id}) ===\n{}",
+ displayable(plan.as_ref()).indent(true)
+ );
+ }
+ Ok(())
+}
+
+async fn run_queries(
+ ctx: &SessionContext,
+ oracle_ctx: Option<&SessionContext>,
+ queries: &AllQueries,
+ query_range: std::ops::RangeInclusive<usize>,
+ iterations: usize,
+ debug: bool,
+ label: &str,
+) -> Result<()> {
+ let mut total_secs = 0.0;
+ for query_id in query_range {
+ let sql = queries.get(query_id)?;
+ println!("Q{query_id}: {sql}");
+
+ let mut per_iter_secs = Vec::with_capacity(iterations);
+ let mut row_count = 0usize;
+ let mut last_batches = vec![];
+ for iteration in 1..=iterations {
+ let start = Instant::now();
+ let batches = ctx.sql(sql).await?.collect().await?;
+ let elapsed = start.elapsed().as_secs_f64();
+ row_count = batches.iter().map(|batch| batch.num_rows()).sum();
+ per_iter_secs.push(elapsed);
+ println!(
+ "Query {query_id} iteration {iteration} took {:.3} s and
returned {row_count} rows ({label})",
+ elapsed
+ );
+ last_batches = batches;
+ }
+ let avg = per_iter_secs.iter().sum::<f64>() / per_iter_secs.len() as
f64;
+ println!("Query {query_id} avg time: {avg:.3} s ({row_count} rows)");
+ total_secs += avg;
+
+ if debug {
+ let plan = ctx.sql(sql).await?.into_optimized_plan()?;
+ println!("=== Optimized logical plan ===\n{plan:?}\n");
+ }
+
+ if let Some(oracle) = oracle_ctx {
+ let oracle_start = Instant::now();
+ let expected = oracle.sql(sql).await?.collect().await?;
+ let oracle_secs = oracle_start.elapsed().as_secs_f64();
+ println!("Query {query_id} oracle (datafusion) took
{oracle_secs:.3} s");
+
+ // h2o queries have no ORDER BY, so Ballista's distributed row
+ // order and DataFusion's single-process order can differ.
+ // Lexsort each side before comparing.
+ let sort_start = Instant::now();
+ let expected_sorted = canonicalize(&expected)?;
+ let actual_sorted = canonicalize(&last_batches)?;
+ let sort_secs = sort_start.elapsed().as_secs_f64();
+
+ let compare_start = Instant::now();
+ compare_results(&expected_sorted, &actual_sorted).map_err(|e| {
+ DataFusionError::Execution(format!(
+ "Query {query_id} verification failed: {e}"
+ ))
+ })?;
+ let compare_secs = compare_start.elapsed().as_secs_f64();
+ println!(
+ "Query {query_id} verified vs DataFusion: OK (sort
{sort_secs:.3}s, compare {compare_secs:.3}s)"
+ );
+ }
+ }
+ println!("Total avg time across queries: {total_secs:.3} s");
+ Ok(())
+}
+
+fn query_range(
+ queries: &AllQueries,
+ query: Option<usize>,
+) -> std::ops::RangeInclusive<usize> {
+ match query {
+ Some(query_id) => query_id..=query_id,
+ None => queries.min_id()..=queries.max_id(),
+ }
+}
+
+/// Concatenate `batches` and lex-sort by every column, so distributed and
+/// single-process row orderings compare equal.
+fn canonicalize(batches: &[RecordBatch]) -> Result<Vec<RecordBatch>> {
+ if batches.is_empty() {
+ return Ok(vec![]);
+ }
+ let schema = batches[0].schema();
+ let combined = datafusion::arrow::compute::concat_batches(&schema, batches)
+ .map_err(|e| DataFusionError::Execution(format!("canonicalize concat:
{e}")))?;
+ let sort_cols: Vec<SortColumn> = combined
+ .columns()
+ .iter()
+ .map(|c| SortColumn {
+ values: c.clone(),
+ options: None,
+ })
+ .collect();
+ let indices = lexsort_to_indices(&sort_cols, None)
+ .map_err(|e| DataFusionError::Execution(format!("canonicalize sort:
{e}")))?;
+ let sorted_cols = combined
+ .columns()
+ .iter()
+ .map(|c| take(c.as_ref(), &indices, None))
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .map_err(|e| DataFusionError::Execution(format!("canonicalize take:
{e}")))?;
+ let sorted = RecordBatch::try_new(schema, sorted_cols)
+ .map_err(|e| DataFusionError::Execution(format!("canonicalize batch:
{e}")))?;
+ Ok(vec![sorted])
+}
+
+#[derive(Copy, Clone, Debug)]
+enum Suite {
+ Groupby,
+ Join,
+ Window,
+}
+
+impl Suite {
+ fn from_queries_path(path: &Path) -> Result<Self> {
+ let name = path
+ .file_name()
+ .and_then(|s| s.to_str())
+ .unwrap_or_default();
+ match name {
+ "groupby.sql" => Ok(Suite::Groupby),
+ "join.sql" => Ok(Suite::Join),
+ "window.sql" => Ok(Suite::Window),
+ other => Err(DataFusionError::Plan(format!(
+ "unknown h2o suite {other:?} — expected groupby.sql, join.sql,
or window.sql"
+ ))),
+ }
+ }
+}
+
+struct SuitePaths {
+ primary: Option<String>,
+ join_csv: String,
+}
+
+async fn register_h2o_tables(
+ ctx: &SessionContext,
+ suite: Suite,
+ paths: &SuitePaths,
+) -> Result<()> {
+ match suite {
+ Suite::Groupby => {
+ let primary = paths.primary.as_deref().ok_or_else(|| {
+ DataFusionError::Plan(
+ "groupby suite requires --path pointing at the primary
data file"
+ .to_string(),
+ )
+ })?;
+ register_table(ctx, "x", primary).await
+ }
+ Suite::Join => {
+ let join_paths: Vec<&str> = paths.join_csv.split(',').collect();
+ if join_paths.len() != 4 {
+ return Err(DataFusionError::Plan(format!(
+ "join suite needs 4 comma-separated paths, got {}",
+ join_paths.len()
+ )));
+ }
+ for (table, path) in ["x", "small", "medium", "large"]
+ .iter()
+ .zip(join_paths.iter())
+ {
+ register_table(ctx, table, path.trim()).await?;
+ }
+ Ok(())
+ }
+ Suite::Window => {
+ // The window suite uses only the `large` table from the join
dataset.
+ let large = paths.join_csv.split(',').nth(3).ok_or_else(|| {
+ DataFusionError::Plan(
+ "window suite: --join-paths must contain 4 comma-separated
paths; \
+ the fourth is registered as `large`"
+ .to_string(),
+ )
+ })?;
+ register_table(ctx, "large", large.trim()).await
+ }
+ }
+}
+
+async fn register_table(ctx: &SessionContext, table: &str, path: &str) ->
Result<()> {
+ let extension = Path::new(path)
+ .extension()
+ .and_then(|s| s.to_str())
+ .unwrap_or_default();
+ match extension {
+ "csv" => ctx
+ .register_csv(table, path, CsvReadOptions::default())
+ .await
+ .map_err(|err| {
+ DataFusionError::Context(
+ format!("registering table {table:?} from {path}"),
+ Box::new(err),
+ )
+ }),
+ "parquet" => ctx
+ .register_parquet(table, path, ParquetReadOptions::default())
+ .await
+ .map_err(|err| {
+ DataFusionError::Context(
+ format!("registering table {table:?} from {path}"),
+ Box::new(err),
+ )
+ }),
+ other => Err(DataFusionError::Plan(format!(
+ "unsupported extension {other:?} for {path}"
+ ))),
+ }
+}
+
+struct AllQueries {
+ queries: Vec<String>,
+}
+
+impl AllQueries {
+ fn from_file(path: &Path) -> Result<Self> {
+ let contents = std::fs::read_to_string(path).map_err(|err| {
+ DataFusionError::Execution(format!("reading {path:?}: {err}"))
+ })?;
+ // Queries are separated by blank lines — matches the datafusion h2o
+ // runner exactly so query indices line up 1:1 across both binaries.
+ let queries = contents.split("\n\n").map(str::to_owned).collect();
+ Ok(Self { queries })
+ }
+
+ fn get(&self, query_id: usize) -> Result<&str> {
+ self.queries
+ .get(query_id - 1)
+ .map(String::as_str)
+ .ok_or_else(|| {
+ DataFusionError::Plan(format!(
+ "invalid query id {query_id}. Must be between {} and {}",
+ self.min_id(),
+ self.max_id()
+ ))
+ })
+ }
+
+ fn min_id(&self) -> usize {
+ 1
+ }
+
+ fn max_id(&self) -> usize {
+ self.queries.len()
+ }
+}
diff --git a/chaos-testing/src/cluster.rs b/chaos-testing/src/cluster.rs
index a4c114382..33fc8d040 100644
--- a/chaos-testing/src/cluster.rs
+++ b/chaos-testing/src/cluster.rs
@@ -218,6 +218,13 @@ impl TestClusterBuilder {
_machine_lock: machine_lock,
};
+ // Executors dial the scheduler over gRPC the moment they start. If
they
+ // are spawned before the scheduler has bound its port, some of them
get
+ // GrpcConnectionError and exit — the scheduler never sees them, and
+ // await_executors then burns its full 120s deadline waiting for ghost
+ // registrations. Gate the executor spawns on scheduler readiness.
+ cluster.await_scheduler_ready().await?;
+
for i in 0..cluster.builder.executors {
cluster.spawn_executor(i)?;
}
@@ -351,6 +358,37 @@ impl TestCluster {
Ok(())
}
+ /// Block until the scheduler is accepting REST/gRPC connections.
+ ///
+ /// Returns early with the log tail if the scheduler process has already
+ /// exited, so a startup crash surfaces immediately instead of masquerading
+ /// as a 30s connect timeout.
+ async fn await_scheduler_ready(&mut self) -> Result<(), String> {
+ let deadline = Instant::now() + Duration::from_secs(30);
+ loop {
+ if let Ok(Some(status)) = self.scheduler.try_wait() {
+ return Err(format!(
+ "scheduler exited before becoming ready: {status}\n{}",
+ self.log_tails()
+ ));
+ }
+ if reqwest::get(format!("{}/api/executors", self.rest_url()))
+ .await
+ .and_then(|r| r.error_for_status())
+ .is_ok()
+ {
+ return Ok(());
+ }
+ if Instant::now() > deadline {
+ return Err(format!(
+ "timed out waiting for scheduler to accept
connections\n{}",
+ self.log_tails()
+ ));
+ }
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ }
+ }
+
/// Block until `n` executors have registered with the scheduler.
///
/// The deadline is generous because a loaded CI runner starts these child
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
