Support grad/jvp/vjp on schema-qualified tables

Mixed-dimension datasets register as schema-qualified tables (e.g.
era5.surface / era5.time_x_level). The autograd rewrite consumes the plan in
a throwaway context that registers an empty table per scanned name, but
register_table("era5.time_x", ...) failed with "failed to resolve schema:
era5" because the namespace did not exist.

Add ensure_schema(): before registering each table, parse its name into a
TableReference and, for qualified names, create the schema namespace
(MemorySchemaProvider) in the default catalog if absent. The Python side
already resolves qualified names via ctx.table(name).schema(); only the Rust
rewrite context needed the namespace.

Tests: a mixed-dimension fixture exercising grad on both the 2D surface and
3D atmosphere tables, against numpy analytic derivatives.

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

Author: claude

src/lib.rs

@@ -55,8 +55,8 @@ use arrow::pyarrow::FromPyArrow;
 use async_stream::try_stream;
 use async_trait::async_trait;
 use datafusion::catalog::streaming::StreamingTable;
-use datafusion::catalog::Session;
-use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
+use datafusion::catalog::{MemorySchemaProvider, Session};
+use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue, TableReference};
 use datafusion::datasource::empty::EmptyTable;
 use datafusion::datasource::TableProvider;
 use datafusion::execution::TaskContext;
@@ -993,6 +993,22 @@ impl LazyArrowStreamTable {
 // Autograd: Substrait-level grad() rewrite
 // ============================================================================
 
+/// Ensure a schema (namespace) exists in the context's catalog, creating an
+/// empty in-memory one if needed. Used so the rewrite context can register
+/// schema-qualified tables (e.g. `era5.surface`) that mixed-dimension datasets
+/// produce.
+fn ensure_schema(ctx: &SessionContext, catalog: Option<&str>, schema: &str) -> DFResult<()> {
+    // A bare TableReference has no catalog; fall back to DataFusion's default.
+    let catalog_name = catalog.unwrap_or("datafusion");
+    let catalog = ctx
+        .catalog(catalog_name)
+        .ok_or_else(|| DataFusionError::Plan(format!("catalog '{catalog_name}' not found")))?;
+    if catalog.schema(schema).is_none() {
+        catalog.register_schema(schema, Arc::new(MemorySchemaProvider::new()))?;
+    }
+    Ok(())
+}
+
 /// Rewrite `grad(expr, column)` calls in a Substrait plan into their symbolic
 /// derivatives.
 ///
@@ -1037,7 +1053,20 @@ fn grad_rewrite<'py>(
             ))
         })?;
         let provider = Arc::new(EmptyTable::new(Arc::new(schema)));
-        ctx.register_table(name.as_str(), provider).map_err(|e| {
+
+        // Schema-qualified names (e.g. "era5.surface", from a mixed-dimension
+        // dataset) need their namespace to exist before the table can be
+        // registered into this throwaway context.
+        let table_ref = TableReference::from(name.as_str());
+        if let Some(schema_name) = table_ref.schema() {
+            ensure_schema(&ctx, table_ref.catalog(), schema_name).map_err(|e| {
+                pyo3::exceptions::PyValueError::new_err(format!(
+                    "grad_rewrite: failed to create schema for table '{name}': {e}"
+                ))
+            })?;
+        }
+
+        ctx.register_table(table_ref, provider).map_err(|e| {
             pyo3::exceptions::PyValueError::new_err(format!(
                 "grad_rewrite: failed to register table '{name}': {e}"
             ))

tests/test_autograd.py

@@ -133,6 +133,46 @@ def test_vjp_reverse_pullback(ctx_xy):
     np.testing.assert_allclose(res["s"], 3.0 * (np.cos(x) * y))
 
 
+@pytest.fixture
+def ctx_mixed():
+    # A mixed-dimension dataset registers as schema-qualified tables:
+    #   era5.time_x        (surface, 2 dims)
+    #   era5.time_x_level  (atmosphere, 3 dims)
+    rng = np.random.default_rng(1)
+    ds = xr.Dataset(
+        {
+            "sfc": (("time", "x"), rng.uniform(0.5, 2.5, (3, 4))),
+            "atm": (("time", "x", "level"), rng.uniform(0.5, 2.5, (3, 4, 2))),
+        },
+        coords={"time": [0, 1, 2], "x": np.arange(4.0), "level": [0, 1]},
+    )
+    context = xql.XarrayContext()
+    context.from_dataset("era5", ds, chunks={"time": 1})
+    return context, ds
+
+
+def test_grad_on_qualified_surface_table(ctx_mixed):
+    context, ds = ctx_mixed
+    res = _ordered(
+        context.sql(
+            "SELECT time, x, sfc, grad(sin(sfc), sfc) AS d FROM era5.time_x"
+        ),
+        key="sfc",
+    )
+    np.testing.assert_allclose(res["d"], np.cos(res["sfc"]))
+
+
+def test_grad_on_qualified_atmosphere_table(ctx_mixed):
+    context, ds = ctx_mixed
+    res = _ordered(
+        context.sql(
+            "SELECT atm, grad(power(atm, 2), atm) AS d FROM era5.time_x_level"
+        ),
+        key="atm",
+    )
+    np.testing.assert_allclose(res["d"], 2.0 * res["atm"])
+
+
 def test_jvp_and_vjp_agree_for_unit_seed(ctx_xy):
     # Forward (unit tangent) and reverse (unit cotangent) coincide for a
     # scalar output -- both contract the same partial derivative.

xarray_sql/sql.py

@@ -251,11 +251,12 @@ def _table_schemas(self) -> list[tuple[str, pa.Schema]]:
         schemas = []
         for name in self._registered_datasets:
             try:
+                # Names may be bare ("air") or schema-qualified ("era5.surface",
+                # from a mixed-dimension dataset); both resolve here.
                 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.
+                # Be defensive: skip a table we can't introspect rather than
+                # failing the whole query.
                 continue
         return schemas