From 73157f2647554e736d33828536c609cc7f5187f2 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 1 Nov 2026 23:24:43 -0800 Subject: [PATCH] Reuse gradient programs and reduce Dataset registration overhead Support prepared SQL gradients, arithmetic chunk-grid counts and direct Arrow registration. Preserve Context metadata through extension installation, expose target partitions and guard physical reuse against value-dependent filters. Co-Authored-By: GPT-6 --- tests/test_ad.py | 55 +++++++++++++++++-- tests/test_prepare.py | 72 +++++++++++++++++++++++++ xarray_sql/ad.py | 50 +++++++++++++++++- xarray_sql/reader.py | 12 ++--- xarray_sql/sql.py | 120 +++++++++++++++++++++++++++++++++++------- 5 files changed, 278 insertions(+), 31 deletions(-) diff --git a/tests/test_ad.py b/tests/test_ad.py index 0f868bc7..df3d2675 100644 --- a/tests/test_ad.py +++ b/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") diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 8cc356bc..1b8b2b53 100644 --- a/tests/test_prepare.py +++ b/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) diff --git a/xarray_sql/ad.py b/xarray_sql/ad.py index d0ff2c8f..eb12555a 100644 --- a/xarray_sql/ad.py +++ b/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 diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 8ec88bcf..db7c82c9 100644 --- a/xarray_sql/reader.py +++ b/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. diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index f64166a4..bf801903 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,6 +1,8 @@ import re import time import uuid +from collections.abc import Callable +from types import ModuleType import numpy as np import pyarrow as pa @@ -9,7 +11,6 @@ from datafusion.catalog import Schema from datafusion.expr import TableScan from datafusion.plan import ExecutionPlan, LogicalPlan -from types import ModuleType from . import cftime as cft from .df import ( @@ -38,14 +39,30 @@ except ImportError: # pragma: no cover - depends on the environment _ad = None +_DYNAMIC_FILTERS = tuple( + f"datafusion.optimizer.enable{kind}_dynamic_filter_pushdown" + for kind in ("", "_join", "_topk", "_aggregate") +) + class Context(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s. Also available under its original name, ``XarrayContext``. + ``threads=`` sets DataFusion's target parallel partitions, overriding + that setting in an explicit ``config``. The default is the CPU count. + It is a concurrency target, rather than a fixed number of OS threads. """ - def __init__(self, config: SessionConfig | None = None, *args, **kwargs): + def __init__( + self, + config: SessionConfig | None = None, + *args, + threads: int | None = None, + **kwargs, + ): + if threads is not None and threads < 1: + raise ValueError("threads must be positive") if config is None: # Split a query's work across partitions only where enough rows # reach it, by DataFusion's estimate. DataFusion otherwise splits @@ -65,6 +82,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): "true", ) ) + if threads is not None: + config = config.with_target_partitions(threads) super().__init__(config, *args, **kwargs) # Track registered xarray Datasets so XarrayDataFrame can recover # defaults (dimension_columns) and metadata (var/dataset attrs, @@ -78,6 +97,8 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): # registered again from Python was deregistered first; SQL that # replaces or drops a table counts too (see sql()). self._deregistrations: dict[str, int] = {} + self._custom_udfs: set[str] = set() + self._trusted_math_udfs: set[str] = set() # With pyproj installed, every context speaks CRS out of the box: # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform. if _proj is not None: @@ -87,6 +108,41 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): if _ad is not None: _ad.register(self) + def register_udf(self, udf): + """Register a scalar function and track overrides for optional lowering.""" + result = super().register_udf(udf) + self._custom_udfs.add(udf.name.lower()) + self._trusted_math_udfs.discard(udf.name.lower()) + return result + + def register_udaf(self, udaf): + """Track aggregate-function overrides for optional numerical lowering.""" + result = super().register_udaf(udaf) + self._custom_udfs.add(udaf.name.lower()) + self._trusted_math_udfs.discard(udaf.name.lower()) + return result + + def register_udwf(self, udwf): + """Track window-function overrides for optional numerical lowering.""" + result = super().register_udwf(udwf) + self._custom_udfs.add(udwf.name.lower()) + self._trusted_math_udfs.discard(udwf.name.lower()) + return result + + def with_extensions(self, *extensions) -> "Context": + """Install DataFusion extension bundles, retaining xarray support. + + Returns another handle on the same session, with this Context's + Dataset metadata and prepared-query invalidation state shared. + Extension codecs belong to the returned handle, so use it for + subsequent queries, as with DataFusion's own method. + """ + session = super().with_extensions(*extensions) + derived = object.__new__(type(self)) + derived.__dict__.update(self.__dict__) + derived.ctx = session.ctx + return derived + def from_dataset( self, name: str, @@ -214,9 +270,13 @@ def _from_dataset( table = read_xarray_table( input_table, chunks, coord_arrays=coord_arrays ) + register(table_name, table) + elif schema is None: + # Uniform Datasets can register their shared Arrow buffers + # directly, without a temporary catalog entry and lookup. + self.register_record_batches(table_name, partitions) else: - table = self._in_memory_table(partitions) - register(table_name, table) + register(table_name, self._in_memory_table(partitions)) self._maybe_register_cftime_udf(input_table) return self @@ -324,16 +384,11 @@ def prepare(self, query: str) -> "PreparedQuery": A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run the query. - Raises: - NotImplementedError: If the query calls ``grad``, ``vjp`` or - ``jvp``, which only ``sql()`` computes for now. + Queries calling ``grad``, ``vjp`` or ``jvp`` retain ddx's programs + through [xarray_sql.ad.prepare][]. Their first call prepares the + programs; later calls check the catalog and settings and refresh + them if necessary. """ - if _ad is not None and _ad.calls_ad(query): - # TODO(#294): prepare the query's programs with ddx. - raise NotImplementedError( - "prepare() doesn't compute grad, vjp or jvp yet; run the " - "query with sql(), or several with xarray_sql.ad.sql_all()" - ) return PreparedQuery(self, query) def _write(self, writes: dict[str, xr.Dataset]) -> None: @@ -415,10 +470,19 @@ class PreparedQuery: query again. A call raises if a table the query reads has been registered again since it was prepared. Functions such as ``now()`` are evaluated once, when the query is prepared. + + A query calling ``grad``, ``vjp`` or ``jvp`` instead retains ddx's + programs. Its first call prepares them, and subsequent calls refresh + them when schemas, views or settings change. Tables may be registered + again under the same names, as with [xarray_sql.ad.prepare][]. """ def __init__(self, ctx: Context, query: str): self._ctx = ctx + self._gradient: Callable[..., list[pa.Table]] | None = None + if _ad is not None and _ad.calls_ad(query): + self._gradient = _ad.prepare(ctx, [query]) + return start = time.perf_counter() df = SessionContext.sql(ctx, query) plan = df.execution_plan() @@ -441,11 +505,13 @@ def __init__(self, ctx: Context, query: str): } # A plan keeps state from a run (a hash join keeps the table it # built), so each call runs a fresh copy of it, decoded from its - # serialized form. Copying a plan copies its in-memory tables, so - # over large ones it's slower than planning again; a plan that reads - # a lazy table, or is nested deeper than protobuf decodes, can't be - # copied at all. Either way, each call plans the query again. + # serialized form. The default codec copies in-memory tables, so + # over large ones it can be slower than planning again; extensions + # can supply a codec that resets state while retaining input buffers. + # A plan the session's codecs cannot copy is planned again each call. self._plan: ExecutionPlan | None = None + if not _safe_physical_reuse(ctx): + return start = time.perf_counter() try: _copy(ctx, plan) @@ -464,6 +530,9 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table: Returns: The query's rows. """ + if self._gradient is not None: + (rows,) = self._gradient(**tables) + return rows for name in tables: if name not in self._registered: raise ValueError( @@ -495,7 +564,22 @@ def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan: # TODO(#283): reset the plan's state with datafusion-python's binding of # reset_plan_states, once it has one, rather than decode it again. # `plan` itself never runs, so it has no state to leave behind. - return ExecutionPlan.from_bytes(ctx, plan.to_bytes()) + return ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + +def _safe_physical_reuse(ctx: SessionContext) -> bool: + """Dynamic filters retain value-dependent bounds even after a reset.""" + names = ", ".join(repr(name) for name in _DYNAMIC_FILTERS) + try: + rows = SessionContext.sql( + ctx, + f"SELECT value FROM information_schema.df_settings WHERE name IN ({names})", + ).to_arrow_table() + except ValueError: + return False + return rows.num_rows == len(_DYNAMIC_FILTERS) and all( + value == "false" for value in rows["value"].to_pylist() + ) def _tables_read(plan: LogicalPlan) -> set[str]: