|
1 | 1 | import uuid |
2 | 2 |
|
| 3 | +import numpy as np |
3 | 4 | import pyarrow as pa |
4 | 5 | import xarray as xr |
5 | 6 | from datafusion import SessionContext |
6 | 7 | from datafusion.catalog import Schema |
| 8 | +from datafusion.plan import ExecutionPlan |
7 | 9 | from types import ModuleType |
8 | 10 |
|
9 | 11 | from . import cftime as cft |
@@ -244,10 +246,188 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: |
244 | 246 | return self._wrap(_ad.sql(self, query)) |
245 | 247 | return self._wrap(super().sql(query, *args, **kwargs)) |
246 | 248 |
|
| 249 | + def prepare(self, query: str) -> "PreparedQuery": |
| 250 | + """Plan a query once, to run it many times. |
| 251 | +
|
| 252 | + A training loop runs the same query on every step, and only the |
| 253 | + values in its tables change. DataFusion plans a query each time it |
| 254 | + runs, and for a query over small tables, planning can take far |
| 255 | + longer than running. ``prepare`` plans the query once, and calling |
| 256 | + the result runs that plan, as ``jax.jit`` traces a function once |
| 257 | + and runs the compiled function after that. A call can first write |
| 258 | + new values into the query's tables:: |
| 259 | +
|
| 260 | + step = ctx.prepare( |
| 261 | + "SELECT x.i, SUM(x.v * w.v) AS y " |
| 262 | + "FROM x JOIN w ON x.j = w.j GROUP BY x.i" |
| 263 | + ) |
| 264 | + for weights in trajectory: # Datasets shaped like w's |
| 265 | + rows = step(w=weights) |
| 266 | +
|
| 267 | + See [PreparedQuery][xarray_sql.PreparedQuery] for what a call |
| 268 | + can write, and which plans are run again rather than planned again. |
| 269 | +
|
| 270 | + Args: |
| 271 | + query: A SQL query string. |
| 272 | +
|
| 273 | + Returns: |
| 274 | + A [PreparedQuery][xarray_sql.PreparedQuery]; call it to run |
| 275 | + the query. |
| 276 | + """ |
| 277 | + return PreparedQuery(self, query) |
| 278 | + |
| 279 | + def _write(self, writes: dict[str, xr.Dataset]) -> None: |
| 280 | + """Write each Dataset's values into the arrays of the Dataset |
| 281 | + registered under its name, all or none of them.""" |
| 282 | + copies = [] |
| 283 | + for name, ds in writes.items(): |
| 284 | + registered = self._registered_datasets[name] |
| 285 | + if set(ds.data_vars) != set(registered.data_vars): |
| 286 | + raise ValueError( |
| 287 | + f"{name!r} has the data variables " |
| 288 | + f"{sorted(registered.data_vars)}, and the new Dataset " |
| 289 | + f"has {sorted(ds.data_vars)}" |
| 290 | + ) |
| 291 | + for dim in registered.dims: |
| 292 | + if dim not in ds.dims or not ds[dim].equals(registered[dim]): |
| 293 | + raise ValueError( |
| 294 | + f"the new Dataset for {name!r} has other coordinates " |
| 295 | + f"along {dim!r}; register it again, and prepare the " |
| 296 | + "query again" |
| 297 | + ) |
| 298 | + for var in registered.data_vars: |
| 299 | + old, new = registered[var].variable, ds[var].variable |
| 300 | + if new.dims != old.dims or new.dtype != old.dtype: |
| 301 | + raise ValueError( |
| 302 | + f"{name}.{var} is {old.dtype} over {old.dims}, and " |
| 303 | + f"the new values are {new.dtype} over {new.dims}" |
| 304 | + ) |
| 305 | + target = old.data |
| 306 | + # A table shares a C-contiguous NumPy array's memory; any |
| 307 | + # other array, it read or copied when it was registered. |
| 308 | + if not ( |
| 309 | + isinstance(target, np.ndarray) |
| 310 | + and target.flags.c_contiguous |
| 311 | + and target.flags.writeable |
| 312 | + ): |
| 313 | + raise ValueError( |
| 314 | + f"{name}.{var} isn't a writable, C-contiguous NumPy " |
| 315 | + "array, which its table would share; register the " |
| 316 | + "new Dataset instead, and prepare the query again" |
| 317 | + ) |
| 318 | + values = np.asarray(new.values) |
| 319 | + # A table works out which values are NULL (NaN) when it is |
| 320 | + # registered, so a write must keep its NaNs in place. |
| 321 | + if values.dtype.kind in "fc" and not np.array_equal( |
| 322 | + np.isnan(values), np.isnan(target) |
| 323 | + ): |
| 324 | + raise ValueError( |
| 325 | + f"the new values of {name}.{var} have NaN in other " |
| 326 | + "places; xarray-sql reads a NaN as NULL, so register " |
| 327 | + "the new Dataset instead, and prepare the query again" |
| 328 | + ) |
| 329 | + copies.append((target, values)) |
| 330 | + for target, values in copies: |
| 331 | + np.copyto(target, values) |
| 332 | + |
247 | 333 | def _wrap(self, inner) -> XarrayDataFrame: |
248 | 334 | """Wrap a DataFusion DataFrame so it can round-trip to xarray.""" |
249 | 335 | return XarrayDataFrame(inner, templates=self._registered_datasets) |
250 | 336 |
|
251 | 337 |
|
| 338 | +class PreparedQuery: |
| 339 | + """A query planned once by [Context.prepare][xarray_sql.Context.prepare]. |
| 340 | +
|
| 341 | + Call it to run the query, and get its rows as a ``pyarrow.Table`` |
| 342 | + ([to_dataset][xarray_sql.to_dataset] turns them back into a Dataset). |
| 343 | + ``query()`` runs it on the values its tables hold now. ``query(w=ds)`` |
| 344 | + first writes ``ds``'s values into the table ``w``. |
| 345 | +
|
| 346 | + A write goes into the arrays of the Dataset registered as ``w``, as |
| 347 | + ``w_ds["v"].values[...] = ds["v"].values`` would. That Dataset changes |
| 348 | + too, and so does what any other query of the table returns. The new |
| 349 | + Dataset must have the registered one's data variables, dtypes, dims and |
| 350 | + coordinates. Its NaNs must be in the same places too, since xarray-sql |
| 351 | + reads a NaN as NULL, and a table works out which values are NULL when it |
| 352 | + is registered. To change anything else, register the new Dataset, and |
| 353 | + prepare the query again: a prepared query reads the tables that were |
| 354 | + registered when it was planned, and refuses to write into a name that |
| 355 | + has been registered again since. |
| 356 | +
|
| 357 | + A plan keeps state from a run: a hash join, for one, keeps the table it |
| 358 | + built. So each call runs a fresh copy of the plan, decoded from its |
| 359 | + serialized form, which is much cheaper than planning. Some plans can't |
| 360 | + be copied that way: a plan that reads a lazy table (a Dask array or a |
| 361 | + lazily opened file), or one nested too deep to decode. A query that |
| 362 | + calls ``grad``, ``vjp`` or ``jvp`` isn't planned ahead either. These are |
| 363 | + planned again on every call, so they return the right rows, without |
| 364 | + the savings. |
| 365 | + """ |
| 366 | + |
| 367 | + def __init__(self, ctx: Context, query: str): |
| 368 | + self._ctx = ctx |
| 369 | + self._query = query |
| 370 | + # The Datasets a call may write into: those registered now. |
| 371 | + self._registered = dict(ctx._registered_datasets) |
| 372 | + self._plan: ExecutionPlan | None = None |
| 373 | + if _ad is not None and _ad.calls_ad(query): |
| 374 | + # TODO(xqlsystems/ddx#202): run ddx's prepared program, whose |
| 375 | + # steps are planned once, rather than rewriting on each call. |
| 376 | + return |
| 377 | + df = SessionContext.sql(ctx, query) |
| 378 | + self._schema = df.schema() |
| 379 | + plan = df.execution_plan() |
| 380 | + try: |
| 381 | + _copy(ctx, plan) |
| 382 | + except Exception: # noqa: BLE001 - DataFusion raises a bare Exception |
| 383 | + # It reads a table the plan's serialized form can't hold, or is |
| 384 | + # nested deeper than protobuf decodes. |
| 385 | + return |
| 386 | + self._plan = plan |
| 387 | + |
| 388 | + def __call__(self, **tables: xr.Dataset) -> pa.Table: |
| 389 | + """Run the query, after writing each Dataset given into the table |
| 390 | + of its name. |
| 391 | +
|
| 392 | + Args: |
| 393 | + **tables: A Dataset for each table to write new values into. |
| 394 | +
|
| 395 | + Returns: |
| 396 | + The query's rows. |
| 397 | + """ |
| 398 | + for name in tables: |
| 399 | + if name not in self._registered: |
| 400 | + raise ValueError( |
| 401 | + f"no Dataset was registered as {name!r} when the query " |
| 402 | + "was prepared" |
| 403 | + ) |
| 404 | + now = self._ctx._registered_datasets.get(name) |
| 405 | + if now is not self._registered[name]: |
| 406 | + raise ValueError( |
| 407 | + f"{name!r} has been registered again since the query was " |
| 408 | + "prepared, and the query reads the table it was planned " |
| 409 | + "with; prepare the query again" |
| 410 | + ) |
| 411 | + self._ctx._write(tables) |
| 412 | + if self._plan is None: |
| 413 | + return self._ctx.sql(self._query).to_arrow_table() |
| 414 | + plan = _copy(self._ctx, self._plan) |
| 415 | + batches = [ |
| 416 | + batch.to_pyarrow() |
| 417 | + for partition in range(plan.partition_count) |
| 418 | + for batch in self._ctx.execute(plan, partition) |
| 419 | + ] |
| 420 | + return pa.Table.from_batches(batches, self._schema) |
| 421 | + |
| 422 | + |
| 423 | +def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan: |
| 424 | + """A copy of ``plan`` with none of its state from a run, reading its |
| 425 | + in-memory tables' values as they are now.""" |
| 426 | + # TODO(#283): reset the plan's state with datafusion-python's binding of |
| 427 | + # reset_plan_states, once it has one, rather than decode it again. |
| 428 | + # `plan` itself never runs, so it has no state to leave behind. |
| 429 | + return ExecutionPlan.from_bytes(ctx, plan.to_bytes()) |
| 430 | + |
| 431 | + |
252 | 432 | XarrayContext = Context |
253 | 433 | """The original name of [Context][xarray_sql.Context], kept for compatibility.""" |
0 commit comments