Skip to content

Commit bd8650d

Browse files
alxmrsclaude
andcommitted
Read arrays that Arrow would copy lazily, so writes in place show
A data variable whose Arrow column can't be its own memory (Fortran order, a transposed or strided view, booleans, strings) was copied into the in-memory table, so a write into the array in place no longer showed in queries, as from_dataset says it does. Such a Dataset is now read lazily, as on main, which reads the array on every scan. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 87752c5 commit bd8650d

3 files changed

Lines changed: 63 additions & 15 deletions

File tree

‎tests/test_sql.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -666,6 +666,31 @@ def test_in_memory_table_shares_the_dataset_s_memory(chunks):
666666
assert ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [7.0] * 24
667667

668668

669+
@pytest.mark.parametrize(
670+
"values, written",
671+
[
672+
(np.asfortranarray(np.zeros((4, 6))), 7.0),
673+
(np.zeros((6, 4)).T, 7.0),
674+
(np.zeros((4, 12))[:, ::2], 7.0),
675+
(np.zeros((4, 6), dtype=bool), True),
676+
(np.full((4, 6), "a"), "b"),
677+
],
678+
ids=["fortran-order", "transposed", "strided", "bool", "str"],
679+
)
680+
def test_in_place_write_shows_where_arrow_would_copy(values, written):
681+
"""A write in place shows in the next query even for an array whose
682+
Arrow column would be a copy, which isn't an in-memory table."""
683+
ds = xr.Dataset({"w": (("a", "b"), values)})
684+
ctx = XarrayContext()
685+
ctx.from_dataset("t", ds)
686+
687+
ds["w"].values[...] = written
688+
689+
assert (
690+
ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [written] * 24
691+
)
692+
693+
669694
def test_in_memory_dataset_over_the_limit_in_arrow_is_read_lazily():
670695
"""The 64 MiB limit is on the table as Arrow, which a long string can
671696
exceed when an estimate of its size doesn't."""

‎xarray_sql/reader.py‎

Lines changed: 35 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,23 @@ def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
222222
return num_rows * sum(widths)
223223

224224

225+
def _shares_memory_with_arrow(values: np.ndarray) -> bool:
226+
"""Whether ``values``'s Arrow column, its rows in C order, can be its
227+
own memory rather than a copy.
228+
229+
Arrow lays out a number or a datetime as NumPy does, in the machine's
230+
byte order. It packs booleans into bits, and keeps strings in buffers
231+
of its own, so those are copies; so is an array that isn't
232+
C-contiguous, such as one in Fortran order or a transposed or strided
233+
view, whose rows in C order aren't one run of memory.
234+
"""
235+
return (
236+
values.dtype.kind in "iufmM"
237+
and values.dtype.isnative
238+
and values.flags.c_contiguous
239+
)
240+
241+
225242
def _in_memory_partitions(
226243
ds: xr.Dataset,
227244
chunks: Chunks = None,
@@ -241,23 +258,29 @@ def _in_memory_partitions(
241258
242259
The rows are in C order, and each partition is a contiguous range of
243260
them, as many partitions as ``chunks`` makes chunks, up to the number
244-
of CPUs. So a data variable whose array is C-contiguous isn't copied,
245-
whatever the chunks: its Arrow column shares the array's memory, as
246-
does DataFusion's table, through Arrow's C data interface. A NaN adds a
247-
validity bitmap beside it, one bit per value, made here: a NaN written
248-
into the array later isn't read as NULL. What the table adds is its
249-
dimension columns, one value per row for each dimension, which
250-
long-form rows need. An array that isn't C-contiguous (Fortran order)
251-
is copied, once.
252-
253-
Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
254-
lazily loaded file), is larger than ``limit``, or has no rows; it is
255-
read lazily instead.
261+
of CPUs. So no data variable is copied, whatever the chunks: its Arrow
262+
column shares the array's memory, as does DataFusion's table, through
263+
Arrow's C data interface, and a write into the array in place shows in
264+
the next query, as it does for a lazy table. A NaN adds a validity
265+
bitmap beside it, one bit per value, made here: a NaN written into the
266+
array later isn't read as NULL. What the table adds is its dimension
267+
columns, one value per row for each dimension, which long-form rows
268+
need.
269+
270+
Returns ``None``, for the Dataset to be read lazily instead, when it
271+
isn't in memory (a Dask array or a lazily loaded file), is larger than
272+
``limit``, or has no rows; or when a data variable's Arrow column
273+
would have to be a copy of its array, which a write in place wouldn't
274+
reach (see `_shares_memory_with_arrow`).
256275
"""
257276
# `_in_memory` is xarray's own test for a variable whose values are
258277
# loaded: a NumPy array, not a Dask array or a lazily indexed file.
259278
if not all(var._in_memory for var in ds.variables.values()):
260279
return None
280+
if not all(
281+
_shares_memory_with_arrow(var.values) for var in ds.data_vars.values()
282+
):
283+
return None
261284
# TODO(#286): register a dimension that no data variable has (only a
262285
# coordinate does) as the lazy reader should, then drop this fallback.
263286
data_vars = list(ds.data_vars.values())

‎xarray_sql/sql.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,9 +80,9 @@ def from_dataset(
8080
A Dataset is read lazily, chunk by chunk, as queries run. One
8181
already in memory (NumPy, not Dask or a lazily opened file) and at
8282
most 64 MiB as a table is instead a DataFusion in-memory table,
83-
which is faster to scan many times. It shares the Dataset's memory,
84-
so changing values in place changes what queries return, as long
85-
as no value becomes or stops being NaN (which reads as NULL).
83+
which is faster to scan many times. Either way, changing values in
84+
place changes what queries return, as long as no value becomes or
85+
stops being NaN (which reads as NULL) in an in-memory table.
8686
8787
For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
8888
julian), a ``cftime()`` scalar UDF is automatically registered so

0 commit comments

Comments
 (0)