Expose grad() in XarrayContext SQL via the Substrait rewrite

Wire the autograd surface into XarrayContext so users can write calculus
directly in SQL:

    ctx.sql("SELECT grad(sin(val), val) AS d_val, sin(val) AS val FROM t")

On construction the context registers the `grad` marker UDF so such queries
parse and plan. XarrayContext.sql() detects `grad(` (a cheap regex gate so
ordinary queries are untouched) and routes through _sql_with_autograd: it
plans the query, produces the logical plan as Substrait, calls the native
grad_rewrite to differentiate every grad(expr, column) symbolically, then
consumes the rewritten Substrait back into an executable DataFrame.

Table scans are resolved by name on the consume side, so _table_schemas()
passes the (name, schema) of each registered table to the rewrite. Schema-
qualified tables (mixed-dimension datasets) are skipped for now and noted as
a follow-up.

Adds tests/test_autograd.py covering sin/cos, product and quotient rules,
power, exp, the non-grad passthrough, and a clear error for unsupported
functions — all checked against numpy analytic derivatives. Existing SQL
tests still pass.

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

Author: claude

tests/test_autograd.py

@@ -0,0 +1,76 @@
+"""Tests for the SQL autograd surface: ``SELECT grad(expr, column) ...``.
+
+These exercise the full path — XarrayContext.sql() -> Substrait -> native
+grad_rewrite -> Substrait -> execute — and compare results against analytic
+derivatives computed with numpy.
+"""
+
+import numpy as np
+import pytest
+import xarray as xr
+
+import xarray_sql as xql
+
+
+@pytest.fixture
+def ctx():
+    val = np.linspace(0.1, 3.0, 16)
+    ds = xr.Dataset(
+        {"val": (("i",), val)},
+        coords={"i": np.arange(16)},
+    )
+    context = xql.XarrayContext()
+    context.from_dataset("t", ds, chunks={"i": 5})
+    return context
+
+
+def _ordered(df, key="i"):
+    """Collect a result DataFrame into a dict of column -> numpy array, sorted
+    by the integer key column so comparisons are index-aligned."""
+    pdf = df.to_pandas().sort_values(key)
+    return {c: pdf[c].to_numpy() for c in pdf.columns}
+
+
+def test_grad_sin_is_cos(ctx):
+    val = np.linspace(0.1, 3.0, 16)
+    res = _ordered(ctx.sql("SELECT i, grad(sin(val), val) AS d FROM t"))
+    np.testing.assert_allclose(res["d"], np.cos(val))
+
+
+def test_grad_product_rule(ctx):
+    val = np.linspace(0.1, 3.0, 16)
+    res = _ordered(ctx.sql("SELECT i, grad(sin(val) * val, val) AS d FROM t"))
+    np.testing.assert_allclose(res["d"], np.cos(val) * val + np.sin(val))
+
+
+def test_grad_exp_equals_value(ctx):
+    val = np.linspace(0.1, 3.0, 16)
+    res = _ordered(
+        ctx.sql("SELECT i, exp(val) AS v, grad(exp(val), val) AS d FROM t")
+    )
+    np.testing.assert_allclose(res["d"], np.exp(val))
+    np.testing.assert_allclose(res["d"], res["v"])
+
+
+def test_grad_quotient_and_power(ctx):
+    val = np.linspace(0.1, 3.0, 16)
+    res = _ordered(
+        ctx.sql(
+            "SELECT i, grad(1.0 / val, val) AS dinv, "
+            "grad(power(val, 3), val) AS dcube FROM t"
+        )
+    )
+    np.testing.assert_allclose(res["dinv"], -1.0 / val**2)
+    np.testing.assert_allclose(res["dcube"], 3.0 * val**2)
+
+
+def test_non_grad_query_is_unaffected(ctx):
+    # Queries without grad() bypass the rewrite and behave normally.
+    res = _ordered(ctx.sql("SELECT i, val FROM t"))
+    np.testing.assert_allclose(res["val"], np.linspace(0.1, 3.0, 16))
+
+
+def test_unsupported_function_raises(ctx):
+    # atan2 has no derivative rule yet -> a clear error, not a wrong answer.
+    with pytest.raises(Exception):
+        ctx.sql("SELECT grad(atan2(val, val), val) AS d FROM t").to_pandas()

xarray_sql/sql.py

@@ -1,13 +1,22 @@
+import re
+
+import pyarrow as pa
 import xarray as xr
-from datafusion import SessionContext
+from datafusion import SessionContext, udf
 from datafusion.catalog import Schema
+from datafusion.substrait import Consumer, Producer, Serde
 from collections import defaultdict
 
+from . import _native
 from . import cftime as cft
 from .df import Chunks
 from .ds import XarrayDataFrame
 from .reader import read_xarray_table
 
+# Matches a call to the autograd marker function ``grad(`` (case-insensitive),
+# used as a cheap gate so ordinary queries skip the Substrait round-trip.
+_GRAD_CALL = re.compile(r"\bgrad\s*\(", re.IGNORECASE)
+
 
 class XarrayContext(SessionContext):
     """A datafusion `SessionContext` that also supports `xarray.Dataset`s."""
@@ -21,6 +30,24 @@ def __init__(self, *args, **kwargs):
         # in SQL (e.g. ``"air"`` for a uniform-dim Dataset, or
         # ``"era5.surface"`` for one entry from a multi-dim-group split).
         self._registered_datasets: dict[str, xr.Dataset] = {}
+        self._register_autograd_udfs()
+
+    def _register_autograd_udfs(self) -> None:
+        """Register the ``grad`` marker UDF used by the autograd rewrite.
+
+        ``grad(expr, column)`` is a *marker*: it lets queries parse and plan
+        with the differentiation request intact. It is never executed — the
+        Substrait rewrite in :meth:`sql` replaces every ``grad(...)`` with the
+        symbolic derivative of ``expr`` before execution.
+        """
+        marker = udf(
+            lambda expr, column: expr,
+            [pa.float64(), pa.float64()],
+            pa.float64(),
+            "immutable",
+            "grad",
+        )
+        self.register_udf(marker)
 
     def from_dataset(
         self,
@@ -174,9 +201,50 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         Returns:
             An :class:`XarrayDataFrame` wrapping the DataFusion DataFrame.
         """
-        inner = super().sql(query, *args, **kwargs)
+        if _GRAD_CALL.search(query):
+            inner = self._sql_with_autograd(query, *args, **kwargs)
+        else:
+            inner = super().sql(query, *args, **kwargs)
         return XarrayDataFrame(inner, templates=self._registered_datasets)
 
+    def _sql_with_autograd(self, query: str, *args, **kwargs):
+        """Plan ``query``, rewrite ``grad(...)`` calls, return a DataFrame.
+
+        The differentiation engine lives in the native (Rust) extension and
+        operates on DataFusion logical expressions. Since that extension links
+        its own copy of DataFusion, the plan crosses the boundary as Substrait:
+        we produce the logical plan as Substrait, hand it to ``grad_rewrite``
+        (which differentiates every ``grad(expr, column)`` symbolically), then
+        consume the rewritten Substrait back into an executable DataFrame.
+        """
+        plan = super().sql(query, *args, **kwargs).logical_plan()
+        substrait_plan = Producer.to_substrait_plan(plan, self)
+        rewritten = _native.grad_rewrite(
+            substrait_plan.encode(), self._table_schemas()
+        )
+        new_plan = Consumer.from_substrait_plan(
+            self, Serde.deserialize_bytes(rewritten)
+        )
+        return self.create_dataframe_from_logical_plan(new_plan)
+
+    def _table_schemas(self) -> list[tuple[str, pa.Schema]]:
+        """Return ``(name, schema)`` for each registered table.
+
+        The Substrait consumer in ``grad_rewrite`` resolves table scans by
+        name, so it needs the schema of every table the plan might reference.
+        Only metadata is read here — never the underlying data.
+        """
+        schemas = []
+        for name in self._registered_datasets:
+            try:
+                schemas.append((name, self.table(name).schema()))
+            except Exception:
+                # Schema-qualified tables (mixed-dimension datasets) aren't
+                # resolvable by a bare name yet; skip rather than fail the
+                # whole query. grad() over those is a follow-up.
+                continue
+        return schemas
+
 
 def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]:
     """Group variables in the dataset based on shared dims.