Differentiation through aggregates: gradient descent in SQL

Differentiating through SUM/AVG needs no new engine code -- it is linearity:

    d/dθ AVG(loss) = AVG(d(loss)/dθ) = AVG(grad(loss, θ))

Writing the grad *inside* the aggregate (SUM(grad(f, x)), AVG(grad(loss, θ)))
composes with SQL scoping: the marker is rewritten to plain SQL before the
aggregate runs, where the columns are still live. The transposed form
grad(SUM(f), x) is rejected by SQL itself (x is gone after aggregation), so we
stay faithful to SQL rather than hacking grad back across the aggregate node.

This is enough to run gradient descent in SQL. Adds:

  * benchmarks/grad_descent.py: fits y ~= a*x + b by minimising MSE, gradients
    w.r.t. the parameters computed via AVG(grad(loss, param)); parameters live
    in a one-row table cross-joined to the data and updated each step. The fit
    matches numpy's least-squares solution.
  * tests: SUM/AVG(grad(...)) equals the aggregate of the derivative, and an
    end-to-end gradient-descent convergence test.
  * docs: module overview + benchmarks README note the aggregate pattern.

No Rust code change -- only a doc comment; the capability already followed from
the existing rewrite.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017mDoFJgsm9kS7SicGoCVF6

Author: claude

benchmarks/README.md

@@ -29,3 +29,20 @@ Two worked cases, each checked against an analytic reference:
 | Saturation vapour pressure | `A*exp(B*tc/(tc+C))` | `grad(e_s, T)` | closed-form Clausius-Clapeyron slope |
 
 Each query round-trips back to an `xarray.Dataset` via `.to_dataset(...)`.
+
+## `grad_descent.py` — gradient descent in SQL
+
+Fits a line `y ~= a*x + b` by minimising the mean squared error, with the
+gradients w.r.t. the parameters computed in SQL. Differentiating through an
+aggregate is just linearity:
+
+```
+d/dθ AVG(loss) = AVG(grad(loss, θ))
+```
+
+so the gradient is an ordinary aggregate of a `grad` expression — no special
+"differentiate through GROUP BY" machinery, since `grad` becomes plain SQL
+before the aggregate runs. The parameters live in a one-row table cross-joined
+to the data; the optimisation loop is plain Python. The fit matches numpy's
+least-squares solution. Self-contained (no network).
+

benchmarks/grad_descent.py

@@ -0,0 +1,106 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "xarray_sql",
+#   "xarray",
+#   "numpy",
+# ]
+#
+# [tool.uv.sources]
+# xarray_sql = { path = "..", editable = true }
+# ///
+"""Gradient descent in SQL.
+
+Fits a line ``y ~= a*x + b`` to noisy data by minimising the mean squared
+error, computing the gradients entirely in SQL with ``grad(...)``.
+
+The key idea is that differentiating through an aggregate is just linearity:
+
+    d/dθ  AVG( loss )  =  AVG( d(loss)/dθ )  =  AVG( grad(loss, θ) )
+
+so the gradient of the loss w.r.t. a parameter is an ordinary aggregate of a
+``grad`` expression. This composes with SQL's scoping (the ``grad`` sits inside
+the aggregate, where the columns are live), which is why no special
+"differentiate through GROUP BY" machinery is needed — it falls out of the
+existing rewrite, which turns ``grad`` into plain SQL before the aggregate runs.
+
+The parameters ``a`` and ``b`` live in a one-row ``params`` table that is
+cross-joined to the data and re-registered with the updated values after each
+step; the optimisation loop itself is plain Python.
+
+Run standalone:
+
+    uv run benchmarks/grad_descent.py
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+import xarray_sql as xql
+
+# A residual r = y - (a*x + b); the per-row loss is r^2.
+RESIDUAL = "(y - (a * x + b))"
+LOSS = f"{RESIDUAL} * {RESIDUAL}"
+
+
+def set_params(ctx: xql.XarrayContext, a: float, b: float) -> None:
+    """(Re)register the one-row params table holding the current a, b."""
+    if "params" in ctx._registered_datasets:
+        ctx.deregister_table("params")
+        del ctx._registered_datasets["params"]
+    params = xr.Dataset(
+        {"a": (("p",), [a]), "b": (("p",), [b])}, coords={"p": [0]}
+    )
+    ctx.from_dataset("params", params, chunks={"p": 1})
+
+
+def loss_and_grads(ctx: xql.XarrayContext) -> tuple[float, float, float]:
+    """Compute (loss, dL/da, dL/db) over the cross-joined data in one query."""
+    row = ctx.sql(
+        f"""
+        SELECT
+          AVG({LOSS})           AS loss,
+          AVG(grad({LOSS}, a))  AS dl_da,
+          AVG(grad({LOSS}, b))  AS dl_db
+        FROM d CROSS JOIN params
+        """
+    ).to_pandas()
+    return float(row["loss"][0]), float(row["dl_da"][0]), float(row["dl_db"][0])
+
+
+def main() -> None:
+    rng = np.random.default_rng(0)
+    n = 500
+    x = rng.uniform(0.0, 1.0, n)
+    a_true, b_true = 2.0, -1.0
+    y = a_true * x + b_true + rng.normal(0.0, 0.01, n)
+
+    data = xr.Dataset(
+        {"x": (("i",), x), "y": (("i",), y)}, coords={"i": np.arange(n)}
+    )
+    ctx = xql.XarrayContext()
+    ctx.from_dataset("d", data, chunks={"i": n})
+
+    a, b, lr, steps = 0.0, 0.0, 0.4, 200
+    for step in range(steps):
+        set_params(ctx, a, b)
+        loss, dl_da, dl_db = loss_and_grads(ctx)
+        a -= lr * dl_da
+        b -= lr * dl_db
+        if step % 40 == 0 or step == steps - 1:
+            print(f"step {step:3d}: loss={loss:.6f}  a={a:.4f}  b={b:.4f}")
+
+    # Reference: the ordinary least-squares solution.
+    a_ols, b_ols = np.polyfit(x, y, 1)
+    print(f"\nSQL gradient descent:  a={a:.4f}  b={b:.4f}")
+    print(f"least-squares (numpy): a={a_ols:.4f}  b={b_ols:.4f}")
+    assert abs(a - a_ols) < 1e-2 and abs(b - b_ols) < 1e-2
+    print(
+        "\nOK: SQL-computed gradients fit the line to the least-squares solution."
+    )
+
+
+if __name__ == "__main__":
+    main()

src/autograd.rs

@@ -40,6 +40,14 @@
 //! Calls nest, giving higher-order derivatives for free: the rewrite walks
 //! bottom-up, so the inner call in `grad(grad(f, x), x)` is differentiated
 //! first and the outer call differentiates that result.
+//!
+//! Differentiation through an aggregate is just linearity and needs no special
+//! handling: write the `grad` *inside* the aggregate, e.g. `SUM(grad(f, x))` or
+//! `AVG(grad(loss, theta))`. Because the marker is rewritten to plain SQL
+//! before the aggregate runs (and the column is in scope there), this is the
+//! relational `d/dθ Σ f = Σ ∂f/∂θ` — enough to run gradient descent in SQL.
+//! (The transposed form `grad(SUM(f), x)` is rejected by SQL's own scoping,
+//! since `x` is gone after aggregation.)
 
 #![allow(dead_code)]
 

tests/test_autograd.py

@@ -117,6 +117,58 @@ def test_unsupported_function_raises(ctx):
         ctx.sql("SELECT grad(atan2(val, val), val) AS d FROM t").to_pandas()
 
 
+def test_grad_inside_aggregate(ctx):
+    # Differentiation through an aggregate is just linearity:
+    #   AGG(grad(f, x)) == d/dx AGG(f). grad rewrites to plain SQL before the
+    #   aggregate runs, so this composes with no special machinery.
+    val = np.linspace(0.1, 3.0, 16)
+    res = ctx.sql(
+        "SELECT SUM(grad(val * val, val)) AS s, "
+        "AVG(grad(sin(val), val)) AS a FROM t"
+    ).to_pandas()
+    np.testing.assert_allclose(res["s"][0], np.sum(2 * val))
+    np.testing.assert_allclose(res["a"][0], np.mean(np.cos(val)))
+
+
+def test_gradient_descent_in_sql():
+    # End to end: fit y ~= a*x + b by minimising MSE, with the gradients
+    # w.r.t. the parameters computed in SQL via AVG(grad(loss, param)).
+    rng = np.random.default_rng(0)
+    n = 200
+    x = rng.uniform(0.0, 1.0, n)
+    a_true, b_true = 2.0, -1.0
+    y = a_true * x + b_true + rng.normal(0.0, 0.01, n)
+    data = xr.Dataset(
+        {"x": (("i",), x), "y": (("i",), y)}, coords={"i": np.arange(n)}
+    )
+    ctx = xql.XarrayContext()
+    ctx.from_dataset("d", data, chunks={"i": n})
+
+    resid = "(y - (a * x + b))"
+    loss = f"{resid} * {resid}"
+    a, b, lr = 0.0, 0.0, 0.4
+    losses = []
+    for _ in range(120):
+        if "params" in ctx._registered_datasets:
+            ctx.deregister_table("params")
+            del ctx._registered_datasets["params"]
+        params = xr.Dataset(
+            {"a": (("p",), [a]), "b": (("p",), [b])}, coords={"p": [0]}
+        )
+        ctx.from_dataset("params", params, chunks={"p": 1})
+        row = ctx.sql(
+            f"SELECT AVG({loss}) AS loss, "
+            f"AVG(grad({loss}, a)) AS dl_da, "
+            f"AVG(grad({loss}, b)) AS dl_db FROM d CROSS JOIN params"
+        ).to_pandas()
+        losses.append(float(row["loss"][0]))
+        a -= lr * float(row["dl_da"][0])
+        b -= lr * float(row["dl_db"][0])
+
+    assert losses[-1] < losses[0]  # loss decreased
+    np.testing.assert_allclose([a, b], [a_true, b_true], atol=0.05)
+
+
 def test_multi_input_grad_columns(ctx_xy):
     # A full Jacobian written as separate scalar grad() columns:
     # f = x*y  ->  df/dx = y, df/dy = x.