Skip to content

Commit d4e0b24

Browse files
alxmrsclaude
andcommitted
Register small in-memory Datasets as native in-memory tables
Every scan of a lazy table goes through datafusion-ffi, which rebuilds the session and the task context, every registered function included, on each scan and each partition's execution: about 1.5 ms per scan (#273). That 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. Context.from_dataset now converts a Dataset that is already in memory and at most 64 MiB as long-form Arrow once, at registration, one partition per chunk, and registers it as a native DataFusion in-memory table, which bypasses FFI. Lazy and larger Datasets keep the lazy table, with partition pruning, projection pushdown and exact statistics. A scan of a 1,040-row Dataset goes from 1.7 ms to 0.26 ms. On the GPT of the ML demos, the forward loss query goes from 9.2 s to 2.5 s. The statistics and pruning tests exercise the lazy table, so they register Dask-backed Datasets. Closes #272 Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 213084a commit d4e0b24

4 files changed

Lines changed: 158 additions & 14 deletions

File tree

‎tests/test_sql.py‎

Lines changed: 64 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,65 @@ 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, and
596+
queries over it give what a lazily read Dataset gives."""
597+
ds = xr.Dataset(
598+
{"air": (("time", "lat"), np.random.rand(6, 4))},
599+
coords={"time": np.arange(6), "lat": np.arange(4) * 2.5},
600+
)
601+
ctx = XarrayContext()
602+
ctx.from_dataset("eager", ds, chunks={"time": 2})
603+
ctx.from_dataset("lazy", ds.chunk(), chunks={"time": 2})
604+
605+
assert "XarrayScanExec" not in _scans(ctx, "eager")
606+
assert "XarrayScanExec" in _scans(ctx, "lazy")
607+
query = "SELECT time, lat, air FROM {} WHERE lat > 2 ORDER BY time, lat"
608+
pd.testing.assert_frame_equal(
609+
ctx.sql(query.format("eager")).to_pandas(),
610+
ctx.sql(query.format("lazy")).to_pandas(),
611+
)
612+
assert set(ctx.catalog().schema().names()) == {"eager", "lazy"}
613+
614+
615+
def test_mixed_dimension_in_memory_dataset_registers_native_tables():
616+
"""Each dimension group of a small in-memory Dataset is a native table
617+
in the Dataset's schema, with nothing else left in the catalog."""
618+
ds = xr.Dataset(
619+
{
620+
"surface": (("time", "lat"), np.random.rand(3, 2)),
621+
"column": (("time", "lat", "level"), np.random.rand(3, 2, 4)),
622+
},
623+
coords={
624+
"time": np.arange(3),
625+
"lat": np.arange(2),
626+
"level": np.arange(4),
627+
},
628+
)
629+
ctx = XarrayContext()
630+
ctx.from_dataset("era5", ds)
631+
632+
assert "XarrayScanExec" not in _scans(ctx, "era5.time_lat_level")
633+
counts = ctx.sql(
634+
"SELECT (SELECT COUNT(*) FROM era5.time_lat) AS surface, "
635+
"(SELECT COUNT(*) FROM era5.time_lat_level) AS column_"
636+
).to_pandas()
637+
assert counts.iloc[0].tolist() == [6, 24]
638+
assert set(ctx.catalog().schema().names()) == set()
639+
640+
641+
def test_large_in_memory_dataset_is_read_lazily():
642+
"""An in-memory Dataset over 64 MiB as a table stays a lazy table, so
643+
it isn't held twice."""
644+
rows = 64 * 2**20 // 16 + 1 # x (int64) and value (float64): 16 bytes
645+
ds = xr.Dataset({"value": (["x"], np.zeros(rows))})
646+
ctx = XarrayContext()
647+
ctx.from_dataset("big", ds)
648+
649+
assert "XarrayScanExec" in _scans(ctx, "big")

‎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: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,59 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader:
193193
return pa.RecordBatchReader.from_stream(reader)
194194

195195

196+
# The largest Dataset, in bytes of long-form Arrow, that Context.from_dataset
197+
# registers as a native DataFusion in-memory table rather than a lazy one.
198+
_IN_MEMORY_TABLE_LIMIT: int = 64 * 2**20
199+
200+
201+
def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
202+
"""An estimate of a table's size in bytes, counting 16 bytes for each
203+
value of a variable-width column."""
204+
widths: list[int] = [
205+
field.type.bit_width // 8 if pa.types.is_primitive(field.type) else 16
206+
for field in schema
207+
]
208+
return num_rows * sum(widths)
209+
210+
211+
def _in_memory_partitions(
212+
ds: xr.Dataset,
213+
chunks: Chunks = None,
214+
*,
215+
limit: int = _IN_MEMORY_TABLE_LIMIT,
216+
batch_size: int = DEFAULT_BATCH_SIZE,
217+
) -> list[list[pa.RecordBatch]] | None:
218+
"""A small in-memory Dataset's rows as Arrow partitions, or ``None``.
219+
220+
A lazy table pays a fixed cost on every scan, through DataFusion's FFI,
221+
of about a millisecond and a half. That is noise for a large Dataset,
222+
but a query that reads small tables many times pays it each time:
223+
DataFusion inlines a CTE or a view at every place that reads it. A
224+
Dataset that is already in memory gains nothing from being read lazily,
225+
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.
228+
229+
Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
230+
lazily loaded file), is larger than ``limit``, or has no rows; it is
231+
read lazily instead.
232+
"""
233+
# `_in_memory` is xarray's own test for a variable whose values are
234+
# loaded: a NumPy array, not a Dask array or a lazily indexed file.
235+
if not all(var._in_memory for var in ds.variables.values()):
236+
return None
237+
ds = _ensure_default_indexes(ds)
238+
schema = _parse_schema(ds)
239+
num_rows = int(np.prod([ds.sizes[d] for d in ds.dims]))
240+
if num_rows == 0 or _arrow_size(schema, num_rows) > limit:
241+
return None
242+
resolved = resolve_chunks(ds, chunks)
243+
return [
244+
list(iter_record_batches(ds.isel(block), schema, batch_size))
245+
for block in _block_slices_from_resolved(ds, resolved)
246+
]
247+
248+
196249
def read_xarray_table(
197250
ds: xr.Dataset,
198251
chunks: Chunks = None,

‎xarray_sql/sql.py‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import pyarrow as pa
12
import xarray as xr
23
from datafusion import SessionContext
34
from datafusion.catalog import Schema
@@ -12,7 +13,7 @@
1213
shared_coord_arrays,
1314
)
1415
from .ds import XarrayDataFrame
15-
from .reader import read_xarray_table
16+
from .reader import _in_memory_partitions, read_xarray_table
1617

1718
_proj: ModuleType | None
1819
try: # pyproj is an optional dependency (`pip install xarray-sql[geo]`).
@@ -74,6 +75,14 @@ def from_dataset(
7475
)
7576
ctx.sql('SELECT * FROM era5.surface')
7677
78+
A Dataset is read lazily, chunk by chunk, as queries run, unless it
79+
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.
85+
7786
For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
7887
julian), a ``cftime()`` scalar UDF is automatically registered so
7988
you can write ergonomic SQL filters::
@@ -158,13 +167,28 @@ def _from_dataset(
158167
register = (
159168
self.register_table if schema is None else schema.register_table
160169
)
161-
register(
162-
table_name,
163-
read_xarray_table(input_table, chunks, coord_arrays=coord_arrays),
164-
)
170+
partitions = _in_memory_partitions(input_table, chunks)
171+
if partitions is None:
172+
table = read_xarray_table(
173+
input_table, chunks, coord_arrays=coord_arrays
174+
)
175+
else:
176+
table = self._in_memory_table(partitions)
177+
register(table_name, table)
165178
self._maybe_register_cftime_udf(input_table)
166179
return self
167180

181+
def _in_memory_table(self, partitions: list[list[pa.RecordBatch]]):
182+
"""A native DataFusion in-memory table of ``partitions``."""
183+
before = set(self.catalog().schema().names())
184+
view = self.create_dataframe(partitions).into_view()
185+
# create_dataframe registers its in-memory table under a generated
186+
# name. The view holds the table itself, not the name, so the name
187+
# can go.
188+
for name in set(self.catalog().schema().names()) - before:
189+
self.deregister_table(name)
190+
return view
191+
168192
def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None:
169193
"""Auto-register a cftime() UDF for non-Gregorian cftime coordinates."""
170194
for coord_name in ds.dims:

0 commit comments

Comments
 (0)