Skip to content

Commit c70b1b2

Browse files
alxmrsclaude
andauthored
Register small in-memory Datasets as native in-memory tables (#275)
Fixes #272. **Stacked on #274** (DataFusion 55): review only the commits after #274's, and merge after #274. ## Why Every scan of a lazy xarray-sql table goes through `datafusion-ffi`, which costs about 1.5 ms per scan. On every scan it rebuilds the session, and on every partition's execution the task context, every registered function included (#273, a draft to re-file upstream). That cost is noise for one scan of a large Dataset. But DataFusion inlines a CTE or a view at every reference, so a query over small tables can scan them thousands of times. The GPT in the ML demos (#269) scans its tables 2,989 times in one forward pass, and this overhead was 6.7 s of the 9.2 s. ## What - `Context.from_dataset` converts a Dataset once, at registration, and registers it as a native DataFusion in-memory table, bypassing FFI, when the Dataset: - is already in memory (every variable is NumPy, per xarray's `Variable._in_memory`; not Dask or a lazily opened file); - is at most 64 MiB as long-form Arrow. The table has a partition per chunk, up to the CPU count, so chunking still sets parallelism. - Every other Dataset keeps the lazy table, with partition pruning, projection pushdown and exact statistics. - Mixed-dimension Datasets register each dimension group in their schema, as an in-memory table taken from `register_record_batches` under a placeholder name, which is then removed. The table itself is registered, not a view of it, so a query's plan scans the registered name; ddx needs that name to find the table it differentiates. - **Memory:** the table shares the Dataset's arrays rather than copying them. NumPy's buffers pass through `pa.array` and Arrow's C data interface into DataFusion unchanged, and partitions are contiguous row ranges, so no chunking forces a copy. The table adds its dimension columns, one value per row for each dimension. A Fortran-ordered array is copied once. ## Results | | before | after | |---|---:|---:| | `SELECT *` over a 1,040-row in-memory Dataset | 1.7 ms | 0.26 ms | | GPT forward loss (DataFusion 54, the same tables registered in memory) | 9.2 s | 2.5 s | | GPT training step through ddx, cached program | 1.17 s | 0.81 s | The GPT numbers come from registering the demo's tables with `register_record_batches`, which is what this change does. The demo itself can't run on this branch until ddxdb supports DataFusion 55. ## Tests - **New:** - a small in-memory Dataset becomes a native table and returns the same rows as the same Dataset read lazily; - a mixed-dimension Dataset registers native tables in its schema and leaves nothing extra in the catalog; - an in-memory Dataset over 64 MiB stays lazy; - the table shares the Dataset's memory under four chunkings: a change made in place shows in the next query, as it does for a lazy table. - **Changed:** the statistics and pruning tests (`test_stats.py`, and the duration-pruning test in `test_sql.py`) exercise the lazy table, so they now register Dask-backed Datasets (`ds.chunk()`). - `pytest -m "not integration"`: 510 passed, 277 skipped. mypy 1.11.2 is clean. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Small, in-memory datasets that meet supported layout requirements can be queried through native tables; larger or chunked datasets continue to be read lazily. * Queries reflect in-place value changes to registered datasets. If missing-value positions change, register the dataset again for accurate results. * **Tests** * Expanded coverage for eager and lazy queries, dataset sizes and dimensions, shared value updates, and statistics. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
1 parent 7f4ed41 commit c70b1b2

4 files changed

Lines changed: 296 additions & 15 deletions

File tree

‎tests/test_sql.py‎

Lines changed: 129 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ def tracking_reader(ds, chunks=None, **kwargs):
4545
)
4646
ctx = XarrayContext()
4747

48-
ctx.from_dataset("forecast", ds, chunks={"lead_time": 2})
48+
# A Dask-backed Dataset, so it is read lazily, and pruned.
49+
ctx.from_dataset("forecast", ds.chunk(), chunks={"lead_time": 2})
4950
result = ctx.sql(
5051
"SELECT value FROM forecast "
5152
"WHERE lead_time >= INTERVAL '4 microseconds' ORDER BY value"
@@ -584,3 +585,130 @@ def test_single_table_all_coordless_dims(self):
584585
pd.testing.assert_frame_equal(
585586
result, expected, check_dtype=False, check_like=True
586587
)
588+
589+
590+
def _scans(ctx, table):
591+
return ctx.sql(f"SELECT * FROM {table}").execution_plan().display_indent()
592+
593+
594+
def test_small_in_memory_dataset_is_a_native_table_with_the_same_rows():
595+
"""A small in-memory Dataset becomes a native in-memory table, which a
596+
query's plan scans by its registered name, and queries over it give what
597+
a lazily read Dataset gives."""
598+
ds = xr.Dataset(
599+
{"air": (("time", "lat"), np.random.rand(6, 4))},
600+
coords={"time": np.arange(6), "lat": np.arange(4) * 2.5},
601+
)
602+
ctx = XarrayContext()
603+
ctx.from_dataset("eager", ds, chunks={"time": 2})
604+
ctx.from_dataset("lazy", ds.chunk(), chunks={"time": 2})
605+
606+
assert "XarrayScanExec" not in _scans(ctx, "eager")
607+
assert "XarrayScanExec" in _scans(ctx, "lazy")
608+
plan = ctx.sql("SELECT * FROM eager").logical_plan().display_indent()
609+
assert plan.splitlines()[-1].strip() == "TableScan: eager"
610+
query = "SELECT time, lat, air FROM {} WHERE lat > 2 ORDER BY time, lat"
611+
pd.testing.assert_frame_equal(
612+
ctx.sql(query.format("eager")).to_pandas(),
613+
ctx.sql(query.format("lazy")).to_pandas(),
614+
)
615+
assert set(ctx.catalog().schema().names()) == {"eager", "lazy"}
616+
617+
618+
def test_mixed_dimension_in_memory_dataset_registers_native_tables():
619+
"""Each dimension group of a small in-memory Dataset is a native table
620+
in the Dataset's schema, with nothing else left in the catalog."""
621+
ds = xr.Dataset(
622+
{
623+
"surface": (("time", "lat"), np.random.rand(3, 2)),
624+
"column": (("time", "lat", "level"), np.random.rand(3, 2, 4)),
625+
},
626+
coords={
627+
"time": np.arange(3),
628+
"lat": np.arange(2),
629+
"level": np.arange(4),
630+
},
631+
)
632+
ctx = XarrayContext()
633+
ctx.from_dataset("era5", ds)
634+
635+
assert "XarrayScanExec" not in _scans(ctx, "era5.time_lat_level")
636+
counts = ctx.sql(
637+
"SELECT (SELECT COUNT(*) FROM era5.time_lat) AS surface, "
638+
"(SELECT COUNT(*) FROM era5.time_lat_level) AS column_"
639+
).to_pandas()
640+
assert counts.iloc[0].tolist() == [6, 24]
641+
assert set(ctx.catalog().schema().names()) == set()
642+
643+
644+
def test_large_in_memory_dataset_is_read_lazily():
645+
"""An in-memory Dataset over 64 MiB as a table stays a lazy table, so
646+
it isn't held twice."""
647+
rows = 64 * 2**20 // 16 + 1 # x (int64) and value (float64): 16 bytes
648+
ds = xr.Dataset({"value": (["x"], np.zeros(rows))})
649+
ctx = XarrayContext()
650+
ctx.from_dataset("big", ds)
651+
652+
assert "XarrayScanExec" in _scans(ctx, "big")
653+
654+
655+
@pytest.mark.parametrize("chunks", [None, {"a": 2}, {"b": 3}, {"a": 2, "b": 2}])
656+
def test_in_memory_table_shares_the_dataset_s_memory(chunks):
657+
"""A native table reads the Dataset's arrays, not a copy, whatever the
658+
chunks: a change made in place shows in the next query, as it does for
659+
a lazy table."""
660+
ds = xr.Dataset({"w": (("a", "b"), np.zeros((4, 6)))})
661+
ctx = XarrayContext()
662+
ctx.from_dataset("t", ds, chunks=chunks)
663+
664+
ds["w"].values[...] = 7.0
665+
666+
assert ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [7.0] * 24
667+
668+
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+
694+
def test_in_memory_dataset_over_the_limit_in_arrow_is_read_lazily():
695+
"""The 64 MiB limit is on the table as Arrow, which a long string can
696+
exceed when an estimate of its size doesn't."""
697+
ds = xr.Dataset({"s": (["x"], np.array([b"a" * 40 * 2**20] * 2))})
698+
ctx = XarrayContext()
699+
ctx.from_dataset("strings", ds)
700+
701+
assert "XarrayScanExec" in _scans(ctx, "strings")
702+
703+
704+
def test_dimension_no_data_variable_has_registers():
705+
"""A coordinate on a dimension that no data variable has doesn't stop
706+
an in-memory Dataset from registering."""
707+
ds = xr.Dataset(
708+
{"t": (("x",), np.arange(3.0))},
709+
coords={"x": [0, 1, 2], "y": [10, 20]},
710+
)
711+
ctx = XarrayContext()
712+
ctx.from_dataset("t", ds)
713+
714+
assert "t" in ctx.catalog().schema().names()

‎tests/test_stats.py‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@
44
so the exact statistics ``XarrayScanExec`` reports are visible to the query
55
optimizer: num_rows (product of a chunk's dimension sizes), total byte size,
66
and per dimension-column min/max bounds. These tests pin that behaviour.
7+
8+
``XarrayScanExec`` scans a lazy table, so each test registers a Dask-backed
9+
Dataset (``ds.chunk()``): a small Dataset that is already in memory becomes
10+
a native DataFusion in-memory table instead.
711
"""
812

913
import numpy as np
@@ -30,7 +34,7 @@ def test_exact_rows_in_scan_statistics():
3034
},
3135
)
3236
ctx = XarrayContext()
33-
ctx.from_dataset("air", ds, chunks={"time": 50})
37+
ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
3438
plan = _explain(ctx, "SELECT lat, lon, air FROM air")
3539
total = 100 * 4 * 5
3640
assert f"Rows=Exact({total})" in plan
@@ -47,7 +51,7 @@ def test_exact_byte_size_in_scan_statistics():
4751
},
4852
)
4953
ctx = XarrayContext()
50-
ctx.from_dataset("air", ds, chunks={"time": 50})
54+
ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
5155
plan = _explain(ctx, "SELECT lat, lon, air FROM air")
5256
# 2000 rows x (lat int64 + lon int64 + air float64) = 2000 x 24 bytes.
5357
assert f"Bytes=Exact({100 * 4 * 5 * 24})" in plan
@@ -69,7 +73,7 @@ def test_dimension_column_min_max_in_scan_statistics():
6973
},
7074
)
7175
ctx = XarrayContext()
72-
ctx.from_dataset("air", ds, chunks={"time": 50})
76+
ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
7377
plan = _explain(ctx, "SELECT lat, lon, air FROM air")
7478
# lat spans 0..3, lon spans 0..4, both never null.
7579
assert "Min=Exact(Int64(0)) Max=Exact(Int64(3)) Null=Exact(0)" in plan
@@ -101,7 +105,7 @@ def test_coordinate_statistics_preserve_logical_type(coord, expected):
101105
coords={"x": coord},
102106
)
103107
ctx = XarrayContext()
104-
ctx.from_dataset("values", ds, chunks={"x": 2})
108+
ctx.from_dataset("values", ds.chunk(), chunks={"x": 2})
105109

106110
plan = _explain(ctx, "SELECT x, value FROM values")
107111

@@ -114,7 +118,7 @@ def test_missing_partition_bound_removes_global_coordinate_statistics():
114118
coords={"x": np.asarray([0.0, 1.0, np.nan, 3.0])},
115119
)
116120
ctx = XarrayContext()
117-
ctx.from_dataset("values", ds, chunks={"x": 2})
121+
ctx.from_dataset("values", ds.chunk(), chunks={"x": 2})
118122

119123
plan = _explain(ctx, "SELECT x, value FROM values")
120124

@@ -134,7 +138,7 @@ def test_count_star_answered_from_statistics():
134138
},
135139
)
136140
ctx = XarrayContext()
137-
ctx.from_dataset("air", ds, chunks={"time": 50})
141+
ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
138142
n = ctx.sql("SELECT COUNT(*) AS n FROM air").to_pandas()["n"][0]
139143
assert int(n) == 100 * 4 * 5
140144

@@ -159,8 +163,8 @@ def test_join_picks_small_build_side():
159163
coords={"lat": np.arange(8), "lon": np.arange(8)},
160164
)
161165
ctx = XarrayContext()
162-
ctx.from_dataset("big", big, chunks={"time": 50})
163-
ctx.from_dataset("small", small, chunks={"lat": 8})
166+
ctx.from_dataset("big", big.chunk(), chunks={"time": 50})
167+
ctx.from_dataset("small", small.chunk(), chunks={"lat": 8})
164168

165169
plan = _explain(
166170
ctx,

‎xarray_sql/reader.py‎

Lines changed: 123 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@
1010

1111
from __future__ import annotations
1212

13-
from collections.abc import Callable, Iterator
13+
import itertools
14+
import os
15+
import sys
16+
from collections.abc import Callable, Iterable, Iterator
1417
from typing import TYPE_CHECKING
1518

1619
import numpy as np
@@ -34,6 +37,17 @@
3437
if TYPE_CHECKING:
3538
from ._native import LazyArrowStreamTable
3639

40+
if sys.version_info >= (3, 12):
41+
_batched = itertools.batched
42+
else:
43+
44+
def _batched(iterable: Iterable, n: int) -> Iterator[tuple]:
45+
"""``itertools.batched``, which Python 3.12 added: ``iterable`` in
46+
tuples of ``n``, the last one shorter."""
47+
it = iter(iterable)
48+
while batch := tuple(itertools.islice(it, n)):
49+
yield batch
50+
3751

3852
class XarrayRecordBatchReader:
3953
"""A lazy Arrow stream reader for xarray Datasets.
@@ -193,6 +207,114 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader:
193207
return pa.RecordBatchReader.from_stream(reader)
194208

195209

210+
# The largest Dataset, in bytes of long-form Arrow, that Context.from_dataset
211+
# registers as a native DataFusion in-memory table rather than a lazy one.
212+
_IN_MEMORY_TABLE_LIMIT: int = 64 * 2**20
213+
214+
215+
def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
216+
"""An estimate of a table's size in bytes, counting 16 bytes for each
217+
value of a variable-width column."""
218+
widths: list[int] = [
219+
field.type.bit_width // 8 if pa.types.is_primitive(field.type) else 16
220+
for field in schema
221+
]
222+
return num_rows * sum(widths)
223+
224+
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+
242+
def _in_memory_partitions(
243+
ds: xr.Dataset,
244+
chunks: Chunks = None,
245+
*,
246+
limit: int = _IN_MEMORY_TABLE_LIMIT,
247+
batch_size: int = DEFAULT_BATCH_SIZE,
248+
) -> list[list[pa.RecordBatch]] | None:
249+
"""A small in-memory Dataset's rows as Arrow partitions, or ``None``.
250+
251+
A lazy table pays a fixed cost on every scan, through DataFusion's FFI,
252+
of about a millisecond and a half. That is noise for a large Dataset,
253+
but a query that reads small tables many times pays it each time:
254+
DataFusion inlines a CTE or a view at every place that reads it. A
255+
Dataset that is already in memory gains nothing from being read lazily,
256+
so when it is also at most ``limit`` bytes as long-form Arrow, its rows
257+
are converted once, for a native in-memory table.
258+
259+
The rows are in C order, and each partition is a contiguous range of
260+
them, as many partitions as ``chunks`` makes chunks, up to the number
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`).
275+
"""
276+
# `_in_memory` is xarray's own test for a variable whose values are
277+
# loaded: a NumPy array, not a Dask array or a lazily indexed file.
278+
if not all(var._in_memory for var in ds.variables.values()):
279+
return None
280+
if not all(
281+
_shares_memory_with_arrow(var.values) for var in ds.data_vars.values()
282+
):
283+
return None
284+
# TODO(#286): register a dimension that no data variable has (only a
285+
# coordinate does) as the lazy reader should, then drop this fallback.
286+
data_vars = list(ds.data_vars.values())
287+
if data_vars and set(ds.dims) - set(data_vars[0].dims):
288+
return None
289+
ds = _ensure_default_indexes(ds)
290+
schema = _parse_schema(ds)
291+
num_rows = int(np.prod([ds.sizes[d] for d in ds.dims]))
292+
if num_rows == 0 or _arrow_size(schema, num_rows) > limit:
293+
return None
294+
# A partition per chunk, for the parallelism the lazy table has, but cut
295+
# from the rows rather than by `isel`: a chunk of a trailing dimension
296+
# isn't contiguous in memory, and flattening it would copy. More
297+
# partitions than cores only adds overhead, since an in-memory table has
298+
# no partitions to prune.
299+
count = min(
300+
sum(
301+
1
302+
for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
303+
),
304+
os.cpu_count() or 1,
305+
)
306+
# `per_partition` batches of `size` rows each make `count` partitions,
307+
# with no batch over `batch_size` rows.
308+
per_partition = -(-num_rows // (count * batch_size))
309+
size = -(-num_rows // (count * per_partition))
310+
batches = list(iter_record_batches(ds, schema, size))
311+
# The estimate undercounts booleans and long strings, so check the
312+
# batches themselves.
313+
if sum(batch.nbytes for batch in batches) > limit:
314+
return None
315+
return [list(p) for p in _batched(batches, per_partition)]
316+
317+
196318
def read_xarray_table(
197319
ds: xr.Dataset,
198320
chunks: Chunks = None,

0 commit comments

Comments
 (0)