Reuse gradient plans and reduce Dataset registration overhead

tests/test_ad.py

@@ -136,9 +136,54 @@ def test_a_mention_in_a_literal_is_not_a_call():
     assert rows == [{"s": "grad(x)", "jvp(v)": 1}]
 
 
-def test_prepare_refuses_a_gradient(problem):
-    ctx, *_ = problem
-    with pytest.raises(NotImplementedError, match="grad"):
-        ctx.prepare(
-            f"WITH {LOSS} SELECT feature, grad(loss, w) AS g FROM weights"
+def test_context_prepare_runs_gradient_steps_and_reads_writes(problem):
+    ctx, x, w, y = problem
+    step = ctx.prepare(
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature"
+    )
+    old = step()
+    original = old["w"].to_numpy().copy()
+    replacement = xr.Dataset(
+        {"w": ("feature", w.copy() * 2)}, coords={"feature": np.arange(3)}
+    )
+    got = step(weights=replacement)
+    np.testing.assert_allclose(
+        got["w"], 2 * x.T @ (x @ replacement.w.values - y)
+    )
+    np.testing.assert_array_equal(old["w"], original)
+    with pytest.raises(ValueError, match="coordinates"):
+        step(weights=replacement.assign_coords(feature=[3, 4, 5]))
+
+
+def test_prepared_gradients_read_new_values_and_keep_old_results(problem):
+    ctx, x, w, y = problem
+    statements = [
+        f"WITH {LOSS} SELECT feature, grad(loss, w) AS w FROM weights ORDER BY feature",
+        f"WITH {LOSS} SELECT feature, w - 0.01 * grad(loss, w) AS w FROM weights ORDER BY feature",
+    ]
+    run = ad.prepare(ctx, statements)
+    [first, _] = run()
+    original = first["w"].to_numpy().copy()
+    for scale in (0.5, -1.0, 2.0):
+        values = xr.Dataset(
+            {"w": ("feature", w.copy() * scale)},
+            coords={"feature": np.arange(3)},
         )
+        gradient, step = run(weights=values)
+        want = 2 * x.T @ (x @ values.w.values - y)
+        np.testing.assert_allclose(gradient["w"], want)
+        np.testing.assert_allclose(step["w"], values.w.values - 0.01 * want)
+        np.testing.assert_array_equal(first["w"], original)
+    # Regular ddx preparation resolves a replacement table by name.
+    ctx.deregister_table("weights")
+    ctx.from_dataset("weights", values * 2)
+    [gradient, _] = run()
+    np.testing.assert_allclose(
+        gradient["w"], 2 * x.T @ (x @ (values.w.values * 2) - y)
+    )
+
+
+def test_prepared_statements_require_a_sequence(problem):
+    ctx, *_ = problem
+    with pytest.raises(TypeError, match="list"):
+        ad.prepare(ctx, "SELECT 1")

tests/test_prepare.py

@@ -9,6 +9,11 @@
 
 from xarray_sql import Context
 
+try:
+    from einfold.datafusion import EinFold
+except ImportError:
+    EinFold = None
+
 # The build side of the join, `a`, is what a reused hash join would keep.
 JOIN = """
 SELECT b.i, SUM(t.v2 * b.w) AS s
@@ -17,6 +22,26 @@
 """
 
 
+@pytest.mark.parametrize("extension", [False, True])
+def test_join_key_data_columns_can_change_their_range(extension):
+    ctx = Context()
+    if extension:
+        if EinFold is None:
+            pytest.skip("requires the optional EinFold planner")
+        ctx = ctx.with_extensions(EinFold())
+    a = xr.Dataset({"key": ("i", np.arange(128)), "v": ("i", np.ones(128))})
+    b = xr.Dataset({"key": ("j", np.arange(256)), "w": ("j", np.arange(256.0))})
+    ctx.from_dataset("a", a)
+    ctx.from_dataset("b", b)
+    sql = "SELECT a.i,b.j,a.v+b.w AS s FROM a JOIN b ON a.key=b.key ORDER BY a.i,b.j"
+    query = ctx.prepare(sql)
+    for start in (0, 128, 16):
+        updated = a.copy(deep=True)
+        updated["key"][:] = np.arange(start, start + 128)
+        rows = query(a=updated)
+        assert rows.equals(ctx.sql(sql).to_arrow_table())
+
+
 def _a(v):
     return xr.Dataset(
         {"v": (("i", "j"), v)},
@@ -262,3 +287,50 @@ def test_a_quoted_table_replaced_in_sql_counts_too():
     ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect()
     with pytest.raises(ValueError, match="registered again"):
         query()
+
+
+def test_extension_handle_keeps_dataset_metadata_and_invalidation(ctx):
+    want = ctx.sql("SELECT * FROM a").to_dataset()
+    derived = ctx.with_extensions()
+    query = derived.prepare("SELECT i, j, v FROM a ORDER BY i, j")
+    xr.testing.assert_identical(
+        derived.sql("SELECT * FROM a").to_dataset(), want
+    )
+    ctx.deregister_table("a")
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+@pytest.mark.parametrize(
+    "chunks", [{"i": 1, "j": 1}, {"i": 1}, {"i": 4, "j": 5}]
+)
+def test_fine_chunk_registration_keeps_rows_metadata_and_writes(chunks):
+    values = np.arange(129 * 33, dtype=np.float64).reshape(129, 33)
+    ds = _a(values)
+    ds.attrs["experiment"] = "fine chunks"
+    ctx = Context()
+    ctx.from_dataset("a", ds, chunks=chunks)
+    sql = "SELECT i,j,v FROM a ORDER BY i,j"
+    xr.testing.assert_identical(ctx.sql(sql).to_dataset(), ds)
+    run = ctx.prepare(sql)
+    replacement = ds * 2
+    rows = run(a=replacement)
+    np.testing.assert_array_equal(rows["v"], replacement.v.values.reshape(-1))
+
+
+@pytest.mark.parametrize("threads", [1, 60])
+def test_thread_target_is_explicit_and_keeps_xarray_support(threads):
+    ctx = Context(threads=threads)
+    assert ctx.sql(
+        "SELECT value FROM information_schema.df_settings "
+        "WHERE name='datafusion.execution.target_partitions'"
+    ).to_pylist() == [{"value": str(threads)}]
+    ctx.from_dataset("a", _a(np.ones((4, 3))))
+    assert ctx.prepare("SELECT SUM(v) AS s FROM a")().to_pylist() == [
+        {"s": 12.0}
+    ]
+
+
+def test_thread_target_must_be_positive():
+    with pytest.raises(ValueError, match="positive"):
+        Context(threads=0)

xarray_sql/ad.py

@@ -43,7 +43,11 @@
 
 [sql_all][xarray_sql.ad.sql_all] runs several statements that take
 ``grad`` of the same loss for the price of one backward pass, which is
-how a training step updates each of a model's parameter tables. For the
+how a training step updates each of a model's parameter tables. For repeated
+training steps, [prepare][xarray_sql.ad.prepare] retains ddx's
+programs between calls and materializes each statement's result.
+[Context.prepare][xarray_sql.Context.prepare] handles one such statement;
+``ad.prepare`` handles several statements sharing a loss. For the
 programs underneath (``grad``, ``vjp``, ``jvp`` and ``run``, to build a
 backward pass once and run it every step), use ``ddxdb.ad`` directly;
 they take a ``Context`` like any other ``SessionContext``.
@@ -52,12 +56,13 @@
 from __future__ import annotations
 
 import re
+import threading
 from collections.abc import Sequence
 
 from datafusion import SessionContext
 from ddxdb import ad as _ddx
 
-__all__ = ["STOP_GRADIENT", "calls_ad", "register", "sql", "sql_all"]
+__all__ = ["STOP_GRADIENT", "calls_ad", "prepare", "register", "sql", "sql_all"]
 
 STOP_GRADIENT = _ddx.STOP_GRADIENT
 """The SQL name of the stop-gradient function, ``ddx_stop_gradient``."""
@@ -138,3 +143,44 @@ def sql_all(ctx: SessionContext, statements: Sequence[str]) -> list:
     frames: list = _ddx.sql_all(ctx, list(statements))
     wrap = getattr(ctx, "_wrap", None)
     return [wrap(f) for f in frames] if wrap else frames
+
+
+def prepare(ctx: SessionContext, statements: Sequence[str]):
+    """Prepare SQL gradient statements for repeated execution.
+
+    ddx retains the differentiation programs and checks table schemas,
+    views and session settings on each call. This avoids rebuilding those
+    programs in a training loop while allowing tables to be registered
+    again. Several statements can share a loss, as with ``sql_all``::
+
+        step = xql.ad.prepare(ctx, [
+            f"WITH loss AS ({LOSS}) "
+            "SELECT i, val - 0.01 * grad(loss, val) AS val FROM w ORDER BY i"
+        ])
+        [rows] = step()
+
+    On a ``Context``, a call may first write Dataset values into registered
+    tables, using the same validation as ``Context.prepare``::
+
+        [rows] = step(w=new_weights)
+
+    Returns a callable producing one materialized ``pyarrow.Table`` per
+    statement. Results remain valid after later calls. Writes and execution
+    are serialized for calls to this prepared function. Other users of the
+    context must coordinate concurrent writes themselves.
+    """
+    if isinstance(statements, str):
+        raise TypeError("statements is a list of SQL statements")
+    run = _ddx.prepare(ctx, list(statements))
+    lock = threading.Lock()
+
+    def execute(**tables):
+        with lock:
+            if tables:
+                write = getattr(ctx, "_write", None)
+                if write is None:
+                    raise TypeError("Dataset writes require xarray_sql.Context")
+                write(tables)
+            return [frame.to_arrow_table() for frame in run()]
+
+    return execute

xarray_sql/reader.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import itertools
+import math
 import os
 import sys
 from collections.abc import Callable, Iterable, Iterator
@@ -21,9 +22,9 @@
 import xarray as xr
 
 from .df import (
+    DEFAULT_BATCH_SIZE,
     Block,
     Chunks,
-    DEFAULT_BATCH_SIZE,
     _block_len,
     _block_metadata,
     _block_slices_from_resolved,
@@ -296,11 +297,10 @@ def _in_memory_partitions(
     # isn't contiguous in memory, and flattening it would copy. More
     # partitions than cores only adds overhead, since an in-memory table has
     # no partitions to prune.
+    # Count the Cartesian chunk grid without constructing slices that this
+    # in-memory path never reads, even for very finely chunked inputs.
     count = min(
-        sum(
-            1
-            for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
-        ),
+        math.prod(len(axis) for axis in resolve_chunks(ds, chunks).values()),
         os.cpu_count() or 1,
     )
     # `per_partition` batches of `size` rows each make `count` partitions,
@@ -324,7 +324,7 @@ def read_xarray_table(
     _iteration_callback: (
         Callable[[Block, list[str] | None], None] | None
     ) = None,
-) -> "LazyArrowStreamTable":
+) -> LazyArrowStreamTable:
     """Create a lazy DataFusion table from an xarray Dataset.
 
     This is the simplest way to register xarray data with DataFusion.