Crate ddx_datafusion 

ddx-datafusion — the DataFusion adapter for ddx.

Write calculus directly in SQL and let DataFusion evaluate the derivative per row — the relational equivalent of jax.vmap(jax.grad(f)):

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

All the calculus lives in ddx_core; this crate only connects it to an engine. It offers the same rewrite by two routes:

	install (Path B)	ddx_sql (Path A)
How	in-engine AnalyzerRule on the bound plan	rewrite the SQL text first
Call style	bare grad(), anywhere	ddx_sql(&ctx, sql)
Works with	SQL and the DataFrame API	SQL strings
Column identity	resolved by the planner	syntactic (guards may fire)
Correlated subqueries	✗ (loud error)	✓
Errors surface at	collect()	the ddx_sql call

Prefer install. Because it runs after binding, columns arrive already resolved, so the qualification-ambiguity errors a pre-binding text rewrite must raise simply cannot occur.

Reach for ddx_sql when a marker sits inside a correlated subquery. That is the one query shape Path B genuinely cannot carry: the bridge re-plans the derivative against the subquery’s own inputs, and an outer reference does not survive that. Path B detects it and says so.

Recursive CTEs are not such a shape: install carries a marker in a recursive term perfectly well, because LogicalPlan::RecursiveQuery is an ordinary node with ordinary inputs.

Where the two paths genuinely differ

They drive the same ddx_core engine, so they agree on the calculus. They do not always agree on what may be the wrt, because they disagree about what a “column” is:

SELECT grad(sum(x) * sum(x), sum(x)) AS d FROM t

Path A refuses this — syntactically sum(x) is a function call, not a bare column, and the wrt must be a bare column. Path B answers 2·sum(x), because by the time it sees the plan the planner has already lowered the aggregate to the bound column sum(t.x), and differentiating with respect to a column is exactly what it does.

Path B’s wrt is any column of the node’s input schema, including planner-derived ones — aggregate outputs, window outputs, computed aliases. That is deliberate. Differentiating with respect to a computed alias is already a supported and correct operation — grad(s*s, s) is 2s — and an aggregate output is the same shape one level down, so refusing it would contradict the case ddx already accepts.

Path B

let ctx = SessionContext::new();
ddx_datafusion::install(&ctx);

ctx.sql("CREATE TABLE t AS VALUES (1.0), (2.0), (3.0)").await?.collect().await?;

// bare grad() — no wrapper
let df = ctx.sql("SELECT grad(column1 * column1, column1) AS d FROM t").await?;

What it supports

Whatever ddx_core supports: + - * /; the unary chain rule for the trig / inverse-trig / exp / log / hyperbolic set plus abs; power with a constant base or exponent; higher-order via nesting; through-aggregate via linearity (AVG(grad(loss, theta))). Anything else is a typed error, never a silently-wrong number. Errors from the engine arrive as DataFusionError::External boxing a ddx_core::DiffError, so you can downcast and match on the variant.

Re-exports

pub use ddx_core;

Structs

DdxAnalyzer

The ddx analyzer rule: rewrites grad/jvp markers away before execution.

Constants

GRAD

The name of the gradient marker, as written in SQL.

JVP

The name of the forward-mode (directional derivative) marker.

Functions

ddx_sql

Rewrite grad/jvp markers in sql and run the result on ctx — the one-liner form of the text rewrite.

ddx_sql_with

ddx_sql driven by a caller-supplied engine — use this when you have registered custom differentiation rules via Ddx::register.

grad_udf

The grad(expr, column) marker: d(expr)/d(column).

install

Install ddx on ctx: register the grad/jvp marker UDFs and the analyzer rule that rewrites them away (Path B).

install_with

install with a caller-configured analyzer — use this to pick up custom differentiation rules (see DdxAnalyzer::with_engine).

jvp_udf

The jvp(expr, column, tangent) marker: d(expr)/d(column) · tangent.

rewrite_sql

Rewrite the markers in sql and return the derivative SQL as text, without running it.

rewrite_sql_with

rewrite_sql driven by a caller-supplied engine.