Track benchmarks over time in CI, starting with one

.github/workflows/ci.yml

@@ -82,7 +82,7 @@ jobs:
         env:
           VIRTUAL_ENV: ${{ github.workspace }}/.venv
         # The ML demos run in their own job, below.
-        run: uv run --no-project pytest -v . -m "not integration" --ignore=tests/test_ml_demos.py
+        run: uv run --no-project pytest -v . -m "not integration and not benchmark" --ignore=tests/test_ml_demos.py
 
   ml-demos:
     # Each ML demo in benchmarks/ml, at its smallest size, checked against
@@ -121,3 +121,125 @@ jobs:
         env:
           VIRTUAL_ENV: ${{ github.workspace }}/.venv
         run: uv run --no-project pytest -v tests/test_ml_demos.py
+
+  benchmarks:
+    # tests/test_benchmarks.py, tracked over time with pytest-benchmark.
+    name: "benchmarks"
+    runs-on: ubuntu-latest
+    # Only on main: a PR from a fork has no permission to push to
+    # gh-pages, and a dispatch from another branch shouldn't record a
+    # point into history that doesn't belong to main's timeline.
+    if: github.ref == 'refs/heads/main' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch')
+    permissions:
+      contents: write
+      deployments: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: dtolnay/rust-toolchain@stable
+        id: rust-toolchain
+
+      - name: Setup sccache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        run: |
+          echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV
+          echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV
+
+      - name: Cache Cargo
+        uses: actions/cache@v4
+        with:
+          path: ~/.cargo
+          key: cargo-cache-${{ steps.rust-toolchain.outputs.cachekey }}-${{ hashFiles('Cargo.lock') }}
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: "3.12"
+          enable-cache: true
+
+      - name: Install xarray_sql
+        run: uv sync --locked --dev --extra benchmark --no-install-package xarray-sql
+      - name: build rust
+        run: uv run --no-project maturin develop --uv
+      - name: Run benchmarks
+        env:
+          VIRTUAL_ENV: ${{ github.workspace }}/.venv
+        run: uv run --no-project pytest tests/test_benchmarks.py -m benchmark --benchmark-only --benchmark-json=benchmark.json
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          name: "Python benchmarks (pytest-benchmark)"
+          tool: "pytest"
+          output-file-path: benchmark.json
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          auto-push: true
+          # mike (docs.yml) owns dev/, 0.x/, latest/, index.html and
+          # versions.json on gh-pages, and replaces dev/ wholesale on
+          # every docs deploy; the default dev/bench would collide
+          # with that and get wiped out. Use a root-level path mike
+          # never touches.
+          benchmark-data-dir-path: benchmarks
+          # Hosted runners are noisy; keep this loose until there's
+          # enough history on gh-pages to pick a tighter threshold.
+          alert-threshold: "200%"
+          comment-on-alert: true
+          fail-on-alert: false
+
+  codspeed:
+    # The same benchmark as the job above, run by CodSpeed
+    # (https://codspeed.io/), which counts instructions rather than
+    # wall time, so it doesn't need repeats to average out a hosted
+    # runner's noise. It can also report on pull requests, including
+    # from forks, since it needs no write access to this repository.
+    #
+    # Dispatch-only for now: CodSpeedHQ/action authenticates against a
+    # CodSpeed project for this repo, which doesn't exist until an org
+    # admin installs CodSpeed's GitHub App on xqlsystems/xarray-sql and
+    # connects it (https://codspeed.io/docs). Until then, `push` or
+    # `pull_request` would just fail. Once it's connected, add `push`
+    # and `pull_request` to the `if:` below (matching the `benchmarks`
+    # job isn't required: CodSpeed's own token doesn't need gh-pages's
+    # write access, so it can run on every PR, not just main).
+    name: "codspeed"
+    runs-on: ubuntu-latest
+    if: github.event_name == 'workflow_dispatch'
+    permissions:
+      contents: read
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: dtolnay/rust-toolchain@stable
+        id: rust-toolchain
+
+      - name: Setup sccache
+        uses: mozilla-actions/sccache-action@v0.0.9
+
+      - name: Configure sccache
+        run: |
+          echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV
+          echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV
+
+      - name: Cache Cargo
+        uses: actions/cache@v4
+        with:
+          path: ~/.cargo
+          key: cargo-cache-${{ steps.rust-toolchain.outputs.cachekey }}-${{ hashFiles('Cargo.lock') }}
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: "3.12"
+          enable-cache: true
+
+      - name: Install xarray_sql
+        run: uv sync --locked --dev --extra benchmark --no-install-package xarray-sql
+      - name: build rust
+        run: uv run --no-project maturin develop --uv
+      - name: Run benchmarks
+        uses: CodSpeedHQ/action@v5
+        env:
+          VIRTUAL_ENV: ${{ github.workspace }}/.venv
+        with:
+          mode: simulation
+          run: uv run --no-project pytest tests/test_benchmarks.py -m benchmark --codspeed

pyproject.toml

@@ -64,6 +64,17 @@ geo = [
 ml = [
   "ddxdb>=0.4.6,<0.5; sys_platform != 'emscripten'",
 ]
+# Benchmarks tracked over time in CI (see tests/test_benchmarks.py).
+# pytest-codspeed recognizes pytest-benchmark's `benchmark` fixture, so
+# a benchmark written against pytest-benchmark runs unmodified under
+# either. Installed together, a --codspeed run disables pytest-
+# benchmark's own instrumentation for that run rather than double-
+# measuring; CodSpeed's CI job (.github/workflows/ci.yml) is
+# dispatch-only until its GitHub App is installed on the repo.
+benchmark = [
+  "pytest-benchmark>=4.0",
+  "pytest-codspeed>=3.0",
+]
 test = [
   "adbc-driver-flightsql>=1.12",
   "adbc-driver-sqlite>=1.12",
@@ -133,6 +144,7 @@ ignore_missing_imports = true
 dev = [
     "xarray_sql[test]",
     "xarray_sql[docs]",
+    "xarray_sql[benchmark]",
     "py-spy>=0.4.0",
     "ruff>=0.15.10",
     "maturin>=1.9.1",
@@ -156,4 +168,5 @@ cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "**
 testpaths = ["tests"]
 markers = [
     "integration: needs network (anonymous GCS); excluded from the CI unit run",
+    "benchmark: tracked by pytest-benchmark in its own CI job; excluded from the CI unit run",
 ]

tests/test_benchmarks.py

@@ -0,0 +1,58 @@
+"""A benchmark tracked over time in CI (see CONTRIBUTING.md).
+
+Starts small: one benchmark, over an in-memory Dataset that needs no
+network, so a regression in xarray-sql, DataFusion or the in-memory
+table adapter shows up in the history github-action-benchmark stores
+on every push to main.
+
+Uses pytest-benchmark's ``benchmark`` fixture. pytest-codspeed
+(https://codspeed.io/) provides a fixture of the same name, so a
+benchmark written this way can run unmodified under either, if
+CodSpeed's instruction-counting is added later.
+"""
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from xarray_sql import Context
+
+pytestmark = pytest.mark.benchmark
+
+
+@pytest.fixture(scope="module")
+def grid() -> xr.Dataset:
+    """A synthetic (time, lat, lon) grid: 100 x 50 x 50 cells, no I/O."""
+    rng = np.random.default_rng(0)
+    time, lat, lon = 100, 50, 50
+    return xr.Dataset(
+        {
+            "temperature": (
+                ("time", "lat", "lon"),
+                rng.normal(size=(time, lat, lon)),
+            ),
+        },
+        coords={
+            "time": np.arange(time),
+            "lat": np.arange(lat),
+            "lon": np.arange(lon),
+        },
+    )
+
+
+def test_groupby_mean_over_time(benchmark, grid):
+    """A climatology query: mean over time, grouped by lat and lon."""
+    ctx = Context()
+    ctx.from_dataset("grid", grid)
+
+    def run():
+        return ctx.sql(
+            "SELECT lat, lon, AVG(temperature) AS mean_temp "
+            'FROM "grid" GROUP BY lat, lon ORDER BY lat, lon'
+        ).to_dataset()
+
+    result = benchmark(run)
+    assert result.sizes["lat"] == grid.sizes["lat"]
+    assert result.sizes["lon"] == grid.sizes["lon"]
+    expected = grid["temperature"].mean("time")
+    np.testing.assert_allclose(result["mean_temp"].values, expected.values)