From c0c4f7c63ea3ce8811b26c78e6c25dd713126f52 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 08:21:53 -0700
Subject: [PATCH 1/7] Plan a query once and run it many times: Context.prepare

A training loop runs the same query on every step, and over small tables
DataFusion's planning takes longer than running. Context.prepare plans a
query once; calling the PreparedQuery runs that plan on the tables'
current values, after writing any Datasets given into their tables in
place. On a join of two 1,024-row tables a call takes 0.33 ms, against
0.91 ms for ctx.sql.

A plan keeps state from a run (a hash join keeps its build side, a
repartition its channels), and datafusion-python can't reset it (#283).
Each call runs a copy decoded from the plan's serialized form. A plan
that can't be serialized (a lazy table) or decoded (nested too deep), and
a query that calls grad (xqlsystems/ddx#202), are planned again on each
call. A write must keep its NaNs in place (#284).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 docs/performance.md    |  21 +++++
 tests/test_prepare.py  | 119 +++++++++++++++++++++++++++
 xarray_sql/__init__.py |   6 +-
 xarray_sql/sql.py      | 180 ++++++++++++++++++++++++++++++++++++++++-
 4 files changed, 323 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_prepare.py

diff --git a/docs/performance.md b/docs/performance.md
index c461de93..dc135f58 100644
--- a/docs/performance.md
+++ b/docs/performance.md
@@ -177,6 +177,27 @@ it or nothing happens:
 ctx.sql("CREATE OR REPLACE TABLE grid_cube AS ...").collect()
 ```
 
+## Plan once: prepare a query you run many times
+
+DataFusion plans a query each time it runs. For a query over small
+tables, as in a training loop, planning can take longer than running.
+`Context.prepare` plans the query once. A call runs that plan on the
+tables' current values, and can first write new values into a table
+in place:
+
+```python
+step = ctx.prepare(
+    "SELECT x.i, SUM(x.v * w.v) AS y FROM x JOIN w ON x.j = w.j GROUP BY x.i"
+)
+for weights in trajectory:  # Datasets shaped like w's
+    rows = step(w=weights)  # a pyarrow.Table
+```
+
+On a join of two 1,024-row tables, a call takes 0.33 ms, and `ctx.sql`
+takes 0.91 ms (one target partition). A write must keep the table's
+shape, dtypes, coordinates and NaN positions. A plan over a lazy table
+(Dask or a lazily opened file) is planned again on each call for now.
+
 ## Round-trip faster with ORDER BY
 
 Results that arrive **grid-ordered** — sorted by the dimension columns,
diff --git a/tests/test_prepare.py b/tests/test_prepare.py
new file mode 100644
index 00000000..f872a1c1
--- /dev/null
+++ b/tests/test_prepare.py
@@ -0,0 +1,119 @@
+"""Tests for queries planned once and run many times (`Context.prepare`)."""
+
+import numpy as np
+import pytest
+import xarray as xr
+from datafusion import SessionConfig
+
+from xarray_sql import Context
+
+# 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
+FROM (SELECT i, j, 2 * v AS v2 FROM a) t JOIN b ON t.j = b.j
+GROUP BY b.i ORDER BY b.i
+"""
+
+
+def _a(v):
+    return xr.Dataset(
+        {"v": (("i", "j"), v)},
+        coords={"i": np.arange(v.shape[0]), "j": np.arange(v.shape[1])},
+    )
+
+
+def _b(w):
+    return xr.Dataset(
+        {"w": (("i", "j"), w)},
+        coords={"i": np.arange(w.shape[0]), "j": np.arange(w.shape[1])},
+    )
+
+
+@pytest.fixture
+def ctx():
+    rng = np.random.default_rng(0)
+    ctx = Context()
+    ctx.from_dataset("a", _a(rng.normal(size=(4, 3))))
+    ctx.from_dataset("b", _b(rng.normal(size=(4, 3))))
+    return ctx
+
+
+def test_a_prepared_query_returns_the_query_rows(ctx):
+    query = ctx.prepare(JOIN)
+    want = ctx.sql(JOIN).to_arrow_table()
+    assert query().equals(want)
+    assert query().equals(want)
+
+
+def test_a_call_reads_the_values_it_writes(ctx):
+    query = ctx.prepare(JOIN)
+    rng = np.random.default_rng(1)
+    for _ in range(3):
+        a, b = _a(rng.normal(size=(4, 3))), _b(rng.normal(size=(4, 3)))
+        rows = query(a=a, b=b)
+        fresh = Context()
+        fresh.from_dataset("a", a)
+        fresh.from_dataset("b", b)
+        assert rows.equals(fresh.sql(JOIN).to_arrow_table())
+
+
+def test_a_write_changes_the_registered_dataset():
+    v = np.zeros((4, 3))
+    ctx = Context()
+    ctx.from_dataset("a", _a(v))
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    assert query(a=_a(np.ones((4, 3)))).column("s").to_pylist() == [12.0]
+    assert (v == 1).all()
+    assert ctx.sql("SELECT SUM(v) AS s FROM a").to_pydict() == {"s": [12.0]}
+
+
+def test_a_plan_with_several_partitions_runs_again():
+    ctx = Context(SessionConfig().with_target_partitions(4))
+    n = 64
+    ctx.from_dataset("a", _a(np.ones((n, 3))), chunks={"i": n // 4})
+    ctx.from_dataset("b", _b(np.ones((n, 3))), chunks={"i": n // 4})
+    query = ctx.prepare(
+        "SELECT a.i, SUM(v * w) AS s FROM a JOIN b ON a.i = b.i AND a.j = b.j "
+        "GROUP BY a.i"
+    )
+    for scale in (1.0, 2.0):
+        rows = query(a=_a(np.full((n, 3), scale)))
+        assert sorted(rows.column("s").to_pylist()) == [3 * scale] * n
+
+
+def test_a_lazy_table_is_planned_on_each_call():
+    pytest.importorskip("dask")
+    ctx = Context()
+    ctx.from_dataset("a", _a(np.ones((4, 3))).chunk({"i": 2}))
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    assert query().column("s").to_pylist() == [12.0]
+    assert query().column("s").to_pylist() == [12.0]
+    with pytest.raises(ValueError, match="NumPy"):
+        query(a=_a(np.ones((4, 3))))
+
+
+@pytest.mark.parametrize(
+    "new, match",
+    [
+        (_a(np.ones((4, 2))), "coordinates"),
+        (_a(np.ones((4, 3), dtype=np.float32)), "float32"),
+        (_a(np.ones((4, 3))).rename({"v": "u"}), "data variables"),
+        (_a(np.full((4, 3), np.nan)), "NaN"),
+    ],
+)
+def test_a_write_must_keep_the_table_shape(ctx, new, match):
+    query = ctx.prepare(JOIN)
+    want = query()
+    with pytest.raises(ValueError, match=match):
+        query(a=new)
+    assert query().equals(want)
+
+
+def test_a_write_needs_the_table_the_query_was_planned_with(ctx):
+    query = ctx.prepare(JOIN)
+    with pytest.raises(ValueError, match="no Dataset"):
+        query(c=_a(np.ones((4, 3))))
+    ctx.deregister_table("a")
+    ctx.from_dataset("a", _a(np.ones((4, 3))))
+    with pytest.raises(ValueError, match="registered again"):
+        query(a=_a(np.ones((4, 3))))
diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py
index af1a85dd..24a1b6f2 100644
--- a/xarray_sql/__init__.py
+++ b/xarray_sql/__init__.py
@@ -14,6 +14,7 @@
 __all__ = [
     "cftime",
     "Context",
+    "PreparedQuery",
     "XarrayContext",
     "FlightSQLServer",
     "read_xarray_table",
@@ -28,17 +29,18 @@
 ]
 
 try:
-    from .sql import Context, XarrayContext
+    from .sql import Context, PreparedQuery, XarrayContext
 except ModuleNotFoundError as e:
     # DataFusion's Python bindings have no wheels for some platforms
     # (e.g. Pyodide); every other engine adapter still works there.
     if e.name != "datafusion":
         raise
     __all__.remove("Context")
+    __all__.remove("PreparedQuery")
     __all__.remove("XarrayContext")
 
     def __getattr__(name: str):
-        if name in ("Context", "XarrayContext"):
+        if name in ("Context", "PreparedQuery", "XarrayContext"):
             raise AttributeError(
                 f"{name} requires the datafusion package, which is "
                 "not installed."
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 632f9bca..c48ba045 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,9 +1,11 @@
 import uuid
 
+import numpy as np
 import pyarrow as pa
 import xarray as xr
 from datafusion import SessionConfig, SessionContext
 from datafusion.catalog import Schema
+from datafusion.plan import ExecutionPlan
 from types import ModuleType
 
 from . import cftime as cft
@@ -238,9 +240,185 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         Returns:
             An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame.
         """
-        inner = super().sql(query, *args, **kwargs)
+        return self._wrap(super().sql(query, *args, **kwargs))
+
+    def prepare(self, query: str) -> "PreparedQuery":
+        """Plan a query once, to run it many times.
+
+        A training loop runs the same query on every step, and only the
+        values in its tables change. DataFusion plans a query each time it
+        runs, and for a query over small tables, planning can take far
+        longer than running. ``prepare`` plans the query once, and calling
+        the result runs that plan, as ``jax.jit`` traces a function once
+        and runs the compiled function after that. A call can first write
+        new values into the query's tables::
+
+            step = ctx.prepare(
+                "SELECT x.i, SUM(x.v * w.v) AS y "
+                "FROM x JOIN w ON x.j = w.j GROUP BY x.i"
+            )
+            for weights in trajectory:  # Datasets shaped like w's
+                rows = step(w=weights)
+
+        See [PreparedQuery][xarray_sql.PreparedQuery] for what a call
+        can write, and which plans are run again rather than planned again.
+
+        Args:
+            query: A SQL query string.
+
+        Returns:
+            A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run
+            the query.
+        """
+        return PreparedQuery(self, query)
+
+    def _write(self, writes: dict[str, xr.Dataset]) -> None:
+        """Write each Dataset's values into the arrays of the Dataset
+        registered under its name, all or none of them."""
+        copies = []
+        for name, ds in writes.items():
+            registered = self._registered_datasets[name]
+            if set(ds.data_vars) != set(registered.data_vars):
+                raise ValueError(
+                    f"{name!r} has the data variables "
+                    f"{sorted(registered.data_vars)}, and the new Dataset "
+                    f"has {sorted(ds.data_vars)}"
+                )
+            for dim in registered.dims:
+                if dim not in ds.dims or not ds[dim].equals(registered[dim]):
+                    raise ValueError(
+                        f"the new Dataset for {name!r} has other coordinates "
+                        f"along {dim!r}; register it again, and prepare the "
+                        "query again"
+                    )
+            for var in registered.data_vars:
+                old, new = registered[var].variable, ds[var].variable
+                if new.dims != old.dims or new.dtype != old.dtype:
+                    raise ValueError(
+                        f"{name}.{var} is {old.dtype} over {old.dims}, and "
+                        f"the new values are {new.dtype} over {new.dims}"
+                    )
+                target = old.data
+                # A table shares a C-contiguous NumPy array's memory; any
+                # other array, it read or copied when it was registered.
+                if not (
+                    isinstance(target, np.ndarray)
+                    and target.flags.c_contiguous
+                    and target.flags.writeable
+                ):
+                    raise ValueError(
+                        f"{name}.{var} isn't a writable, C-contiguous NumPy "
+                        "array, which its table would share; register the "
+                        "new Dataset instead, and prepare the query again"
+                    )
+                values = np.asarray(new.values)
+                # A table works out which values are NULL (NaN) when it is
+                # registered, so a write must keep its NaNs in place.
+                if values.dtype.kind in "fc" and not np.array_equal(
+                    np.isnan(values), np.isnan(target)
+                ):
+                    raise ValueError(
+                        f"the new values of {name}.{var} have NaN in other "
+                        "places; xarray-sql reads a NaN as NULL, so register "
+                        "the new Dataset instead, and prepare the query again"
+                    )
+                copies.append((target, values))
+        for target, values in copies:
+            np.copyto(target, values)
+
+    def _wrap(self, inner) -> XarrayDataFrame:
+        """Wrap a DataFusion DataFrame so it can round-trip to xarray."""
         return XarrayDataFrame(inner, templates=self._registered_datasets)
 
 
+class PreparedQuery:
+    """A query planned once by [Context.prepare][xarray_sql.Context.prepare].
+
+    Call it to run the query, and get its rows as a ``pyarrow.Table``
+    ([to_dataset][xarray_sql.to_dataset] turns them back into a Dataset).
+    ``query()`` runs it on the values its tables hold now. ``query(w=ds)``
+    first writes ``ds``'s values into the table ``w``.
+
+    A write goes into the arrays of the Dataset registered as ``w``, as
+    ``w_ds["v"].values[...] = ds["v"].values`` would. That Dataset changes
+    too, and so does what any other query of the table returns. The new
+    Dataset must have the registered one's data variables, dtypes, dims and
+    coordinates. Its NaNs must be in the same places too, since xarray-sql
+    reads a NaN as NULL, and a table works out which values are NULL when it
+    is registered. To change anything else, register the new Dataset, and
+    prepare the query again: a prepared query reads the tables that were
+    registered when it was planned, and refuses to write into a name that
+    has been registered again since.
+
+    A plan keeps state from a run: a hash join, for one, keeps the table it
+    built. So each call runs a fresh copy of the plan, decoded from its
+    serialized form, which is much cheaper than planning. Some plans can't
+    be copied that way: a plan that reads a lazy table (a Dask array or a
+    lazily opened file), or one nested too deep to decode. These are
+    planned again on every call, so they return the right rows, without
+    the savings.
+    """
+
+    def __init__(self, ctx: Context, query: str):
+        self._ctx = ctx
+        self._query = query
+        # The Datasets a call may write into: those registered now.
+        self._registered = dict(ctx._registered_datasets)
+        self._plan: ExecutionPlan | None = None
+        df = SessionContext.sql(ctx, query)
+        self._schema = df.schema()
+        plan = df.execution_plan()
+        try:
+            _copy(ctx, plan)
+        except Exception:  # noqa: BLE001 - DataFusion raises a bare Exception
+            # It reads a table the plan's serialized form can't hold, or is
+            # nested deeper than protobuf decodes.
+            return
+        self._plan = plan
+
+    def __call__(self, **tables: xr.Dataset) -> pa.Table:
+        """Run the query, after writing each Dataset given into the table
+        of its name.
+
+        Args:
+            **tables: A Dataset for each table to write new values into.
+
+        Returns:
+            The query's rows.
+        """
+        for name in tables:
+            if name not in self._registered:
+                raise ValueError(
+                    f"no Dataset was registered as {name!r} when the query "
+                    "was prepared"
+                )
+            now = self._ctx._registered_datasets.get(name)
+            if now is not self._registered[name]:
+                raise ValueError(
+                    f"{name!r} has been registered again since the query was "
+                    "prepared, and the query reads the table it was planned "
+                    "with; prepare the query again"
+                )
+        self._ctx._write(tables)
+        if self._plan is None:
+            return self._ctx.sql(self._query).to_arrow_table()
+        plan = _copy(self._ctx, self._plan)
+        batches = [
+            batch.to_pyarrow()
+            for partition in range(plan.partition_count)
+            for batch in self._ctx.execute(plan, partition)
+        ]
+        return pa.Table.from_batches(batches, self._schema)
+
+
+def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan:
+    """A copy of ``plan`` with none of its state from a run, reading its
+    in-memory tables' values as they are now."""
+    # 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())
+
+
 XarrayContext = Context
 """The original name of [Context][xarray_sql.Context], kept for compatibility."""

From df52375c79ad46f9533ec34f53a15d7f9bedd100 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 14:29:44 -0700
Subject: [PATCH 2/7] Address review: write a prepared query's tables through
 xarray

PreparedQuery's writes go through xarray's own assignment,
registered[var][...] = new[var], rather than np.copyto, so they match
dims by name as xarray does: a Dataset with its dims in another order
is written in the right places. A write may now go into any writable
NumPy array, since a table reads one either way: shared, or on every
scan.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_prepare.py | 18 ++++++++++++++++++
 xarray_sql/sql.py     | 27 ++++++++++++++-------------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index f872a1c1..05c80b6d 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -67,6 +67,24 @@ def test_a_write_changes_the_registered_dataset():
     assert ctx.sql("SELECT SUM(v) AS s FROM a").to_pydict() == {"s": [12.0]}
 
 
+def test_a_write_matches_dims_by_name(ctx):
+    query = ctx.prepare(JOIN)
+    rng = np.random.default_rng(2)
+    a = _a(rng.normal(size=(4, 3)))
+    want = query(a=a)
+    assert query(a=a.transpose("j", "i")).equals(want)
+
+
+def test_a_write_reaches_a_table_read_lazily():
+    """An array that a table reads on every scan, rather than sharing (here,
+    one in Fortran order), takes writes too."""
+    v = np.asfortranarray(np.zeros((4, 3)))
+    ctx = Context()
+    ctx.from_dataset("a", _a(v))
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    assert query(a=_a(np.ones((4, 3)))).column("s").to_pylist() == [12.0]
+
+
 def test_a_plan_with_several_partitions_runs_again():
     ctx = Context(SessionConfig().with_target_partitions(4))
     n = 64
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index c48ba045..970f0687 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -293,25 +293,24 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
                     )
             for var in registered.data_vars:
                 old, new = registered[var].variable, ds[var].variable
-                if new.dims != old.dims or new.dtype != old.dtype:
+                # Dims are matched by name, as xarray matches them.
+                if set(new.dims) != set(old.dims) or new.dtype != old.dtype:
                     raise ValueError(
                         f"{name}.{var} is {old.dtype} over {old.dims}, and "
                         f"the new values are {new.dtype} over {new.dims}"
                     )
                 target = old.data
-                # A table shares a C-contiguous NumPy array's memory; any
-                # other array, it read or copied when it was registered.
+                # A table reads a NumPy array's memory, shared or on every
+                # scan; a Dask array or a lazily opened file has none to
+                # write into.
                 if not (
-                    isinstance(target, np.ndarray)
-                    and target.flags.c_contiguous
-                    and target.flags.writeable
+                    isinstance(target, np.ndarray) and target.flags.writeable
                 ):
                     raise ValueError(
-                        f"{name}.{var} isn't a writable, C-contiguous NumPy "
-                        "array, which its table would share; register the "
-                        "new Dataset instead, and prepare the query again"
+                        f"{name}.{var} isn't a writable NumPy array; register "
+                        "the new Dataset instead, and prepare the query again"
                     )
-                values = np.asarray(new.values)
+                values = new.transpose(*old.dims).values
                 # A table works out which values are NULL (NaN) when it is
                 # registered, so a write must keep its NaNs in place.
                 if values.dtype.kind in "fc" and not np.array_equal(
@@ -322,9 +321,11 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
                         "places; xarray-sql reads a NaN as NULL, so register "
                         "the new Dataset instead, and prepare the query again"
                     )
-                copies.append((target, values))
+                copies.append((registered[var], ds[var]))
+        # Assigning into a registered DataArray writes into its array, which
+        # the table shares.
         for target, values in copies:
-            np.copyto(target, values)
+            target[...] = values
 
     def _wrap(self, inner) -> XarrayDataFrame:
         """Wrap a DataFusion DataFrame so it can round-trip to xarray."""
@@ -340,7 +341,7 @@ class PreparedQuery:
     first writes ``ds``'s values into the table ``w``.
 
     A write goes into the arrays of the Dataset registered as ``w``, as
-    ``w_ds["v"].values[...] = ds["v"].values`` would. That Dataset changes
+    ``w_ds["v"][...] = ds["v"]`` would, matching dims by name. That Dataset changes
     too, and so does what any other query of the table returns. The new
     Dataset must have the registered one's data variables, dtypes, dims and
     coordinates. Its NaNs must be in the same places too, since xarray-sql

From 059af80f5ea59095b959012f39cdd4ebc401288e Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 19:25:08 -0700
Subject: [PATCH 3/7] Address review: NaT, large tables, tables registered
 again, now()

- A write must keep NaT in place, as it keeps NaN: a table reads both
  as NULL, from a validity bitmap made when it's registered.
- Keep the plan's copy only where copying is faster than planning.
  Copying a plan copies its in-memory tables, so over a 36 MB table a
  call took 45 ms against ctx.sql's 1 ms; it now plans again.
- Every call checks that the tables the query reads, through any views,
  are still the ones it was prepared with, on either path.
  Context.deregister_table forgets the Dataset, so a table replaced by
  any means counts as registered again.
- Say that now() is evaluated when the query is prepared.
- Shorten PreparedQuery's docstring to what a caller needs; how plans
  are copied moves to a comment.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 docs/performance.md   |  14 +++--
 tests/test_prepare.py |  55 +++++++++++++++++++
 xarray_sql/sql.py     | 125 ++++++++++++++++++++++++++----------------
 3 files changed, 144 insertions(+), 50 deletions(-)

diff --git a/docs/performance.md b/docs/performance.md
index dc135f58..60d7135a 100644
--- a/docs/performance.md
+++ b/docs/performance.md
@@ -193,10 +193,16 @@ for weights in trajectory:  # Datasets shaped like w's
     rows = step(w=weights)  # a pyarrow.Table
 ```
 
-On a join of two 1,024-row tables, a call takes 0.33 ms, and `ctx.sql`
-takes 0.91 ms (one target partition). A write must keep the table's
-shape, dtypes, coordinates and NaN positions. A plan over a lazy table
-(Dask or a lazily opened file) is planned again on each call for now.
+On a join of two 1,024-row tables, a call takes 0.41 ms, and `ctx.sql`
+takes 0.99 ms (one target partition). A write must keep the table's
+shape, dtypes, coordinates, and NaN and NaT positions. Functions such as
+`now()` are evaluated once, when the query is prepared.
+
+The savings are for small tables. Each call runs a copy of the plan,
+which copies the in-memory tables it reads. Where that's slower than
+planning, as it is over tables of more than about a megabyte, a call
+plans the query again and runs as fast as `ctx.sql`. So does a call of
+a plan over a lazy table (Dask or a lazily opened file).
 
 ## Round-trip faster with ORDER BY
 
diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index 05c80b6d..43a58c3b 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -1,5 +1,7 @@
 """Tests for queries planned once and run many times (`Context.prepare`)."""
 
+import time
+
 import numpy as np
 import pytest
 import xarray as xr
@@ -135,3 +137,56 @@ def test_a_write_needs_the_table_the_query_was_planned_with(ctx):
     ctx.from_dataset("a", _a(np.ones((4, 3))))
     with pytest.raises(ValueError, match="registered again"):
         query(a=_a(np.ones((4, 3))))
+
+
+def _times(*days):
+    t = np.array(
+        [np.datetime64("NaT") if d is None else np.datetime64(d) for d in days],
+        dtype="datetime64[ns]",
+    )
+    return xr.Dataset({"t": ("i", t)}, coords={"i": np.arange(len(t))})
+
+
+def test_a_write_must_keep_nat_in_place():
+    ctx = Context()
+    ctx.from_dataset("a", _times("2000-01-01", None, "2000-01-03"))
+    query = ctx.prepare("SELECT COUNT(t) AS n, MIN(t) AS m FROM a")
+    with pytest.raises(ValueError, match="NaT"):
+        query(a=_times(None, "2000-01-02", "2000-01-03"))
+    rows = query(a=_times("2000-01-04", None, "2000-01-02")).to_pydict()
+    assert rows["n"] == [2]
+    assert str(rows["m"][0]) == "2000-01-02 00:00:00"
+
+
+def test_a_query_reads_its_tables_as_they_were_prepared():
+    pytest.importorskip("dask")
+    ctx = Context()
+    ctx.from_dataset("a", _a(np.ones((4, 3))).chunk({"i": 2}))
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    ctx.deregister_table("a")
+    ctx.from_dataset("a", _a(np.full((4, 3), 2.0)))
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+def test_a_table_replaced_by_datafusion_counts_as_registered_again(ctx):
+    query = ctx.prepare(JOIN)
+    ctx.deregister_table("a")
+    ctx.register_record_batches("a", [ctx.sql("SELECT * FROM b").collect()])
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+def test_a_table_the_query_doesnt_read_can_be_registered_again(ctx):
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    want = query()
+    ctx.deregister_table("b")
+    ctx.from_dataset("b", _b(np.ones((4, 3))))
+    assert query().equals(want)
+
+
+def test_now_is_the_time_the_query_was_prepared(ctx):
+    query = ctx.prepare("SELECT now() AS t")
+    first = query()
+    time.sleep(0.01)
+    assert query().equals(first)
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 970f0687..f3d40377 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,3 +1,4 @@
+import time
 import uuid
 
 import numpy as np
@@ -5,7 +6,8 @@
 import xarray as xr
 from datafusion import SessionConfig, SessionContext
 from datafusion.catalog import Schema
-from datafusion.plan import ExecutionPlan
+from datafusion.expr import TableScan
+from datafusion.plan import ExecutionPlan, LogicalPlan
 from types import ModuleType
 
 from . import cftime as cft
@@ -242,6 +244,12 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         """
         return self._wrap(super().sql(query, *args, **kwargs))
 
+    def deregister_table(self, name: str) -> None:
+        """Remove a table from the session, and forget the Dataset
+        registered under its name."""
+        super().deregister_table(name)
+        self._registered_datasets.pop(name, None)
+
     def prepare(self, query: str) -> "PreparedQuery":
         """Plan a query once, to run it many times.
 
@@ -261,7 +269,7 @@ def prepare(self, query: str) -> "PreparedQuery":
                 rows = step(w=weights)
 
         See [PreparedQuery][xarray_sql.PreparedQuery] for what a call
-        can write, and which plans are run again rather than planned again.
+        can write.
 
         Args:
             query: A SQL query string.
@@ -311,15 +319,17 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
                         "the new Dataset instead, and prepare the query again"
                     )
                 values = new.transpose(*old.dims).values
-                # A table works out which values are NULL (NaN) when it is
-                # registered, so a write must keep its NaNs in place.
-                if values.dtype.kind in "fc" and not np.array_equal(
-                    np.isnan(values), np.isnan(target)
+                # A table works out which values are NULL (NaN or NaT) when
+                # it is registered, so a write must keep them in place.
+                missing = _missing(values)
+                if missing is not None and not np.array_equal(
+                    missing, _missing(target)
                 ):
                     raise ValueError(
-                        f"the new values of {name}.{var} have NaN in other "
-                        "places; xarray-sql reads a NaN as NULL, so register "
-                        "the new Dataset instead, and prepare the query again"
+                        f"the new values of {name}.{var} have NaN or NaT in "
+                        "other places; xarray-sql reads them as NULL, so "
+                        "register the new Dataset instead, and prepare the "
+                        "query again"
                     )
                 copies.append((registered[var], ds[var]))
         # Assigning into a registered DataArray writes into its array, which
@@ -335,47 +345,48 @@ def _wrap(self, inner) -> XarrayDataFrame:
 class PreparedQuery:
     """A query planned once by [Context.prepare][xarray_sql.Context.prepare].
 
-    Call it to run the query, and get its rows as a ``pyarrow.Table``
-    ([to_dataset][xarray_sql.to_dataset] turns them back into a Dataset).
-    ``query()`` runs it on the values its tables hold now. ``query(w=ds)``
-    first writes ``ds``'s values into the table ``w``.
-
-    A write goes into the arrays of the Dataset registered as ``w``, as
-    ``w_ds["v"][...] = ds["v"]`` would, matching dims by name. That Dataset changes
-    too, and so does what any other query of the table returns. The new
-    Dataset must have the registered one's data variables, dtypes, dims and
-    coordinates. Its NaNs must be in the same places too, since xarray-sql
-    reads a NaN as NULL, and a table works out which values are NULL when it
-    is registered. To change anything else, register the new Dataset, and
-    prepare the query again: a prepared query reads the tables that were
-    registered when it was planned, and refuses to write into a name that
-    has been registered again since.
-
-    A plan keeps state from a run: a hash join, for one, keeps the table it
-    built. So each call runs a fresh copy of the plan, decoded from its
-    serialized form, which is much cheaper than planning. Some plans can't
-    be copied that way: a plan that reads a lazy table (a Dask array or a
-    lazily opened file), or one nested too deep to decode. These are
-    planned again on every call, so they return the right rows, without
-    the savings.
+    ``query()`` runs it on the values its tables hold now, and returns its
+    rows as a ``pyarrow.Table``. ``query(w=ds)`` first writes ``ds``'s
+    values into the Dataset registered as ``w``, as
+    ``w_ds["v"][...] = ds["v"]`` would, so other queries of ``w`` see them
+    too. ``ds`` must match ``w_ds`` in all but its values: the same data
+    variables, dtypes, dims and coordinates, and NaN or NaT in the same
+    places, since xarray-sql reads them as NULL.
+
+    To change anything else, register the new Dataset and prepare the
+    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.
     """
 
     def __init__(self, ctx: Context, query: str):
         self._ctx = ctx
         self._query = query
-        # The Datasets a call may write into: those registered now.
-        self._registered = dict(ctx._registered_datasets)
-        self._plan: ExecutionPlan | None = None
+        start = time.perf_counter()
         df = SessionContext.sql(ctx, query)
-        self._schema = df.schema()
         plan = df.execution_plan()
+        planning = time.perf_counter() - start
+        self._schema = df.schema()
+        # The Datasets the query reads, as they are registered now.
+        self._registered = {
+            name: ctx._registered_datasets[name]
+            for name in _tables_read(df.logical_plan())
+            if name in ctx._registered_datasets
+        }
+        # 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.
+        self._plan: ExecutionPlan | None = None
+        start = time.perf_counter()
         try:
             _copy(ctx, plan)
         except Exception:  # noqa: BLE001 - DataFusion raises a bare Exception
-            # It reads a table the plan's serialized form can't hold, or is
-            # nested deeper than protobuf decodes.
             return
-        self._plan = plan
+        if time.perf_counter() - start < planning:
+            self._plan = plan
 
     def __call__(self, **tables: xr.Dataset) -> pa.Table:
         """Run the query, after writing each Dataset given into the table
@@ -390,19 +401,18 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table:
         for name in tables:
             if name not in self._registered:
                 raise ValueError(
-                    f"no Dataset was registered as {name!r} when the query "
-                    "was prepared"
+                    f"the query reads no Dataset registered as {name!r} "
+                    "when it was prepared"
                 )
-            now = self._ctx._registered_datasets.get(name)
-            if now is not self._registered[name]:
+        for name, registered in self._registered.items():
+            if self._ctx._registered_datasets.get(name) is not registered:
                 raise ValueError(
                     f"{name!r} has been registered again since the query was "
-                    "prepared, and the query reads the table it was planned "
-                    "with; prepare the query again"
+                    "prepared; prepare the query again"
                 )
         self._ctx._write(tables)
         if self._plan is None:
-            return self._ctx.sql(self._query).to_arrow_table()
+            return SessionContext.sql(self._ctx, self._query).to_arrow_table()
         plan = _copy(self._ctx, self._plan)
         batches = [
             batch.to_pyarrow()
@@ -421,5 +431,28 @@ def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan:
     return ExecutionPlan.from_bytes(ctx, plan.to_bytes())
 
 
+def _tables_read(plan: LogicalPlan) -> set[str]:
+    """The names of the tables ``plan`` scans, through any views."""
+    names = set()
+    stack = [plan]
+    while stack:
+        node = stack.pop()
+        variant = node.to_variant()
+        if isinstance(variant, TableScan):
+            names.add(variant.table_name())
+        stack.extend(node.inputs())
+    return names
+
+
+def _missing(values: np.ndarray) -> np.ndarray | None:
+    """Where ``values`` holds NaN or NaT, which xarray-sql reads as NULL, or
+    ``None`` for a dtype that holds neither."""
+    if values.dtype.kind in "fc":
+        return np.isnan(values)
+    if values.dtype.kind in "mM":
+        return np.isnat(values)
+    return None
+
+
 XarrayContext = Context
 """The original name of [Context][xarray_sql.Context], kept for compatibility."""

From 6719421937581c36c9c2ef4b05e372f910262e6e Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 19:37:34 -0700
Subject: [PATCH 4/7] Address review: every coordinate, every table read, now()
 when planning again

- A write must keep every coordinate, not only each dim's: a coordinate
  along two dims went unchecked.
- Track every table a prepared query reads, not only the Datasets, by
  how often its name has been removed or replaced: by deregister_table,
  or by SQL that creates or drops a table, which DataFusion runs without
  deregistering.
- A call that plans again starts from the optimized plan kept when the
  query was prepared, so it reads the tables it was prepared with, and
  now() keeps its value, as on the other path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_prepare.py | 40 ++++++++++++++++++++++++++++
 xarray_sql/sql.py     | 61 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 92 insertions(+), 9 deletions(-)

diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index 43a58c3b..b54f90e0 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -190,3 +190,43 @@ def test_now_is_the_time_the_query_was_prepared(ctx):
     first = query()
     time.sleep(0.01)
     assert query().equals(first)
+
+
+def test_a_write_must_keep_every_coordinate(ctx):
+    query = ctx.prepare(JOIN)
+    # A coordinate along both dims, which neither dim's coordinate holds.
+    labelled = _a(np.ones((4, 3))).assign_coords(
+        area=(("i", "j"), np.ones((4, 3)))
+    )
+    with pytest.raises(ValueError, match="coordinates"):
+        query(a=labelled)
+
+
+def test_a_table_not_from_a_dataset_counts_too(ctx):
+    ctx.register_record_batches("c", [ctx.sql("SELECT i, j FROM a").collect()])
+    query = ctx.prepare("SELECT COUNT(*) AS n FROM a JOIN c USING (i, j)")
+    assert query().column("n").to_pylist() == [12]
+    ctx.deregister_table("c")
+    ctx.register_record_batches("c", [ctx.sql("SELECT i, j FROM b").collect()])
+    with pytest.raises(ValueError, match="registered again"):
+        query()
+
+
+def test_now_is_the_time_a_query_over_a_lazy_table_was_prepared():
+    pytest.importorskip("dask")
+    ctx = Context()
+    ctx.from_dataset("a", _a(np.ones((4, 3))).chunk({"i": 2}))
+    query = ctx.prepare("SELECT now() AS t, SUM(v) AS s FROM a")
+    first = query()
+    time.sleep(0.01)
+    assert query().equals(first)
+
+
+def test_a_table_replaced_in_sql_counts_too():
+    ctx = Context()
+    ctx.sql("CREATE TABLE c AS VALUES (1), (2)").collect()
+    query = ctx.prepare("SELECT COUNT(*) AS n FROM c")
+    assert query().column("n").to_pylist() == [2]
+    ctx.sql("CREATE OR REPLACE TABLE c AS VALUES (1), (2), (3)").collect()
+    with pytest.raises(ValueError, match="registered again"):
+        query()
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index f3d40377..1f2cb971 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,3 +1,4 @@
+import re
 import time
 import uuid
 
@@ -64,6 +65,11 @@ def __init__(self, config: SessionConfig | None = None, *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] = {}
+        # How many times each name's table has been removed or replaced.
+        # DataFusion refuses to register a name that's taken, so a table
+        # registered again from Python was deregistered first; SQL that
+        # creates or drops a table counts too (see sql()).
+        self._deregistrations: dict[str, int] = {}
         # 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:
@@ -242,6 +248,13 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         Returns:
             An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame.
         """
+        # DataFusion runs a statement that creates or drops a table as it
+        # plans it, replacing a table without deregistering it.
+        replaced = _table_replaced(query)
+        if replaced is not None:
+            self._deregistrations[replaced] = (
+                self._deregistrations.get(replaced, 0) + 1
+            )
         return self._wrap(super().sql(query, *args, **kwargs))
 
     def deregister_table(self, name: str) -> None:
@@ -249,6 +262,7 @@ def deregister_table(self, name: str) -> None:
         registered under its name."""
         super().deregister_table(name)
         self._registered_datasets.pop(name, None)
+        self._deregistrations[name] = self._deregistrations.get(name, 0) + 1
 
     def prepare(self, query: str) -> "PreparedQuery":
         """Plan a query once, to run it many times.
@@ -292,11 +306,13 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
                     f"{sorted(registered.data_vars)}, and the new Dataset "
                     f"has {sorted(ds.data_vars)}"
                 )
-            for dim in registered.dims:
-                if dim not in ds.dims or not ds[dim].equals(registered[dim]):
+            for coord in set(registered.coords) | set(ds.coords):
+                if coord not in ds.coords or not ds[coord].equals(
+                    registered.coords.get(coord)
+                ):
                     raise ValueError(
                         f"the new Dataset for {name!r} has other coordinates "
-                        f"along {dim!r}; register it again, and prepare the "
+                        f"({coord!r}); register it again, and prepare the "
                         "query again"
                     )
             for var in registered.data_vars:
@@ -361,16 +377,24 @@ class PreparedQuery:
 
     def __init__(self, ctx: Context, query: str):
         self._ctx = ctx
-        self._query = query
         start = time.perf_counter()
         df = SessionContext.sql(ctx, query)
         plan = df.execution_plan()
         planning = time.perf_counter() - start
         self._schema = df.schema()
-        # The Datasets the query reads, as they are registered now.
+        # Planned again from its optimized form, the query reads the tables
+        # it was prepared with, and functions such as now() keep the values
+        # they were folded into.
+        self._optimized = df.optimized_logical_plan()
+        # Each table the query reads, by how often its name had been
+        # deregistered; and of those, the Datasets a call may write into.
+        self._read = {
+            name: ctx._deregistrations.get(name, 0)
+            for name in _tables_read(df.logical_plan())
+        }
         self._registered = {
             name: ctx._registered_datasets[name]
-            for name in _tables_read(df.logical_plan())
+            for name in self._read
             if name in ctx._registered_datasets
         }
         # A plan keeps state from a run (a hash join keeps the table it
@@ -404,15 +428,16 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table:
                     f"the query reads no Dataset registered as {name!r} "
                     "when it was prepared"
                 )
-        for name, registered in self._registered.items():
-            if self._ctx._registered_datasets.get(name) is not registered:
+        for name, deregistrations in self._read.items():
+            if self._ctx._deregistrations.get(name, 0) != deregistrations:
                 raise ValueError(
                     f"{name!r} has been registered again since the query was "
                     "prepared; prepare the query again"
                 )
         self._ctx._write(tables)
         if self._plan is None:
-            return SessionContext.sql(self._ctx, self._query).to_arrow_table()
+            df = self._ctx.create_dataframe_from_logical_plan(self._optimized)
+            return df.to_arrow_table()
         plan = _copy(self._ctx, self._plan)
         batches = [
             batch.to_pyarrow()
@@ -444,6 +469,24 @@ def _tables_read(plan: LogicalPlan) -> set[str]:
     return names
 
 
+_TABLE_DDL = re.compile(
+    r"\s*(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
+    r"\s+(?:IF\s+(?:NOT\s+)?EXISTS\s+)?(\"[^\"]+\"|[\w.]+)",
+    re.IGNORECASE,
+)
+
+
+def _table_replaced(query: str) -> str | None:
+    """The name of the table ``query`` creates or drops, as DataFusion
+    resolves it, or ``None`` for any other statement."""
+    match = _TABLE_DDL.match(query)
+    if match is None:
+        return None
+    name = match.group(1)
+    # DataFusion keeps a quoted name's case, and lowers an unquoted one.
+    return name[1:-1] if name.startswith('"') else name.lower()
+
+
 def _missing(values: np.ndarray) -> np.ndarray | None:
     """Where ``values`` holds NaN or NaT, which xarray-sql reads as NULL, or
     ``None`` for a dtype that holds neither."""

From 6940527e8feb91e8ec1591700c00cbf8a432be51 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Mon, 26 Oct 2026 03:05:06 -0700
Subject: [PATCH 5/7] Address review: SQL that replaces a table forgets its
 Dataset; skip comments

- CREATE OR REPLACE TABLE over a name registered from a Dataset forgets
  the Dataset, as deregister_table does, so a query prepared after it
  can't write into the old arrays.
- Comments before a statement that creates or drops a table no longer
  hide it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_prepare.py | 14 ++++++++++++--
 xarray_sql/sql.py     | 15 ++++++++++-----
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index b54f90e0..6c313758 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -222,11 +222,21 @@ def test_now_is_the_time_a_query_over_a_lazy_table_was_prepared():
     assert query().equals(first)
 
 
-def test_a_table_replaced_in_sql_counts_too():
+@pytest.mark.parametrize("comment", ["", "-- refresh\n", "/* refresh */ "])
+def test_a_table_replaced_in_sql_counts_too(comment):
     ctx = Context()
     ctx.sql("CREATE TABLE c AS VALUES (1), (2)").collect()
     query = ctx.prepare("SELECT COUNT(*) AS n FROM c")
     assert query().column("n").to_pylist() == [2]
-    ctx.sql("CREATE OR REPLACE TABLE c AS VALUES (1), (2), (3)").collect()
+    ctx.sql(
+        comment + "CREATE OR REPLACE TABLE c AS VALUES (1), (2), (3)"
+    ).collect()
     with pytest.raises(ValueError, match="registered again"):
         query()
+
+
+def test_a_dataset_replaced_in_sql_takes_no_writes(ctx):
+    ctx.sql("CREATE OR REPLACE TABLE a AS SELECT * FROM b").collect()
+    query = ctx.prepare("SELECT SUM(w) AS s FROM a")
+    with pytest.raises(ValueError, match="no Dataset"):
+        query(a=_a(np.ones((4, 3))))
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 1f2cb971..583fd0bb 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -252,15 +252,18 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         # plans it, replacing a table without deregistering it.
         replaced = _table_replaced(query)
         if replaced is not None:
-            self._deregistrations[replaced] = (
-                self._deregistrations.get(replaced, 0) + 1
-            )
+            self._forget(replaced)
         return self._wrap(super().sql(query, *args, **kwargs))
 
     def deregister_table(self, name: str) -> None:
         """Remove a table from the session, and forget the Dataset
         registered under its name."""
         super().deregister_table(name)
+        self._forget(name)
+
+    def _forget(self, name: str) -> None:
+        """Forget the Dataset registered as ``name``, whose table has been
+        removed or replaced."""
         self._registered_datasets.pop(name, None)
         self._deregistrations[name] = self._deregistrations.get(name, 0) + 1
 
@@ -470,9 +473,11 @@ def _tables_read(plan: LogicalPlan) -> set[str]:
 
 
 _TABLE_DDL = re.compile(
-    r"\s*(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
+    # Any comments before the statement, then the statement.
+    r"(?:\s+|--[^\n]*|/\*.*?\*/)*"
+    r"(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
     r"\s+(?:IF\s+(?:NOT\s+)?EXISTS\s+)?(\"[^\"]+\"|[\w.]+)",
-    re.IGNORECASE,
+    re.IGNORECASE | re.DOTALL,
 )
 
 

From 2102e038bb903bbf958533a086645e293aa9135f Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Mon, 26 Oct 2026 03:19:04 -0700
Subject: [PATCH 6/7] Address review: forget a table only when SQL replaces or
 drops it

CREATE TABLE IF NOT EXISTS over an existing table changes nothing, and
a plain CREATE TABLE over one fails, yet both forgot the table's Dataset
and failed the prepared queries that read it. Only CREATE OR REPLACE
TABLE and DROP TABLE forget a table now, and only one that existed
before the statement ran, once it has run.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_prepare.py | 22 ++++++++++++++++++++++
 xarray_sql/sql.py     | 31 ++++++++++++++++---------------
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/tests/test_prepare.py b/tests/test_prepare.py
index 6c313758..8cc356bc 100644
--- a/tests/test_prepare.py
+++ b/tests/test_prepare.py
@@ -240,3 +240,25 @@ def test_a_dataset_replaced_in_sql_takes_no_writes(ctx):
     query = ctx.prepare("SELECT SUM(w) AS s FROM a")
     with pytest.raises(ValueError, match="no Dataset"):
         query(a=_a(np.ones((4, 3))))
+
+
+@pytest.mark.parametrize(
+    "statement",
+    [
+        "CREATE TABLE IF NOT EXISTS a AS SELECT * FROM b",
+        "DROP TABLE IF EXISTS c",
+    ],
+)
+def test_sql_that_changes_no_table_keeps_a_prepared_query(ctx, statement):
+    query = ctx.prepare("SELECT SUM(v) AS s FROM a")
+    ctx.sql(statement).collect()
+    assert query(a=_a(np.ones((4, 3)))).column("s").to_pylist() == [12.0]
+
+
+def test_a_quoted_table_replaced_in_sql_counts_too():
+    ctx = Context()
+    ctx.sql('CREATE TABLE "Mixed" AS VALUES (1)').collect()
+    query = ctx.prepare('SELECT COUNT(*) AS n FROM "Mixed"')
+    ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect()
+    with pytest.raises(ValueError, match="registered again"):
+        query()
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 583fd0bb..e9e77bc9 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -68,7 +68,7 @@ def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
         # How many times each name's table has been removed or replaced.
         # DataFusion refuses to register a name that's taken, so a table
         # registered again from Python was deregistered first; SQL that
-        # creates or drops a table counts too (see sql()).
+        # replaces or drops a table counts too (see sql()).
         self._deregistrations: dict[str, int] = {}
         # With pyproj installed, every context speaks CRS out of the box:
         # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform.
@@ -248,12 +248,17 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         Returns:
             An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame.
         """
-        # DataFusion runs a statement that creates or drops a table as it
-        # plans it, replacing a table without deregistering it.
+        # DataFusion runs a statement that replaces or drops a table as it
+        # plans it, without deregistering the table.
         replaced = _table_replaced(query)
-        if replaced is not None:
-            self._forget(replaced)
-        return self._wrap(super().sql(query, *args, **kwargs))
+        existed = replaced is not None and self.table_exist(replaced)
+        df = super().sql(query, *args, **kwargs)
+        if existed:
+            # DataFusion keeps a quoted name's case, and lowers an unquoted
+            # one.
+            quoted = replaced.startswith('"')
+            self._forget(replaced[1:-1] if quoted else replaced.lower())
+        return self._wrap(df)
 
     def deregister_table(self, name: str) -> None:
         """Remove a table from the session, and forget the Dataset
@@ -475,21 +480,17 @@ def _tables_read(plan: LogicalPlan) -> set[str]:
 _TABLE_DDL = re.compile(
     # Any comments before the statement, then the statement.
     r"(?:\s+|--[^\n]*|/\*.*?\*/)*"
-    r"(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
-    r"\s+(?:IF\s+(?:NOT\s+)?EXISTS\s+)?(\"[^\"]+\"|[\w.]+)",
+    r"(?:CREATE\s+OR\s+REPLACE\s+(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
+    r"\s+(?:IF\s+EXISTS\s+)?(\"[^\"]+\"|[\w.]+)",
     re.IGNORECASE | re.DOTALL,
 )
 
 
 def _table_replaced(query: str) -> str | None:
-    """The name of the table ``query`` creates or drops, as DataFusion
-    resolves it, or ``None`` for any other statement."""
+    """The name of the table ``query`` replaces or drops, as it's written,
+    or ``None`` for any other statement."""
     match = _TABLE_DDL.match(query)
-    if match is None:
-        return None
-    name = match.group(1)
-    # DataFusion keeps a quoted name's case, and lowers an unquoted one.
-    return name[1:-1] if name.startswith('"') else name.lower()
+    return None if match is None else match.group(1)
 
 
 def _missing(values: np.ndarray) -> np.ndarray | None:

From 560fa8e6d8c197d423a47bf240f845150c9b6553 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Mon, 26 Oct 2026 03:19:20 -0700
Subject: [PATCH 7/7] Narrow the replaced name's type for mypy

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 xarray_sql/sql.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index e9e77bc9..0c56d039 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -253,7 +253,7 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
         replaced = _table_replaced(query)
         existed = replaced is not None and self.table_exist(replaced)
         df = super().sql(query, *args, **kwargs)
-        if existed:
+        if replaced is not None and existed:
             # DataFusion keeps a quoted name's case, and lowers an unquoted
             # one.
             quoted = replaced.startswith('"')