ddxdb

Write calculus directly in SQL and let the database evaluate the derivative, row by row, alongside everything else:

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

grad and jvp are markers, not row functions. They are rewritten away into ordinary derivative SQL before the engine sees them, so what runs is a plain expression — the relational equivalent of jax.vmap(jax.grad(f)), with the rows as the batch dimension.

This is the Python distribution of ddx, a thin wrapper over the ddx-core engine.

Install

pip install ddxdb                  # everything below except Context
pip install "ddxdb[datafusion]"    # + the DataFusion Context

rewrite_sql is the whole library

Text in, text out — so it works with any engine that accepts SQL. Pass the result wherever you would have passed the original:

import ddxdb

ddxdb.rewrite_sql("SELECT grad(sin(x), x) AS d FROM t")
# 'SELECT (cos(x)) AS d FROM t'

con.sql(ddxdb.rewrite_sql(q, "duckdb"))        # DuckDB
session.sql(ddxdb.rewrite_sql(q, "spark"))     # Spark
ctx.sql(ddxdb.rewrite_sql(q))                  # DataFusion

Accepted dialects: generic, datafusion, postgres, ansi, snowflake, oracle, duckdb, mysql, sqlite, bigquery, redshift, hive, spark, databricks, mssql, teradata, clickhouse.

Pick the one that matches the engine you will run on, not just the one that parses your SQL. The dialect also decides which column an identifier names, and engines disagree three ways:

Getting this wrong does not raise. grad("X" * "X", X) is 2X on Snowflake and 0 on Postgres — both correct, for different engines — so ddx keeps a table rather than a default, and refuses a dialect whose rule it has not established.

Because the rewrite happens in your process, on your connection, it sees your temp tables, session settings and open transaction — anything the query itself could see. A rewrite performed inside the database, on a connection of its own, would not.

Context, for DataFusion

A real SessionContext subclass whose .sql() rewrites first — every inherited method, property and constructor argument works unchanged:

ctx = ddxdb.Context()
ctx.sql("SELECT grad(x * x, x) AS d FROM t").collect()      # → 2x

It lives in ddxdb.datafusion (a subclass needs its base class at import time, so it cannot sit beside rewrite_sql without dragging DataFusion in) and is re-exported as ddxdb.Context, imported on first use. import ddxdb still needs no engine.

There is sugar for DataFusion and not for other engines because DataFusion is ddx's integration target. Everything else uses the one-liner above, which is why there are no per-engine helpers here to drift out of date.

What you can write

+ - * /; the chain rule for the trig / inverse-trig / exp / log / hyperbolic set plus abs; power with a constant base or exponent. Higher order falls out of nesting — grad(grad(f, x), x) just works. Differentiating through an aggregate is linearity, so the marker goes inside it, which is what makes a gradient-descent step expressible in SQL:

SELECT theta - 0.01 * AVG(grad(loss, theta)) FROM batch

A marker rewrites in place, so it is legal anywhere a scalar expression is — including inside a recursive CTE, which is how a whole training loop fits in one query.

Two other functions

ddxdb.differentiate_sql("x * y", "x")     # 'y' — the derivative as text

The escape hatch, for assembling SQL where a marker cannot reach — inside a recursive term you are building programmatically, or a query some other tool emits. Everything else should use rewrite_sql.

ddxdb.supported_functions()               # ['abs', 'acos', 'asin', ...]

The unary functions ddx has a rule for, read from the engine rather than restated. Note that a name being present does not by itself make an expression differentiable — the surrounding constructs matter too — so catching the typed error below remains the general answer to "can ddx handle this?".

Errors are typed

An unsupported construct is always an error, never a silently wrong number — this is a numerical-correctness library, and a plausible-looking wrong derivative is the worst thing it could produce. The kind of failure is a class, so you can catch the one you can act on:

try:
    ddxdb.rewrite_sql(query)
except ddxdb.UnsupportedExpression:
    ...   # no rule for something in there — fall back
except ddxdb.AmbiguousColumn:
    ...   # the query needs a qualifier — a fix the caller makes

All of them derive from ddxdb.DdxError. The full set is UnsupportedExpression, InvalidMarker, AmbiguousColumn, ProjectionBoundary and SqlParseError, plus NotScalar, UnknownColumn and InvalidColumn from whole-query grad.

Gradients of whole queries: grad(loss, table.column)

grad(expr, column) differentiates one expression. Training a model needs the gradient of a whole query, and that is grad too, in a FROM clause: the gradient of the loss a CTE computes, as a relation shaped like the table.

from ddxdb import Context

ctx = Context()
# ... register x(sample, inp, val), w(inp, out, val), y(sample, out, val) ...

ctx.sql("""
WITH h AS (
  SELECT x.sample, w.out, tanh(SUM(x.val * w.val)) AS val
  FROM x JOIN w ON x.inp = w.inp GROUP BY x.sample, w.out),
loss AS (
  SELECT SUM(power(h.val - y.val, 2)) AS l
  FROM h JOIN y ON h.sample = y.sample AND h.out = y.out)
SELECT w.inp, w.out, w.val - 0.1 * g.val AS val
FROM w JOIN grad(loss, w.val) g ON w.inp = g.inp AND w.out = g.out
""")

That is one SGD step: params - lr * grad(loss)(params), written as a join. Nothing in the loss is labelled for ddx; the one function it gives a meaning to is ddx_stop_gradient(x), JAX's lax.stop_gradient. On a plain SessionContext, ddxdb.ad.sql(ctx, statement) does the same, and ddxdb.ad.sql_all runs several statements that take grad of one loss for the price of one backward pass. ddxdb.ad.grad and ddxdb.ad.vjp give the underlying program. A query that is not a loss raises NotScalar, a column the loss does not read raises UnknownColumn, and a column that cannot be differentiated (not a float, or in a table whose rows do not have unique dims) raises InvalidColumn. This needs DataFusion.

Another engine needs no DataFusion: ddxdb.grad_plan(plan_bytes, wrt) differentiates the serialized Substrait plan its producer writes, and ddxdb.run(backend, program) runs the result on any object with four methods, select_all, returns_rows, materialize and drop_table (the ddxdb.Backend protocol). Both are ddx-ad's own Rust, the same code the Rust adapter runs. Pass namespace="__ddx_mine_" to get the same program from the same plan every time.

One thing to know

grad does not see through a CTE or a view. Differentiation stops at column references, so a column computed upstream is a constant to it:

WITH v AS (SELECT x, sin(x) AS s FROM t)
SELECT grad(s * x, x) FROM v       -- ds/dx is treated as 0

That is defensible relational semantics and a real trap, so ddx refuses the worst case rather than quietly dropping the term: referencing a computed CTE alias as a non-wrt term raises ProjectionBoundary and tells you to differentiate inside the CTE instead. Differentiating with respect to such an alias is fine — every occurrence is then the differentiation leaf, and grad(s * s, s) is exactly 2s.

Development

pip install maturin pytest
maturin develop --uv
python -m pytest tests/

Building compiles protoc from source for the substrait crate (it needs cmake and a C++ compiler), which takes a few minutes the first time.

License

Licensed under Apache-2.0, the same as the rest of ddx.