diff --git a/.github/workflows/pr-bench-dispatch.yml b/.github/workflows/pr-bench-dispatch.yml index 3cb60c619bf..47d28e98718 100644 --- a/.github/workflows/pr-bench-dispatch.yml +++ b/.github/workflows/pr-bench-dispatch.yml @@ -151,3 +151,21 @@ jobs: secrets: inherit with: matrix_preset: "pr-compact" + + remove-spatial-label: + runs-on: ubuntu-latest + timeout-minutes: 10 + if: github.event.label.name == 'action/bench-spatial' + steps: + - uses: actions-ecosystem/action-remove-labels@2ce5d41b4b6aa8503e285553f75ed56e0a40bae0 # v1 + if: github.event.pull_request.head.repo.full_name == 'vortex-data/vortex' + with: + labels: action/bench-spatial + fail_on_error: true + + spatial-bench: + needs: remove-spatial-label + uses: ./.github/workflows/pr-bench-sql.yml + secrets: inherit + with: + matrix_preset: "pr-spatial" diff --git a/.github/workflows/pr-bench-sql.yml b/.github/workflows/pr-bench-sql.yml index 40789d631a8..865c0b56551 100644 --- a/.github/workflows/pr-bench-sql.yml +++ b/.github/workflows/pr-bench-sql.yml @@ -29,6 +29,7 @@ on: - "pr-compact" - "pr-all" - "pr-full" + - "pr-spatial" permissions: contents: read diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index cae3630d6b8..621357f0cfc 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -127,6 +127,20 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: ${{ (inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false) && 'true' || 'false' }} + - name: Checkout SpatialBench generator + if: matrix.subcommand == 'spatialbench' + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + repository: apache/sedona-spatialbench + ref: b9221a9c4b02b10db20611d79b4019d2b3c4b68e + path: spatialbench-generator + persist-credentials: false + - name: Build SpatialBench generator + if: matrix.subcommand == 'spatialbench' + working-directory: spatialbench-generator + run: | + cargo build --release --bin spatialbench-cli + echo "SPATIALBENCH_CLI=${GITHUB_WORKSPACE}/spatialbench-generator/target/release/spatialbench-cli" >> "${GITHUB_ENV}" - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: @@ -206,6 +220,7 @@ jobs: --output results.json \ --ingest-jsonl results.ingest.jsonl \ --no-build \ + ${{ matrix.exclude_queries && format('--exclude-queries {0}', matrix.exclude_queries) || '' }} \ ${{ matrix.iterations && format('--iterations {0}', matrix.iterations) || '' }} \ ${{ matrix.scale_factor && format('--opt scale-factor={0}', matrix.scale_factor) || '' }} @@ -226,6 +241,7 @@ jobs: --output results.json \ --ingest-jsonl results.ingest.jsonl \ --no-build \ + ${{ matrix.exclude_queries && format('--exclude-queries {0}', matrix.exclude_queries) || '' }} \ ${{ matrix.iterations && format('--iterations {0}', matrix.iterations) || '' }} \ --opt remote-data-dir="$REMOTE_STORAGE" \ ${{ matrix.scale_factor && format('--opt scale-factor={0}', matrix.scale_factor) || '' }} @@ -255,7 +271,9 @@ jobs: cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR - if: inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false + if: >- + inputs.mode == 'pr' && github.event_name != 'workflow_dispatch' + && github.event.pull_request.head.repo.fork == false uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 with: file-path: comment.md @@ -266,7 +284,9 @@ jobs: comment-tag: bench-pr-comment-${{ matrix.id }}-${{ inputs.matrix_preset }} - name: Comment PR on failure - if: failure() && inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false + if: >- + failure() && inputs.mode == 'pr' && github.event_name != 'workflow_dispatch' + && github.event.pull_request.head.repo.fork == false uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 with: message: | diff --git a/Cargo.lock b/Cargo.lock index 43862759c38..29c6f3b82bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1524,7 +1524,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2113,6 +2113,7 @@ dependencies = [ "datafusion-common", "datafusion-physical-plan", "futures", + "geodatafusion", "itertools 0.14.0", "object_store", "opentelemetry", @@ -2917,7 +2918,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3112,7 +3113,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3559,6 +3560,20 @@ dependencies = [ "wkt 0.14.0", ] +[[package]] +name = "geoarrow-expr-geo" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e4a62ac19c86827c6ec81ea584594b3ee96db5a8119b9774d3466c6b373c434" +dependencies = [ + "arrow-array 58.4.0", + "arrow-buffer 58.4.0", + "geo", + "geo-traits", + "geoarrow-array", + "geoarrow-schema", +] + [[package]] name = "geoarrow-schema" version = "0.8.0" @@ -3572,6 +3587,27 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "geodatafusion" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fecbdd00d0fff2b04635c1b1e4129c217908f0c2d17539e0a2275308afce2552" +dependencies = [ + "arrow-arith 58.4.0", + "arrow-array 58.4.0", + "arrow-buffer 58.4.0", + "arrow-schema 58.4.0", + "datafusion", + "geo", + "geo-traits", + "geoarrow-array", + "geoarrow-expr-geo", + "geoarrow-schema", + "geohash", + "thiserror 1.0.69", + "wkt 0.14.0", +] + [[package]] name = "geographiclib-rs" version = "0.2.7" @@ -3581,6 +3617,16 @@ dependencies = [ "libm", ] +[[package]] +name = "geohash" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f58890382f70caccc5fa388981f7ac80c913795042afce9f3e065695d8f7464" +dependencies = [ + "geo-types", + "libm", +] + [[package]] name = "geojson" version = "0.24.2" @@ -4396,7 +4442,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5886,7 +5932,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7129,7 +7175,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -7858,7 +7904,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -7916,7 +7962,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -8787,7 +8833,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -8797,7 +8843,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -10964,7 +11010,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 199796a32c5..a72f51cd442 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -169,6 +169,7 @@ geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" geoarrow-cast = "0.8.0" +geodatafusion = "0.5.0" get_dir = "0.5.0" glob = "0.3.2" goldenfile = "1" diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py index 9f8f0fee915..2470d95fa22 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py @@ -17,6 +17,7 @@ "pr-compact": "Pull-request SQL benchmarks for Vortex Compact plus Parquet controls.", "pr-all": "The union of the focused PR and PR Compact benchmark matrices.", "pr-full": "Every regular SQL benchmark at full PR target coverage.", + "pr-spatial": "SpatialBench across DataFusion and DuckDB targets.", "nightly": "Large-scale SF=100 TPC-H on NVMe and S3 at default targets.", } @@ -59,6 +60,11 @@ Format.VORTEX_COMPACT, ) COMPACT_DUCKDB_TARGETS = duck(Format.PARQUET, Format.VORTEX_COMPACT) +SPATIAL_TARGETS = df(Format.PARQUET, Format.VORTEX) | duck( + Format.PARQUET, + Format.VORTEX, + Format.VORTEX_SPATIAL_NATIVE, +) DEFAULT = Coverage(DEFAULT_TARGETS) STANDARD = Coverage(STANDARD_TARGETS) @@ -77,6 +83,11 @@ DATAFUSION_VORTEX = Coverage(DATAFUSION_VORTEX_TARGETS) COMPACT = Coverage(COMPACT_TARGETS) COMPACT_DUCKDB = Coverage(COMPACT_DUCKDB_TARGETS) +SPATIAL = Coverage( + SPATIAL_TARGETS, + data_formats=(Format.PARQUET, Format.VORTEX, Format.VORTEX_SPATIAL_NATIVE), + exclude_queries=(5, 7, 9), +) # Concrete benchmark cases @@ -196,6 +207,14 @@ "develop": STANDARD_WITH_DUCKDB, }, ), + BenchmarkCase( + id="spatialbench-nvme", + benchmark=Benchmark.SPATIALBENCH, + name="SpatialBench SF=0.6 on NVME", + scale_factor=0.6, + iterations=1, + runs={"pr-spatial": SPATIAL}, + ), BenchmarkCase( id="statpopgen", benchmark=Benchmark.STATPOPGEN, diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/model.py b/bench-orchestrator/bench_orchestrator/ci_matrix/model.py index 587f7f36da4..b6b5b54c2a3 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/model.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/model.py @@ -24,6 +24,7 @@ class Coverage: targets: TargetSet data_formats: tuple[Format, ...] | None = None + exclude_queries: tuple[int, ...] = () @dataclass(frozen=True) diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/render.py b/bench-orchestrator/bench_orchestrator/ci_matrix/render.py index 8b814349a40..8153635511d 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/render.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/render.py @@ -42,6 +42,8 @@ def _matrix_entry(benchmark: BenchmarkCase, coverage: Coverage) -> dict[str, obj entry["local_dir"] = benchmark.local_dir if benchmark.remote_key is not None: entry["remote_key"] = benchmark.remote_key + if coverage.exclude_queries: + entry["exclude_queries"] = ",".join(str(query) for query in coverage.exclude_queries) return entry diff --git a/bench-orchestrator/bench_orchestrator/config.py b/bench-orchestrator/bench_orchestrator/config.py index 4aceaa59ac5..54dfd822300 100644 --- a/bench-orchestrator/bench_orchestrator/config.py +++ b/bench-orchestrator/bench_orchestrator/config.py @@ -83,11 +83,10 @@ class Benchmark(Enum): Engine.LANCE: [Format.LANCE], } -# Engines each benchmark can run on. Benchmarks default to *every* engine; list one here only to -# restrict it. SpatialBench's queries use DuckDB-specific `ST_*` spatial SQL that DataFusion has no -# functions for yet. +# Engines each benchmark can run on. Benchmarks default to every engine. Add an entry only to +# restrict a benchmark. BENCHMARK_ENGINES: dict[Benchmark, frozenset[Engine]] = { - Benchmark.SPATIALBENCH: frozenset({Engine.DUCKDB}), + Benchmark.SPATIALBENCH: frozenset({Engine.DATAFUSION, Engine.DUCKDB}), } diff --git a/bench-orchestrator/tests/test_config.py b/bench-orchestrator/tests/test_config.py index e2b09b06141..bbdebca8dbc 100644 --- a/bench-orchestrator/tests/test_config.py +++ b/bench-orchestrator/tests/test_config.py @@ -57,8 +57,7 @@ def test_resolve_axis_targets_filters_unsupported_combinations() -> None: assert warnings == ["Format lance is not supported by engine duckdb"] -def test_resolve_axis_targets_skips_engines_a_benchmark_cannot_run() -> None: - # SpatialBench is DuckDB-only (ST_* spatial SQL), so the DataFusion axis is dropped with a warning. +def test_resolve_axis_targets_supports_spatialbench_datafusion() -> None: targets, warnings = resolve_axis_targets( [Engine.DATAFUSION, Engine.DUCKDB], [Format.PARQUET, Format.VORTEX], @@ -66,10 +65,12 @@ def test_resolve_axis_targets_skips_engines_a_benchmark_cannot_run() -> None: ) assert targets == [ + BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.PARQUET), + BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.VORTEX), BenchmarkTarget(engine=Engine.DUCKDB, format=Format.PARQUET), BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX), ] - assert warnings == ["Benchmark spatialbench does not support engine datafusion"] + assert warnings == [] def test_resolve_axis_targets_supports_spatialbench_three_lanes() -> None: @@ -89,14 +90,14 @@ def test_resolve_axis_targets_supports_spatialbench_three_lanes() -> None: assert warnings == [] -def test_validate_targets_rejects_engine_a_benchmark_cannot_run() -> None: +def test_validate_targets_accepts_spatialbench_datafusion() -> None: errors = validate_targets( [BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.PARQUET)], {}, Benchmark.SPATIALBENCH, ) - assert errors == ["Benchmark spatialbench does not support engine datafusion"] + assert errors == [] def test_validate_targets_rejects_remote_lance() -> None: diff --git a/bench-orchestrator/tests/test_matrix.py b/bench-orchestrator/tests/test_matrix.py index d0f765be246..532262b1b72 100644 --- a/bench-orchestrator/tests/test_matrix.py +++ b/bench-orchestrator/tests/test_matrix.py @@ -47,6 +47,7 @@ "pr-compact": COMPACT_IDS, "pr-all": PR_ALL_IDS, "pr-full": REGULAR_IDS, + "pr-spatial": ("spatialbench-nvme",), "nightly": ("tpch-nvme", "tpch-s3"), } @@ -76,6 +77,7 @@ def test_pr_target_selection() -> None: pr = {entry["id"]: entry for entry in _entries("pr")} pr_compact = {entry["id"]: entry for entry in _entries("pr-compact")} pr_full = {entry["id"]: entry for entry in _entries("pr-full")} + pr_spatial = {entry["id"]: entry for entry in _entries("pr-spatial")} assert _targets(pr["tpch-nvme"]) == { ("datafusion", "parquet"), @@ -86,6 +88,14 @@ def test_pr_target_selection() -> None: assert ("datafusion", "lance") in _targets(develop["tpch-nvme"]) assert all(("datafusion", "lance") not in _targets(entry) for entry in pr_full.values()) assert "vortex-compact" in cast("list[str]", pr_full["clickbench-nvme"]["data_formats"]) + assert _targets(pr_spatial["spatialbench-nvme"]) == { + ("datafusion", "parquet"), + ("datafusion", "vortex"), + ("duckdb", "parquet"), + ("duckdb", "vortex"), + ("duckdb", "vortex-spatial-native"), + } + assert pr_spatial["spatialbench-nvme"]["exclude_queries"] == "5,7,9" for entry in pr_compact.values(): targets = _targets(entry) assert {file_format for _engine, file_format in targets} == {"parquet", "vortex-compact"} diff --git a/benchmarks/datafusion-bench/Cargo.toml b/benchmarks/datafusion-bench/Cargo.toml index 306669754d8..a419ee4da05 100644 --- a/benchmarks/datafusion-bench/Cargo.toml +++ b/benchmarks/datafusion-bench/Cargo.toml @@ -38,6 +38,7 @@ datafusion = { workspace = true, features = [ datafusion-common = { workspace = true } datafusion-physical-plan = { workspace = true } futures.workspace = true +geodatafusion.workspace = true itertools.workspace = true object_store = { workspace = true, features = ["aws", "gcp"] } opentelemetry.workspace = true diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index aeea953ba3d..37696da439e 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -23,6 +23,7 @@ use object_store::aws::AmazonS3Builder; use object_store::gcp::GoogleCloudStorageBuilder; use object_store::local::LocalFileSystem; use url::Url; +use vortex_bench::Benchmark; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_datafusion::VortexFormat; @@ -65,6 +66,13 @@ pub fn get_session_context() -> SessionContext { SessionContext::new_with_state(session_state_builder.build()) } +/// Register functions required by one benchmark in a DataFusion session. +pub fn setup_session(session: &SessionContext, benchmark: &dyn Benchmark) { + if benchmark.dataset_name() == "spatialbench" { + geodatafusion::register(session); + } +} + pub fn make_object_store( session: &SessionContext, source: &Url, diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 59cb68539e0..844f58e72c9 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -133,8 +133,9 @@ async fn main() -> anyhow::Result<()> { let benchmark = create_benchmark(args.benchmark, &opts)?; + let query_corpus = benchmark.query_corpus(Engine::DataFusion)?; let filtered_queries = filter_queries( - benchmark.queries()?, + query_corpus, args.queries.as_ref(), args.exclude_queries.as_ref(), ); @@ -182,6 +183,7 @@ async fn main() -> anyhow::Result<()> { let benchmark = &*benchmark; async move { let session = datafusion_bench::get_session_context(); + datafusion_bench::setup_session(&session, benchmark); datafusion_bench::make_object_store(&session, benchmark.data_url())?; register_benchmark_tables(&session, benchmark, format).await?; Ok((session, format)) diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index c680d62f836..d46b78474c1 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -111,8 +111,9 @@ fn main() -> anyhow::Result<()> { let benchmark = create_benchmark(args.benchmark, &opts)?; + let query_corpus = benchmark.query_corpus(Engine::DuckDB)?; let filtered_queries = filter_queries( - benchmark.queries()?, + query_corpus, args.queries.as_ref(), args.exclude_queries.as_ref(), ); diff --git a/vortex-bench/sql/spatialbench.md b/vortex-bench/sql/spatialbench.md index 4b67405aaeb..1355d0f70d4 100644 --- a/vortex-bench/sql/spatialbench.md +++ b/vortex-bench/sql/spatialbench.md @@ -1,21 +1,26 @@ # SpatialBench benchmark The [Apache Sedona SpatialBench](https://sedona.apache.org/spatialbench/) spatial -analytics benchmark: twelve queries (Q1 ... Q12 in [`spatialbench.sql`](./spatialbench.sql), -DuckDB dialect) over a trips/zones schema, exercising spatial predicates and functions such -as `ST_DWithin`, `ST_Intersects`, and `ST_Distance`. The query logic matches upstream -`sedona-spatialbench`; only formatting differs. +analytics benchmark has twelve queries over a trips/zones schema. The queries exercise spatial +predicates and functions such as `ST_DWithin`, `ST_Intersects`, and `ST_Distance`. + +[`spatialbench/duckdb.sql`](./spatialbench/duckdb.sql) contains the DuckDB dialect. +[`spatialbench/datafusion.sql`](./spatialbench/datafusion.sql) contains the equivalent DataFusion +dialect. The query logic matches upstream `sedona-spatialbench`. + +Engine dialects use the `sql//.sql` path. +The harness selects the matching file automatically. The harness lives in [`src/spatialbench`](../src/spatialbench). -## Running locally +## Local use ```bash vx-bench run spatialbench ``` -The default command compares the Parquet and Vortex WKB representations with DuckDB. To run the -native Vortex spatial representation explicitly: +The default command compares Parquet and Vortex with DataFusion and DuckDB. To run the native +Vortex spatial representation explicitly: ```bash vx-bench run spatialbench --engine duckdb --format vortex-spatial-native @@ -26,3 +31,9 @@ To compare all three representations in one run: ```bash vx-bench run spatialbench --engine duckdb --format parquet,vortex,vortex-spatial-native ``` + +To compare DataFusion over Parquet and Vortex: + +```bash +vx-bench run spatialbench --engine datafusion --format parquet,vortex +``` diff --git a/vortex-bench/sql/spatialbench/datafusion.sql b/vortex-bench/sql/spatialbench/datafusion.sql new file mode 100644 index 00000000000..7c1f8b426d8 --- /dev/null +++ b/vortex-bench/sql/spatialbench/datafusion.sql @@ -0,0 +1,174 @@ +-- SpatialBench queries for DataFusion with GeoDataFusion functions. +-- Geometry columns contain WKB, so each spatial expression decodes its operands. + +-- Q1 +SELECT + t.t_tripkey, + ST_X(ST_GeomFromWKB(t.t_pickuploc)) AS pickup_lon, + ST_Y(ST_GeomFromWKB(t.t_pickuploc)) AS pickup_lat, + t.t_pickuptime, + ST_Distance(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromText('POINT (-111.7610 34.8697)')) AS distance_to_center +FROM trip t +WHERE ST_Distance(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromText('POINT (-111.7610 34.8697)')) <= 0.45 +ORDER BY distance_to_center ASC, t.t_tripkey ASC; + +-- Q2 +SELECT COUNT(*) AS trip_count_in_coconino_county +FROM trip t +JOIN ( + SELECT z.z_boundary + FROM zone z + WHERE z.z_name = 'Coconino County' + LIMIT 1 +) county ON ST_Intersects( + ST_GeomFromWKB(t.t_pickuploc), + ST_GeomFromWKB(county.z_boundary) +); + +-- Q3 +SELECT + DATE_TRUNC('month', t.t_pickuptime) AS pickup_month, + COUNT(t.t_tripkey) AS total_trips, + AVG(t.t_distance) AS avg_distance, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration, + AVG(t.t_fare) AS avg_fare +FROM trip t +WHERE ST_Distance( + ST_GeomFromWKB(t.t_pickuploc), + ST_GeomFromText('POLYGON((-111.9060 34.7347, -111.6160 34.7347, -111.6160 35.0047, -111.9060 35.0047, -111.9060 34.7347))') +) <= 0.045 +GROUP BY pickup_month +ORDER BY pickup_month; + +-- Q4 +SELECT z.z_zonekey, z.z_name, COUNT(*) AS trip_count +FROM zone z +JOIN ( + SELECT t.t_pickuploc + FROM trip t + ORDER BY t.t_tip DESC, t.t_tripkey ASC + LIMIT 1000 +) top_trips ON ST_Within(ST_GeomFromWKB(top_trips.t_pickuploc), ST_GeomFromWKB(z.z_boundary)) +GROUP BY z.z_zonekey, z.z_name +ORDER BY trip_count DESC, z.z_zonekey ASC; + +-- Q5 +SELECT + c.c_custkey, + c.c_name AS customer_name, + DATE_TRUNC('month', t.t_pickuptime) AS pickup_month, + ST_Area(ST_ConvexHull(ST_Collect(ARRAY_AGG(ST_GeomFromWKB(t.t_dropoffloc))))) AS monthly_travel_hull_area, + COUNT(*) AS dropoff_count +FROM trip t +JOIN customer c ON t.t_custkey = c.c_custkey +GROUP BY c.c_custkey, c.c_name, pickup_month +HAVING dropoff_count > 5 +ORDER BY dropoff_count DESC, c.c_custkey ASC; + +-- Q6 +SELECT + z.z_zonekey, + z.z_name, + COUNT(t.t_tripkey) AS total_pickups, + AVG(t.t_totalamount) AS avg_distance, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration +FROM trip t, zone z +WHERE ST_Intersects( + ST_GeomFromText('POLYGON((-112.2110 34.4197, -111.3110 34.4197, -111.3110 35.3197, -112.2110 35.3197, -112.2110 34.4197))'), + ST_GeomFromWKB(z.z_boundary) +) +AND ST_Within(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(z.z_boundary)) +GROUP BY z.z_zonekey, z.z_name +ORDER BY total_pickups DESC, z.z_zonekey ASC; + +-- Q7 +WITH trip_lengths AS ( + SELECT + t.t_tripkey, + t.t_distance AS reported_distance_m, + ST_Length(ST_MakeLine(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(t.t_dropoffloc))) / 0.000009 AS line_distance_m + FROM trip t +) +SELECT + t.t_tripkey, + t.reported_distance_m, + t.line_distance_m, + t.reported_distance_m / NULLIF(t.line_distance_m, 0) AS detour_ratio +FROM trip_lengths t +ORDER BY detour_ratio DESC NULLS LAST, reported_distance_m DESC, t_tripkey ASC; + +-- Q8 +SELECT b.b_buildingkey, b.b_name, COUNT(*) AS nearby_pickup_count +FROM trip t +JOIN building b ON ST_Distance(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(b.b_boundary)) <= 0.0045 +GROUP BY b.b_buildingkey, b.b_name +ORDER BY nearby_pickup_count DESC, b.b_buildingkey ASC; + +-- Q9 +WITH b1 AS ( + SELECT b_buildingkey AS id, ST_GeomFromWKB(b_boundary) AS geom FROM building +), +b2 AS ( + SELECT b_buildingkey AS id, ST_GeomFromWKB(b_boundary) AS geom FROM building +), +pairs AS ( + SELECT + b1.id AS building_1, + b2.id AS building_2, + ST_Area(b1.geom) AS area1, + ST_Area(b2.geom) AS area2, + ST_Area(ST_Intersection(b1.geom, b2.geom)) AS overlap_area + FROM b1 + JOIN b2 ON b1.id < b2.id AND ST_Intersects(b1.geom, b2.geom) +) +SELECT + building_1, + building_2, + area1, + area2, + overlap_area, + CASE + WHEN overlap_area = 0 THEN 0.0 + WHEN (area1 + area2 - overlap_area) = 0 THEN 1.0 + ELSE overlap_area / (area1 + area2 - overlap_area) + END AS iou +FROM pairs +ORDER BY iou DESC, building_1 ASC, building_2 ASC; + +-- Q10 +SELECT + z.z_zonekey, + z.z_name AS pickup_zone, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration, + AVG(t.t_distance) AS avg_distance, + COUNT(t.t_tripkey) AS num_trips +FROM zone z +LEFT JOIN trip t ON ST_Within(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(z.z_boundary)) +GROUP BY z.z_zonekey, z.z_name +ORDER BY avg_duration DESC NULLS LAST, z.z_zonekey ASC; + +-- Q11 +SELECT COUNT(*) AS cross_zone_trip_count +FROM trip t +JOIN zone pickup_zone ON ST_Within(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(pickup_zone.z_boundary)) +JOIN zone dropoff_zone ON ST_Within(ST_GeomFromWKB(t.t_dropoffloc), ST_GeomFromWKB(dropoff_zone.z_boundary)) +WHERE pickup_zone.z_zonekey != dropoff_zone.z_zonekey; + +-- Q12 +SELECT + t.t_tripkey, + t.t_pickuploc, + nb.b_buildingkey, + nb.building_name, + nb.distance_to_building +FROM trip t +CROSS JOIN LATERAL ( + SELECT + b.b_buildingkey, + b.b_name AS building_name, + ST_Distance(ST_GeomFromWKB(t.t_pickuploc), ST_GeomFromWKB(b.b_boundary)) AS distance_to_building + FROM building b + ORDER BY distance_to_building + LIMIT 5 +) AS nb +ORDER BY nb.distance_to_building, nb.b_buildingkey; diff --git a/vortex-bench/sql/spatialbench.sql b/vortex-bench/sql/spatialbench/duckdb.sql similarity index 100% rename from vortex-bench/sql/spatialbench.sql rename to vortex-bench/sql/spatialbench/duckdb.sql diff --git a/vortex-bench/src/benchmark.rs b/vortex-bench/src/benchmark.rs index 47de30a5faf..d10230751a8 100644 --- a/vortex-bench/src/benchmark.rs +++ b/vortex-bench/src/benchmark.rs @@ -3,6 +3,7 @@ //! Core benchmark trait and types. +use std::fs; use std::path::Path; use arrow_schema::Schema; @@ -12,6 +13,18 @@ use url::Url; use crate::BenchmarkDataset; use crate::Engine; use crate::Format; +use crate::workspace_root; + +pub(crate) fn read_query_file(path: &Path) -> anyhow::Result> { + let contents = fs::read_to_string(path)?; + Ok(contents + .split_terminator(';') + .map(str::trim) + .filter(|statement| !statement.is_empty()) + .enumerate() + .map(|(index, statement)| (index + 1, statement.to_owned())) + .collect()) +} /// Specification for a table in a benchmark dataset. #[derive(Debug)] @@ -35,6 +48,22 @@ pub trait Benchmark: Send + Sync { /// Get all available queries for this benchmark fn queries(&self) -> anyhow::Result>; + /// Get the equivalent query corpus for one engine dialect. + /// + /// If `vortex-bench/sql//.sql` exists, this method loads that file. + /// Otherwise, it uses the benchmark's default queries. + fn query_corpus(&self, engine: Engine) -> anyhow::Result> { + let engine_queries = workspace_root() + .join("vortex-bench") + .join("sql") + .join(self.dataset_name()) + .join(format!("{engine}.sql")); + if engine_queries.try_exists()? { + return read_query_file(&engine_queries); + } + self.queries() + } + /// SQL an `engine` must run before this benchmark's queries (e.g. loading engine /// extensions). Runners replay these after every (re)open. Default: none. fn engine_init_sql(&self, _engine: Engine) -> Vec { diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index 693a174eae1..86a104d5cf2 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -3,7 +3,6 @@ //! SpatialBench benchmark implementation -use std::fs; use std::path::Path; use url::Url; @@ -13,6 +12,7 @@ use crate::BenchmarkDataset; use crate::Engine; use crate::Format; use crate::TableSpec; +use crate::benchmark::read_query_file; use crate::spatialbench::datagen; use crate::spatialbench::datagen::Table; use crate::utils::file::resolve_data_url; @@ -62,22 +62,15 @@ impl Benchmark for SpatialBenchBenchmark { "vortex-bench/sql/spatialbench.md" } - /// All SpatialBench queries, numbered started at Q1 in `spatialbench.sql` file order. + /// All SpatialBench queries, numbered from Q1 in `spatialbench/duckdb.sql` file order. fn queries(&self) -> anyhow::Result> { - // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. - let queries_file = workspace_root() - .join("vortex-bench") - .join("sql") - .join("spatialbench") - .with_extension("sql"); - let contents = fs::read_to_string(queries_file)?; - Ok(contents - .split_terminator(';') - .map(str::trim) - .filter(|stmt| !stmt.is_empty()) - .enumerate() - .map(|(idx, stmt)| (idx + 1, stmt.to_string())) - .collect()) + read_query_file( + &workspace_root() + .join("vortex-bench") + .join("sql") + .join("spatialbench") + .join("duckdb.sql"), + ) } async fn generate_base_data(&self) -> anyhow::Result<()> { @@ -204,3 +197,22 @@ fn zone_parquet_present(parquet_dir: &Path) -> bool { .map(|mut paths| paths.next().is_some()) .unwrap_or(false) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn discovers_engine_query_corpora() -> anyhow::Result<()> { + let benchmark = SpatialBenchBenchmark::new("0.6".to_string(), None)?; + + let duckdb = benchmark.query_corpus(Engine::DuckDB)?; + let datafusion = benchmark.query_corpus(Engine::DataFusion)?; + + assert_eq!(duckdb.len(), 12); + assert_eq!(datafusion.len(), 12); + assert!(duckdb[0].1.contains("ST_X(t.t_pickuploc)")); + assert!(datafusion[0].1.contains("ST_GeomFromWKB")); + Ok(()) + } +}