ddx_datafusion/
sql.rs

1// SPDX-FileCopyrightText: 2026 Alexander Merose <al@merose.com> & ddx Authors
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! The SQL source-to-source rewrite.
6//!
7//! Rewrite every `grad`/`jvp` marker in the SQL *text* before it reaches the
8//! engine, then hand plain SQL to a stock [`SessionContext`]. This is the
9//! universal path: it runs before planning, so it works for every query shape
10//! the parser accepts — recursive CTEs, DML, subqueries — which is what lets a
11//! whole training loop live in one query. Path B (in-engine, [`crate::analyzer`])
12//! is the more ergonomic one but is bounded by what a `LogicalPlan` can carry.
13
14use datafusion::error::Result;
15use datafusion::prelude::{DataFrame, SessionContext};
16use ddx_core::sqlparser::dialect::GenericDialect;
17use ddx_core::Ddx;
18
19use crate::error::to_df_err;
20
21/// Rewrite `grad`/`jvp` markers in `sql` and run the result on `ctx` — the
22/// one-liner form of the text rewrite.
23///
24/// The context needs no ddx setup at all: no marker UDFs, no analyzer rule. By
25/// the time the engine sees the statement the markers are gone, replaced by
26/// ordinary derivative SQL.
27///
28/// A statement containing no marker is passed through byte-identical and is
29/// never even parsed by ddx, so wrapping every query in
30/// `ddx_sql` costs essentially nothing.
31///
32/// ```
33/// # use datafusion::prelude::SessionContext;
34/// # use ddx_datafusion::ddx_sql;
35/// # #[tokio::main]
36/// # async fn main() -> datafusion::error::Result<()> {
37/// let ctx = SessionContext::new();
38/// ctx.sql("CREATE TABLE t AS VALUES (1.0), (2.0), (3.0)").await?.collect().await?;
39///
40/// // d(x*x)/dx = 2x, computed by the engine as an ordinary column.
41/// let df = ddx_sql(&ctx, "SELECT grad(column1 * column1, column1) AS d FROM t").await?;
42/// let batches = df.collect().await?;
43/// assert_eq!(batches[0].num_rows(), 3);
44/// # Ok(())
45/// # }
46/// ```
47pub async fn ddx_sql(ctx: &SessionContext, sql: &str) -> Result<DataFrame> {
48    ddx_sql_with(ctx, sql, &Ddx::for_datafusion()).await
49}
50
51/// [`ddx_sql`] driven by a caller-supplied engine — use this when you have
52/// registered custom differentiation rules via [`Ddx::register`].
53pub async fn ddx_sql_with(ctx: &SessionContext, sql: &str, ddx: &Ddx) -> Result<DataFrame> {
54    ctx.sql(&rewrite_sql_with(sql, ddx)?).await
55}
56
57/// Rewrite the markers in `sql` and return the derivative SQL as text, without
58/// running it.
59///
60/// Useful for logging what will execute, for feeding another tool, or for the
61/// `ddxdb` Python shim, which does exactly this and then calls a stock
62/// `Context.sql()`.
63pub fn rewrite_sql(sql: &str) -> Result<String> {
64    rewrite_sql_with(sql, &Ddx::for_datafusion())
65}
66
67/// [`rewrite_sql`] driven by a caller-supplied engine.
68///
69/// `GenericDialect` is the parser DataFusion itself uses for SQL, and
70/// [`Ddx::for_datafusion`] supplies the matching identifier-folding policy
71/// (unquoted folds, quoted keeps case) — the two must agree or column matching
72/// silently diverges from the engine's own.
73pub fn rewrite_sql_with(sql: &str, ddx: &Ddx) -> Result<String> {
74    ddx.rewrite_sql(sql, &GenericDialect {}).map_err(to_df_err)
75}
ddx_datafusion/sql.rs

ddx_datafusion/
sql.rs
