Skip to content

Commit ddf8a47

Browse files
alxmrsclaude
andcommitted
Share the Dataset's memory in an in-memory table, whatever the chunks
A data variable's Arrow column already shared its NumPy array's memory, and DataFusion's table keeps those buffers through Arrow's C data interface. But a partition per chunk flattened each chunk with isel, and a chunk of a trailing dimension isn't contiguous, so it was copied. Partitions are now contiguous ranges of the C-order rows, as many as there are chunks, up to the number of CPUs, so no chunking copies a C-contiguous array. The table adds only its dimension columns. The docstring no longer says the Dataset is held twice. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent f22d938 commit ddf8a47

3 files changed

Lines changed: 51 additions & 10 deletions

File tree

‎tests/test_sql.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -647,3 +647,17 @@ def test_large_in_memory_dataset_is_read_lazily():
647647
ctx.from_dataset("big", ds)
648648

649649
assert "XarrayScanExec" in _scans(ctx, "big")
650+
651+
652+
@pytest.mark.parametrize("chunks", [None, {"a": 2}, {"b": 3}, {"a": 2, "b": 2}])
653+
def test_in_memory_table_shares_the_dataset_s_memory(chunks):
654+
"""A native table reads the Dataset's arrays, not a copy, whatever the
655+
chunks: a change made in place shows in the next query, as it does for
656+
a lazy table."""
657+
ds = xr.Dataset({"w": (("a", "b"), np.zeros((4, 6)))})
658+
ctx = XarrayContext()
659+
ctx.from_dataset("t", ds, chunks=chunks)
660+
661+
ds["w"].values[...] = 7.0
662+
663+
assert ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [7.0] * 24

‎xarray_sql/reader.py‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
from __future__ import annotations
1212

13+
import os
1314
from collections.abc import Callable, Iterator
1415
from typing import TYPE_CHECKING
1516

@@ -223,8 +224,16 @@ def _in_memory_partitions(
223224
DataFusion inlines a CTE or a view at every place that reads it. A
224225
Dataset that is already in memory gains nothing from being read lazily,
225226
so when it is also at most ``limit`` bytes as long-form Arrow, its rows
226-
are converted once, one partition per chunk, for a native in-memory
227-
table.
227+
are converted once, for a native in-memory table.
228+
229+
The rows are in C order, and each partition is a contiguous range of
230+
them, as many partitions as ``chunks`` makes chunks, up to the number
231+
of CPUs. So a data variable whose array is C-contiguous isn't copied,
232+
whatever the chunks: its Arrow column shares the array's memory, as
233+
does DataFusion's table, through Arrow's C data interface. A NaN adds a validity bitmap beside it, one
234+
bit per value. What the table adds is its dimension columns, one value
235+
per row for each dimension, which long-form rows need. An array that
236+
isn't C-contiguous (Fortran order) is copied, once.
228237
229238
Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
230239
lazily loaded file), is larger than ``limit``, or has no rows; it is
@@ -239,10 +248,26 @@ def _in_memory_partitions(
239248
num_rows = int(np.prod([ds.sizes[d] for d in ds.dims]))
240249
if num_rows == 0 or _arrow_size(schema, num_rows) > limit:
241250
return None
242-
resolved = resolve_chunks(ds, chunks)
251+
# A partition per chunk, for the parallelism the lazy table has, but cut
252+
# from the rows rather than by `isel`: a chunk of a trailing dimension
253+
# isn't contiguous in memory, and flattening it would copy. More
254+
# partitions than cores only adds overhead, since an in-memory table has
255+
# no partitions to prune.
256+
count = min(
257+
sum(
258+
1
259+
for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
260+
),
261+
os.cpu_count() or 1,
262+
)
263+
# `per_partition` batches of `size` rows each make `count` partitions,
264+
# with no batch over `batch_size` rows.
265+
per_partition = -(-num_rows // (count * batch_size))
266+
size = -(-num_rows // (count * per_partition))
267+
batches = list(iter_record_batches(ds, schema, size))
243268
return [
244-
list(iter_record_batches(ds.isel(block), schema, batch_size))
245-
for block in _block_slices_from_resolved(ds, resolved)
269+
batches[i : i + per_partition]
270+
for i in range(0, len(batches), per_partition)
246271
]
247272

248273

‎xarray_sql/sql.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,11 +77,13 @@ def from_dataset(
7777
7878
A Dataset is read lazily, chunk by chunk, as queries run, unless it
7979
is already in memory (NumPy arrays, not Dask or a lazily opened
80-
file) and at most 64 MiB as a table. Such a Dataset is converted
81-
to Arrow once, here, and registered as a native DataFusion
82-
in-memory table, which is faster to scan many times, as a query
83-
does when it reads a table in several places. It is held twice,
84-
as arrays and as Arrow.
80+
file) and at most 64 MiB as a table. Such a Dataset is registered
81+
as a native DataFusion in-memory table, which is faster to scan
82+
many times, as a query does when it reads a table in several
83+
places. The table shares its data variables' memory with the
84+
Dataset rather than copying it, so, as with a lazy table, changing
85+
the Dataset's values in place changes what queries return. It adds
86+
a column per dimension, one value per row.
8587
8688
For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
8789
julian), a ``cftime()`` scalar UDF is automatically registered so

0 commit comments

Comments
 (0)