diff --git a/docs/performance.md b/docs/performance.md index c461de93..60d7135a 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -177,6 +177,33 @@ 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.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 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..8cc356bc --- /dev/null +++ b/tests/test_prepare.py @@ -0,0 +1,264 @@ +"""Tests for queries planned once and run many times (`Context.prepare`).""" + +import time + +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_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 + 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)))) + + +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) + + +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) + + +@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( + 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)))) + + +@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/__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..0c56d039 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,9 +1,14 @@ +import re +import time 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.expr import TableScan +from datafusion.plan import ExecutionPlan, LogicalPlan from types import ModuleType from . import cftime as cft @@ -60,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 + # 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. if _proj is not None: @@ -238,9 +248,260 @@ 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) + # DataFusion runs a statement that replaces or drops a table as it + # plans it, without deregistering the table. + replaced = _table_replaced(query) + existed = replaced is not None and self.table_exist(replaced) + df = super().sql(query, *args, **kwargs) + if replaced is not None and 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 + 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 + + 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. + + 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 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"({coord!r}); register it again, and prepare the " + "query again" + ) + for var in registered.data_vars: + old, new = registered[var].variable, ds[var].variable + # 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 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.writeable + ): + raise ValueError( + f"{name}.{var} isn't a writable NumPy array; register " + "the new Dataset instead, and prepare the query again" + ) + values = new.transpose(*old.dims).values + # 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 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 + # the table shares. + for target, values in copies: + 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]. + + ``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 + start = time.perf_counter() + df = SessionContext.sql(ctx, query) + plan = df.execution_plan() + planning = time.perf_counter() - start + self._schema = df.schema() + # 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 self._read + 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 + return + 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 + 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"the query reads no Dataset registered as {name!r} " + "when it was prepared" + ) + 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: + 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() + 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()) + + +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 + + +_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+EXISTS\s+)?(\"[^\"]+\"|[\w.]+)", + re.IGNORECASE | re.DOTALL, +) + + +def _table_replaced(query: str) -> str | None: + """The name of the table ``query`` replaces or drops, as it's written, + or ``None`` for any other statement.""" + match = _TABLE_DDL.match(query) + return None if match is None else match.group(1) + + +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."""