.github/workflows/benchmark.yml

name: SpatialBench Benchmark

on:
  # Run every other week on Monday at 6:00 UTC
  schedule:
    - cron: '0 6 1-7,15-21 * 1'
  # Run on PRs that modify benchmark code
  pull_request:
    paths:
      - 'benchmark/**'
      - 'spatialbench-queries/**'
      - '.github/workflows/benchmark.yml'
  # Run on pushes to main that modify benchmark code
  push:
    branches: ["main"]
    paths:
      - 'benchmark/**'
      - 'spatialbench-queries/**'
      - '.github/workflows/benchmark.yml'
  # Allow manual triggering with extended options
  workflow_dispatch:
    inputs:
      scale_factor:
        description: 'Scale factors to benchmark (comma-separated, e.g. "1,10")'
        required: false
        default: '1,10'
        type: string
      engines:
        description: 'Engines to benchmark (comma-separated)'
        required: false
        default: 'duckdb,geopandas,sedonadb,spatial_polars'
        type: string
      sedonadb_version:
        description: 'SedonaDB version (e.g., 1.0.0, leave empty for latest)'
        required: false
        default: ''
        type: string
      duckdb_version:
        description: 'DuckDB version (e.g., 1.0.0, leave empty for latest)'
        required: false
        default: ''
        type: string
      geopandas_version:
        description: 'GeoPandas version (e.g., 1.0.0, leave empty for latest)'
        required: false
        default: ''
        type: string
      spatial_polars_version:
        description: 'Spatial Polars version (e.g., 1.0.0, leave empty for latest)'
        required: false
        default: ''
        type: string
      runs:
        description: 'Number of runs per query (average taken for fair comparison)'
        required: false
        default: '3'
        type: choice
        options:
          - '1'
          - '3'
          - '5'
      sedonadb_nightly:
        description: 'Use SedonaDB nightly build from Gemfury (ignores version if true)'
        required: false
        default: false
        type: boolean
      duckdb_nightly:
        description: 'Use DuckDB pre-release/nightly build (ignores version if true)'
        required: false
        default: false
        type: boolean

concurrency:
  group: ${{ github.repository }}-${{ github.ref }}-benchmark
  cancel-in-progress: true

env:
  CARGO_TERM_COLOR: always
  BENCHMARK_ENGINES: ${{ github.event.inputs.engines || 'duckdb,geopandas,sedonadb,spatial_polars' }}
  BENCHMARK_RUNS: ${{ github.event.inputs.runs || '3' }}
  # Package versions (empty = latest, can be overridden via workflow_dispatch)
  SEDONADB_VERSION: ${{ github.event.inputs.sedonadb_version }}
  DUCKDB_VERSION: ${{ github.event.inputs.duckdb_version }}
  GEOPANDAS_VERSION: ${{ github.event.inputs.geopandas_version }}
  SPATIAL_POLARS_VERSION: ${{ github.event.inputs.spatial_polars_version }}
  # Nightly build options (default: false, use stable releases)
  SEDONADB_NIGHTLY: ${{ github.event.inputs.sedonadb_nightly || 'false' }}
  DUCKDB_NIGHTLY: ${{ github.event.inputs.duckdb_nightly || 'false' }}
  QUERY_TIMEOUT: ${{ github.event.inputs.timeout || '600' }}
  # Hugging Face dataset for benchmark data
  HF_DATASET: apache-sedona/spatialbench
  HF_DATA_VERSION: v0.1.0

jobs:
  # Parse scale factors into a JSON array for matrix strategy
  parse-scale-factors:
    name: Parse Scale Factors
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.parse.outputs.matrix }}
    steps:
      - name: Parse scale factor input
        id: parse
        run: |
          # Default: "1,10" for automatic runs, or user-provided for workflow_dispatch
          INPUT="${{ github.event.inputs.scale_factor || '1,10' }}"
          # Convert comma-separated string to JSON array: "1,10" -> ["1","10"]
          MATRIX=$(echo "$INPUT" | tr ',' '\n' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | jq -R . | jq -s -c .)
          echo "matrix=$MATRIX"
          echo "matrix=$MATRIX" >> $GITHUB_OUTPUT

  # Download benchmark data from Hugging Face
  download-data:
    name: Download Data (SF${{ matrix.scale_factor }})
    needs: parse-scale-factors
    runs-on: ubuntu-latest
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
    steps:
      - uses: actions/checkout@v7

      - name: Cache benchmark data
        id: cache-data
        uses: actions/cache@v6.1.0
        with:
          path: benchmark-data-sf${{ matrix.scale_factor }}
          key: benchmark-data-${{ env.HF_DATA_VERSION }}-sf${{ matrix.scale_factor }}

      - name: Setup Python
        if: steps.cache-data.outputs.cache-hit != 'true'
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'

      - name: Install huggingface-hub
        if: steps.cache-data.outputs.cache-hit != 'true'
        run: pip install huggingface-hub

      - name: Download benchmark data from Hugging Face
        if: steps.cache-data.outputs.cache-hit != 'true'
        env:
          SCALE_FACTOR: ${{ matrix.scale_factor }}
        run: |
          # Map scale factor to HF folder name
          SF="${{ matrix.scale_factor }}"
          if [ "$SF" = "0.1" ]; then
            HF_SF="sf0.1"
          else
            HF_SF="sf${SF}"
          fi

          echo "Downloading data from HF: ${{ env.HF_DATASET }}/${{ env.HF_DATA_VERSION }}/${HF_SF}"

          python -c "
          from huggingface_hub import snapshot_download
          import os

          sf = os.environ['SCALE_FACTOR']
          hf_sf = 'sf0.1' if sf == '0.1' else f'sf{sf}'

          snapshot_download(
              repo_id='${{ env.HF_DATASET }}',
              repo_type='dataset',
              local_dir='hf-data',
              allow_patterns=[f'${{ env.HF_DATA_VERSION }}/{hf_sf}/**'],
          )
          "

          # Move data to expected location
          mkdir -p benchmark-data-sf${{ matrix.scale_factor }}

          SF="${{ matrix.scale_factor }}"
          if [ "$SF" = "0.1" ]; then
            HF_SF="sf0.1"
          else
            HF_SF="sf${SF}"
          fi

          cp -r hf-data/${{ env.HF_DATA_VERSION }}/${HF_SF}/* benchmark-data-sf${{ matrix.scale_factor }}/

          echo "Downloaded data structure:"
          find benchmark-data-sf${{ matrix.scale_factor }} -type f -name "*.parquet" | head -20
          echo ""
          echo "Directory contents:"
          ls -la benchmark-data-sf${{ matrix.scale_factor }}/
          echo ""
          echo "Total size:"
          du -sh benchmark-data-sf${{ matrix.scale_factor }}/

      - name: Show cached data info
        if: steps.cache-data.outputs.cache-hit == 'true'
        run: |
          echo "Using cached benchmark data"
          echo "Directory contents:"
          ls -la benchmark-data-sf${{ matrix.scale_factor }}/
          echo ""
          echo "Total size:"
          du -sh benchmark-data-sf${{ matrix.scale_factor }}/

  # ── Per-query benchmark jobs ──
  # Each query runs in its own job (separate runner) so that if one query
  # OOMs and kills the runner, the remaining queries still execute.
  # max-parallel: 1 ensures queries run sequentially per engine to avoid
  # overloading the CI and to keep results orderly.

  benchmark-duckdb:
    name: DuckDB ${{ matrix.query }} (SF${{ matrix.scale_factor }})
    needs: [parse-scale-factors, download-data]
    runs-on: ubuntu-latest
    if: contains(github.event.inputs.engines || 'duckdb,geopandas,sedonadb,spatial_polars', 'duckdb')
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
        query: [q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12]
      fail-fast: false
      max-parallel: 2  # 1 per scale factor
    steps:
      - uses: actions/checkout@v7

      - name: Restore benchmark data from cache
        uses: actions/cache/restore@v6.1.0
        with:
          path: benchmark-data-sf${{ matrix.scale_factor }}
          key: benchmark-data-${{ env.HF_DATA_VERSION }}-sf${{ matrix.scale_factor }}
          fail-on-cache-miss: true

      - name: Setup Python
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'
          cache: 'pip'
          cache-dependency-path: .github/workflows/benchmark.yml

      - name: Install dependencies
        run: |
          if [ "${{ env.DUCKDB_NIGHTLY }}" = "true" ]; then
            pip install duckdb --pre pyarrow pandas
          elif [ -n "${{ env.DUCKDB_VERSION }}" ]; then
            pip install "duckdb==${{ env.DUCKDB_VERSION }}" pyarrow pandas
          else
            pip install duckdb pyarrow pandas
          fi
          echo "Installed DuckDB version: $(python -c 'import duckdb; print(duckdb.__version__)')"

      - name: Install DuckDB spatial extension
        run: |
          # INSTALL is a no-op on DuckDB 1.5 stable (spatial bundled natively) but required for nightly builds.
          python -c "import duckdb; con = duckdb.connect(); con.execute('INSTALL spatial'); con.execute('LOAD spatial'); print('DuckDB spatial extension installed and loaded')"

      - name: Run DuckDB ${{ matrix.query }}
        run: |
          python benchmark/run_benchmark.py \
            --data-dir benchmark-data-sf${{ matrix.scale_factor }} \
            --engines duckdb \
            --queries ${{ matrix.query }} \
            --runs ${{ env.BENCHMARK_RUNS }} \
            --timeout ${{ env.QUERY_TIMEOUT }} \
            --scale-factor ${{ matrix.scale_factor }} \
            --output duckdb_${{ matrix.query }}_results.json

      - name: Upload results
        if: always() && hashFiles(format('duckdb_{0}_results.json', matrix.query)) != ''
        uses: actions/upload-artifact@v7
        with:
          name: duckdb-${{ matrix.query }}-results-sf${{ matrix.scale_factor }}
          path: duckdb_${{ matrix.query }}_results.json
          retention-days: 30

  benchmark-geopandas:
    name: GeoPandas ${{ matrix.query }} (SF${{ matrix.scale_factor }})
    needs: [parse-scale-factors, download-data]
    runs-on: ubuntu-latest
    if: contains(github.event.inputs.engines || 'duckdb,geopandas,sedonadb,spatial_polars', 'geopandas')
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
        query: [q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12]
      fail-fast: false
      max-parallel: 2
    steps:
      - uses: actions/checkout@v7

      - name: Restore benchmark data from cache
        uses: actions/cache/restore@v6.1.0
        with:
          path: benchmark-data-sf${{ matrix.scale_factor }}
          key: benchmark-data-${{ env.HF_DATA_VERSION }}-sf${{ matrix.scale_factor }}
          fail-on-cache-miss: true

      - name: Setup Python
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'
          cache: 'pip'
          cache-dependency-path: .github/workflows/benchmark.yml

      - name: Install dependencies
        run: |
          if [ -n "${{ env.GEOPANDAS_VERSION }}" ]; then
            pip install "geopandas==${{ env.GEOPANDAS_VERSION }}" pandas pyarrow shapely
          else
            pip install geopandas pandas pyarrow shapely
          fi
          echo "Installed GeoPandas version: $(python -c 'from importlib.metadata import version; print(version("geopandas"))')"

      - name: Run GeoPandas ${{ matrix.query }}
        run: |
          python benchmark/run_benchmark.py \
            --data-dir benchmark-data-sf${{ matrix.scale_factor }} \
            --engines geopandas \
            --queries ${{ matrix.query }} \
            --runs ${{ env.BENCHMARK_RUNS }} \
            --timeout ${{ env.QUERY_TIMEOUT }} \
            --scale-factor ${{ matrix.scale_factor }} \
            --output geopandas_${{ matrix.query }}_results.json

      - name: Upload results
        if: always() && hashFiles(format('geopandas_{0}_results.json', matrix.query)) != ''
        uses: actions/upload-artifact@v7
        with:
          name: geopandas-${{ matrix.query }}-results-sf${{ matrix.scale_factor }}
          path: geopandas_${{ matrix.query }}_results.json
          retention-days: 30

  benchmark-sedonadb:
    name: SedonaDB ${{ matrix.query }} (SF${{ matrix.scale_factor }})
    needs: [parse-scale-factors, download-data]
    runs-on: ubuntu-latest
    if: contains(github.event.inputs.engines || 'duckdb,geopandas,sedonadb,spatial_polars', 'sedonadb')
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
        query: [q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12]
      fail-fast: false
      max-parallel: 2
    steps:
      - uses: actions/checkout@v7

      - name: Restore benchmark data from cache
        uses: actions/cache/restore@v6.1.0
        with:
          path: benchmark-data-sf${{ matrix.scale_factor }}
          key: benchmark-data-${{ env.HF_DATA_VERSION }}-sf${{ matrix.scale_factor }}
          fail-on-cache-miss: true

      - name: Setup Python
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'
          cache: 'pip'
          cache-dependency-path: .github/workflows/benchmark.yml

      - name: Install dependencies
        run: |
          if [ "${{ env.SEDONADB_NIGHTLY }}" = "true" ]; then
            pip install "sedonadb[geopandas]" pandas pyarrow pyproj \
              --pre \
              --index-url https://repo.fury.io/sedona-nightlies/ \
              --extra-index-url https://pypi.org/simple/
          elif [ -n "${{ env.SEDONADB_VERSION }}" ]; then
            pip install "sedonadb[geopandas]==${{ env.SEDONADB_VERSION }}" pandas pyarrow pyproj
          else
            pip install "sedonadb[geopandas]" pandas pyarrow pyproj
          fi
          echo "Installed SedonaDB version: $(python -c 'from importlib.metadata import version; print(version("sedonadb"))')"

      - name: Run SedonaDB ${{ matrix.query }}
        run: |
          python benchmark/run_benchmark.py \
            --data-dir benchmark-data-sf${{ matrix.scale_factor }} \
            --engines sedonadb \
            --queries ${{ matrix.query }} \
            --runs ${{ env.BENCHMARK_RUNS }} \
            --timeout ${{ env.QUERY_TIMEOUT }} \
            --scale-factor ${{ matrix.scale_factor }} \
            --output sedonadb_${{ matrix.query }}_results.json

      - name: Upload results
        if: always() && hashFiles(format('sedonadb_{0}_results.json', matrix.query)) != ''
        uses: actions/upload-artifact@v7
        with:
          name: sedonadb-${{ matrix.query }}-results-sf${{ matrix.scale_factor }}
          path: sedonadb_${{ matrix.query }}_results.json
          retention-days: 30

  benchmark-spatial-polars:
    name: Spatial Polars ${{ matrix.query }} (SF${{ matrix.scale_factor }})
    needs: [parse-scale-factors, download-data]
    runs-on: ubuntu-latest
    if: contains(github.event.inputs.engines || 'duckdb,geopandas,sedonadb,spatial_polars', 'spatial_polars')
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
        query: [q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12]
      fail-fast: false
      max-parallel: 2
    steps:
      - uses: actions/checkout@v7

      - name: Restore benchmark data from cache
        uses: actions/cache/restore@v6.1.0
        with:
          path: benchmark-data-sf${{ matrix.scale_factor }}
          key: benchmark-data-${{ env.HF_DATA_VERSION }}-sf${{ matrix.scale_factor }}
          fail-on-cache-miss: true

      - name: Setup Python
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'
          cache: 'pip'
          cache-dependency-path: .github/workflows/benchmark.yml

      - name: Install dependencies
        run: |
          if [ -n "${{ env.SPATIAL_POLARS_VERSION }}" ]; then
            pip install "spatial-polars[knn]==${{ env.SPATIAL_POLARS_VERSION }}" pyarrow
          else
            pip install "spatial-polars[knn]" pyarrow
          fi
          echo "Installed Spatial Polars version: $(python -c 'from importlib.metadata import version; print(version("spatial-polars"))')"

      - name: Run Spatial Polars ${{ matrix.query }}
        run: |
          python benchmark/run_benchmark.py \
            --data-dir benchmark-data-sf${{ matrix.scale_factor }} \
            --engines spatial_polars \
            --queries ${{ matrix.query }} \
            --runs ${{ env.BENCHMARK_RUNS }} \
            --timeout ${{ env.QUERY_TIMEOUT }} \
            --scale-factor ${{ matrix.scale_factor }} \
            --output spatial_polars_${{ matrix.query }}_results.json

      - name: Upload results
        if: always() && hashFiles(format('spatial_polars_{0}_results.json', matrix.query)) != ''
        uses: actions/upload-artifact@v7
        with:
          name: spatial_polars-${{ matrix.query }}-results-sf${{ matrix.scale_factor }}
          path: spatial_polars_${{ matrix.query }}_results.json
          retention-days: 30

  summarize-results:
    name: Summarize Results (SF${{ matrix.scale_factor }})
    needs: [parse-scale-factors, benchmark-duckdb, benchmark-geopandas, benchmark-sedonadb, benchmark-spatial-polars]
    if: always() && (needs.benchmark-duckdb.result != 'cancelled' || needs.benchmark-geopandas.result != 'cancelled' || needs.benchmark-sedonadb.result != 'cancelled' || needs.benchmark-spatial-polars.result != 'cancelled')
    runs-on: ubuntu-latest
    strategy:
      matrix:
        scale_factor: ${{ fromJson(needs.parse-scale-factors.outputs.matrix) }}
    steps:
      - uses: actions/checkout@v7

      - name: Download all results for this scale factor
        uses: actions/download-artifact@v8
        with:
          pattern: '*-results-sf${{ matrix.scale_factor }}'
          path: results
          merge-multiple: true
        continue-on-error: true

      - name: Setup Python
        uses: actions/setup-python@v6
        with:
          python-version: '3.11'

      - name: Generate summary
        run: |
          python benchmark/summarize_results.py \
            --results-dir results \
            --timeout ${{ env.QUERY_TIMEOUT }} \
            --runs ${{ env.BENCHMARK_RUNS }} \
            --engines ${{ env.BENCHMARK_ENGINES }} \
            --output benchmark_summary.md

      - name: Display summary
        run: cat benchmark_summary.md

      - name: Add summary to job output
        run: cat benchmark_summary.md >> $GITHUB_STEP_SUMMARY

      - name: Upload combined results
        uses: actions/upload-artifact@v7
        with:
          name: benchmark-summary-sf${{ matrix.scale_factor }}
          path: |
            results/
            benchmark_summary.md
          retention-days: 90