Skip to content

Commit df52375

Browse files
alxmrsclaude
andcommitted
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 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent c0c4f7c commit df52375

2 files changed

Lines changed: 32 additions & 13 deletions

File tree

‎tests/test_prepare.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,24 @@ def test_a_write_changes_the_registered_dataset():
6767
assert ctx.sql("SELECT SUM(v) AS s FROM a").to_pydict() == {"s": [12.0]}
6868

6969

70+
def test_a_write_matches_dims_by_name(ctx):
71+
query = ctx.prepare(JOIN)
72+
rng = np.random.default_rng(2)
73+
a = _a(rng.normal(size=(4, 3)))
74+
want = query(a=a)
75+
assert query(a=a.transpose("j", "i")).equals(want)
76+
77+
78+
def test_a_write_reaches_a_table_read_lazily():
79+
"""An array that a table reads on every scan, rather than sharing (here,
80+
one in Fortran order), takes writes too."""
81+
v = np.asfortranarray(np.zeros((4, 3)))
82+
ctx = Context()
83+
ctx.from_dataset("a", _a(v))
84+
query = ctx.prepare("SELECT SUM(v) AS s FROM a")
85+
assert query(a=_a(np.ones((4, 3)))).column("s").to_pylist() == [12.0]
86+
87+
7088
def test_a_plan_with_several_partitions_runs_again():
7189
ctx = Context(SessionConfig().with_target_partitions(4))
7290
n = 64

‎xarray_sql/sql.py‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -293,25 +293,24 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
293293
)
294294
for var in registered.data_vars:
295295
old, new = registered[var].variable, ds[var].variable
296-
if new.dims != old.dims or new.dtype != old.dtype:
296+
# Dims are matched by name, as xarray matches them.
297+
if set(new.dims) != set(old.dims) or new.dtype != old.dtype:
297298
raise ValueError(
298299
f"{name}.{var} is {old.dtype} over {old.dims}, and "
299300
f"the new values are {new.dtype} over {new.dims}"
300301
)
301302
target = old.data
302-
# A table shares a C-contiguous NumPy array's memory; any
303-
# other array, it read or copied when it was registered.
303+
# A table reads a NumPy array's memory, shared or on every
304+
# scan; a Dask array or a lazily opened file has none to
305+
# write into.
304306
if not (
305-
isinstance(target, np.ndarray)
306-
and target.flags.c_contiguous
307-
and target.flags.writeable
307+
isinstance(target, np.ndarray) and target.flags.writeable
308308
):
309309
raise ValueError(
310-
f"{name}.{var} isn't a writable, C-contiguous NumPy "
311-
"array, which its table would share; register the "
312-
"new Dataset instead, and prepare the query again"
310+
f"{name}.{var} isn't a writable NumPy array; register "
311+
"the new Dataset instead, and prepare the query again"
313312
)
314-
values = np.asarray(new.values)
313+
values = new.transpose(*old.dims).values
315314
# A table works out which values are NULL (NaN) when it is
316315
# registered, so a write must keep its NaNs in place.
317316
if values.dtype.kind in "fc" and not np.array_equal(
@@ -322,9 +321,11 @@ def _write(self, writes: dict[str, xr.Dataset]) -> None:
322321
"places; xarray-sql reads a NaN as NULL, so register "
323322
"the new Dataset instead, and prepare the query again"
324323
)
325-
copies.append((target, values))
324+
copies.append((registered[var], ds[var]))
325+
# Assigning into a registered DataArray writes into its array, which
326+
# the table shares.
326327
for target, values in copies:
327-
np.copyto(target, values)
328+
target[...] = values
328329

329330
def _wrap(self, inner) -> XarrayDataFrame:
330331
"""Wrap a DataFusion DataFrame so it can round-trip to xarray."""
@@ -340,7 +341,7 @@ class PreparedQuery:
340341
first writes ``ds``'s values into the table ``w``.
341342
342343
A write goes into the arrays of the Dataset registered as ``w``, as
343-
``w_ds["v"].values[...] = ds["v"].values`` would. That Dataset changes
344+
``w_ds["v"][...] = ds["v"]`` would, matching dims by name. That Dataset changes
344345
too, and so does what any other query of the table returns. The new
345346
Dataset must have the registered one's data variables, dtypes, dims and
346347
coordinates. Its NaNs must be in the same places too, since xarray-sql

0 commit comments

Comments
 (0)