Differentiable SQL: grad()/jvp()/vjp() over xarray and any DataFusion table

[alxmrs]

Differentiable SQL for xarray-sql: write grad(expr, column) in a query and get
its derivative back as an ordinary column, evaluated row-by-row by DataFusion
alongside everything else. Because each row of a table is an independent
evaluation point, differentiating a column expression and letting DataFusion
evaluate it per row is the relational equivalent of jax.vmap(jax.grad(f)) — the
rows are the batch dimension. This turns SQL into a place you can express
gradients, directional derivatives, and whole training loops.

The surface

Three scalar operations, each returning one value per row (staying in the
long/tidy data model):

• grad(expr, column) — the partial derivative d(expr)/d(column).
• jvp(expr, column, tangent) — forward-mode directional derivative,
  d(expr)/d(column) * tangent (seed a tangent on an input). A multi-input
  directional derivative is a sum of jvp terms.
• vjp(expr, column, cotangent) — reverse-mode pullback,
  cotangent * d(expr)/d(column) (seed a cotangent on the output).

A full gradient or Jacobian is expressed as several scalar columns rather than a
nested array, which would break the one-value-per-coordinate model:

SELECT i, grad(x * y, x) AS dfdx, grad(x * y, y) AS dfdy FROM g

What you can write:

• Higher-order derivatives, by nesting: grad(grad(sin(x), x), x).

• Differentiation through an aggregate — it's just linearity, so put the
  marker inside the aggregate: AVG(grad(loss, theta)) is the relational
  d/dθ (Σ loss) / N. Enough to run gradient descent in SQL.

• grad inside any query shape — recursive CTEs, DML, and subqueries — so a
  training loop can live entirely in one declarative query. For example,
  Newton's method for √2 driven by grad in the recursive term:

  WITH RECURSIVE newton AS (
    SELECT 0 AS step, CAST(1.0 AS DOUBLE) AS x
    UNION ALL
    SELECT step + 1, x - (x * x - 2.0) / grad(x * x - 2.0, x)
    FROM newton WHERE step < 20)
  SELECT x FROM newton ORDER BY step DESC LIMIT 1

• grad over any registered table, not just xarray datasets: plain in-memory
  tables holding model parameters, and schema-qualified tables (era5.surface)
  from mixed-dimension datasets all work.

Also exported: differentiate_sql(expr, wrt, columns) — a calculus compiler
that differentiates a single expression and hands back the derivative as SQL
text, for when you want the update rule as a string to embed yourself.

How it works

The core is src/autograd.rs, a symbolic differentiation engine over DataFusion
logical Expr trees. The design mirrors JAX's per-primitive rule registry
(defjvp and friends): every node type has a differentiation rule and the chain
rule composes them as the tree is walked. A small 0/1-folding simplifier
keeps the output compact, playing the role of JAX's Zero tangents and
add_tangents. grad and jvp share the chain rule and differ only in their
leaf rule (a one-hot seed vs. an arbitrary per-input tangent); vjp scales the
partial by the cotangent.

Rules are implemented for +, -, *, /, the unary chain rule for
sin/cos/tan, asin/acos/atan, exp/ln/log2/log10/sqrt,
sinh/cosh/tanh, abs, and power() with a constant base or exponent. An
unsupported node or function returns a clear NotImplemented error rather than a
silently wrong derivative.

The markers are differentiated as a SQL source-to-source rewrite, before the
query is planned: XarrayContext.sql() hands a query containing a marker to
the native rewrite_grad_sql, which parses the statement with sqlparser,
differentiates each grad/jvp/vjp call, and renders the derivative back into
the query in place. The result is ordinary SQL the stock datafusion-python
context plans and executes directly. Running before binding is what makes every
query shape the parser accepts — recursive CTEs, DML, subqueries — work
uniformly, with no special cases per plan type.

This needs no DataFusion fork and no custom datafusion-python wheel: it runs
against the published package. There is no Substrait round-trip and no protoc
build dependency.

Things to know

• Marker arguments use unqualified column names (e.g. grad(y - a*x - b, a)),
  matching how the differentiation reads its variable. This is the one
  consequence of differentiating syntactically, before binding.
• Put grad inside an aggregate (AVG(grad(f, x))), not outside. The
  transposed form grad(SUM(f), x) is rejected by SQL's own scoping, since the
  per-row column is gone after aggregation.

Tests

Covered end to end by the Python suite (tests/test_autograd.py) — derivatives
checked against numpy analytics, gradient descent and a recursive-CTE training
loop converging to their closed-form solutions, jvp/vjp agreement, multi-
input Jacobians, higher-order derivatives, schema-qualified tables — and by the
Rust unit tests for the differentiation rules and the SQL rewrite.

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

[alxmrs]

Any ideas for what a more SQL native update would be? What if we had a parameter table and added new rows per each update?

[alxmrs]

Great idea — implemented it that way. The optimiser trajectory is now a params(step, a, b) table that grows one generation per step, and the update itself is computed in SQL by descending from the current row:

WITH cur AS (SELECT a, b FROM params WHERE step = :k)
SELECT cur.a - lr * AVG(grad(loss, a)) AS a,
       cur.b - lr * AVG(grad(loss, b)) AS b
FROM d CROSS JOIN cur
GROUP BY cur.a, cur.b

So each step appends the next generation, and the whole loss curve is a single GROUP BY over the history joined to the data — the trajectory is a relation you can query. One honest caveat: xarray-backed tables are read-only to SQL, so "append a row" is done by re-registering the grown history; a true in-place INSERT would need a mutable table provider. (I also tried WHERE step = (SELECT MAX(step) ...), but scalar subqueries don't convert to Substrait yet, so the loop passes the current step as a literal.)

I moved the runnable demos off this branch to keep it focused on the engine — grad_descent.py (with this change) now lives on the stacked claude/xarray-sql-era5-demo branch, and the MNIST trainer on claude/xarray-sql-mnist-demo stacked above it.

---

Generated by Claude Code