Register small in-memory Datasets as native in-memory tables

tests/test_sql.py

@@ -45,7 +45,8 @@ def tracking_reader(ds, chunks=None, **kwargs):
     )
     ctx = XarrayContext()
 
-    ctx.from_dataset("forecast", ds, chunks={"lead_time": 2})
+    # A Dask-backed Dataset, so it is read lazily, and pruned.
+    ctx.from_dataset("forecast", ds.chunk(), chunks={"lead_time": 2})
     result = ctx.sql(
         "SELECT value FROM forecast "
         "WHERE lead_time >= INTERVAL '4 microseconds' ORDER BY value"
@@ -584,3 +585,130 @@ def test_single_table_all_coordless_dims(self):
         pd.testing.assert_frame_equal(
             result, expected, check_dtype=False, check_like=True
         )
+
+
+def _scans(ctx, table):
+    return ctx.sql(f"SELECT * FROM {table}").execution_plan().display_indent()
+
+
+def test_small_in_memory_dataset_is_a_native_table_with_the_same_rows():
+    """A small in-memory Dataset becomes a native in-memory table, which a
+    query's plan scans by its registered name, and queries over it give what
+    a lazily read Dataset gives."""
+    ds = xr.Dataset(
+        {"air": (("time", "lat"), np.random.rand(6, 4))},
+        coords={"time": np.arange(6), "lat": np.arange(4) * 2.5},
+    )
+    ctx = XarrayContext()
+    ctx.from_dataset("eager", ds, chunks={"time": 2})
+    ctx.from_dataset("lazy", ds.chunk(), chunks={"time": 2})
+
+    assert "XarrayScanExec" not in _scans(ctx, "eager")
+    assert "XarrayScanExec" in _scans(ctx, "lazy")
+    plan = ctx.sql("SELECT * FROM eager").logical_plan().display_indent()
+    assert plan.splitlines()[-1].strip() == "TableScan: eager"
+    query = "SELECT time, lat, air FROM {} WHERE lat > 2 ORDER BY time, lat"
+    pd.testing.assert_frame_equal(
+        ctx.sql(query.format("eager")).to_pandas(),
+        ctx.sql(query.format("lazy")).to_pandas(),
+    )
+    assert set(ctx.catalog().schema().names()) == {"eager", "lazy"}
+
+
+def test_mixed_dimension_in_memory_dataset_registers_native_tables():
+    """Each dimension group of a small in-memory Dataset is a native table
+    in the Dataset's schema, with nothing else left in the catalog."""
+    ds = xr.Dataset(
+        {
+            "surface": (("time", "lat"), np.random.rand(3, 2)),
+            "column": (("time", "lat", "level"), np.random.rand(3, 2, 4)),
+        },
+        coords={
+            "time": np.arange(3),
+            "lat": np.arange(2),
+            "level": np.arange(4),
+        },
+    )
+    ctx = XarrayContext()
+    ctx.from_dataset("era5", ds)
+
+    assert "XarrayScanExec" not in _scans(ctx, "era5.time_lat_level")
+    counts = ctx.sql(
+        "SELECT (SELECT COUNT(*) FROM era5.time_lat) AS surface, "
+        "(SELECT COUNT(*) FROM era5.time_lat_level) AS column_"
+    ).to_pandas()
+    assert counts.iloc[0].tolist() == [6, 24]
+    assert set(ctx.catalog().schema().names()) == set()
+
+
+def test_large_in_memory_dataset_is_read_lazily():
+    """An in-memory Dataset over 64 MiB as a table stays a lazy table, so
+    it isn't held twice."""
+    rows = 64 * 2**20 // 16 + 1  # x (int64) and value (float64): 16 bytes
+    ds = xr.Dataset({"value": (["x"], np.zeros(rows))})
+    ctx = XarrayContext()
+    ctx.from_dataset("big", ds)
+
+    assert "XarrayScanExec" in _scans(ctx, "big")
+
+
+@pytest.mark.parametrize("chunks", [None, {"a": 2}, {"b": 3}, {"a": 2, "b": 2}])
+def test_in_memory_table_shares_the_dataset_s_memory(chunks):
+    """A native table reads the Dataset's arrays, not a copy, whatever the
+    chunks: a change made in place shows in the next query, as it does for
+    a lazy table."""
+    ds = xr.Dataset({"w": (("a", "b"), np.zeros((4, 6)))})
+    ctx = XarrayContext()
+    ctx.from_dataset("t", ds, chunks=chunks)
+
+    ds["w"].values[...] = 7.0
+
+    assert ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [7.0] * 24
+
+
+@pytest.mark.parametrize(
+    "values, written",
+    [
+        (np.asfortranarray(np.zeros((4, 6))), 7.0),
+        (np.zeros((6, 4)).T, 7.0),
+        (np.zeros((4, 12))[:, ::2], 7.0),
+        (np.zeros((4, 6), dtype=bool), True),
+        (np.full((4, 6), "a"), "b"),
+    ],
+    ids=["fortran-order", "transposed", "strided", "bool", "str"],
+)
+def test_in_place_write_shows_where_arrow_would_copy(values, written):
+    """A write in place shows in the next query even for an array whose
+    Arrow column would be a copy, which isn't an in-memory table."""
+    ds = xr.Dataset({"w": (("a", "b"), values)})
+    ctx = XarrayContext()
+    ctx.from_dataset("t", ds)
+
+    ds["w"].values[...] = written
+
+    assert (
+        ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [written] * 24
+    )
+
+
+def test_in_memory_dataset_over_the_limit_in_arrow_is_read_lazily():
+    """The 64 MiB limit is on the table as Arrow, which a long string can
+    exceed when an estimate of its size doesn't."""
+    ds = xr.Dataset({"s": (["x"], np.array([b"a" * 40 * 2**20] * 2))})
+    ctx = XarrayContext()
+    ctx.from_dataset("strings", ds)
+
+    assert "XarrayScanExec" in _scans(ctx, "strings")
+
+
+def test_dimension_no_data_variable_has_registers():
+    """A coordinate on a dimension that no data variable has doesn't stop
+    an in-memory Dataset from registering."""
+    ds = xr.Dataset(
+        {"t": (("x",), np.arange(3.0))},
+        coords={"x": [0, 1, 2], "y": [10, 20]},
+    )
+    ctx = XarrayContext()
+    ctx.from_dataset("t", ds)
+
+    assert "t" in ctx.catalog().schema().names()

tests/test_stats.py

@@ -4,6 +4,10 @@
 so the exact statistics ``XarrayScanExec`` reports are visible to the query
 optimizer: num_rows (product of a chunk's dimension sizes), total byte size,
 and per dimension-column min/max bounds. These tests pin that behaviour.
+
+``XarrayScanExec`` scans a lazy table, so each test registers a Dask-backed
+Dataset (``ds.chunk()``): a small Dataset that is already in memory becomes
+a native DataFusion in-memory table instead.
 """
 
 import numpy as np
@@ -30,7 +34,7 @@ def test_exact_rows_in_scan_statistics():
         },
     )
     ctx = XarrayContext()
-    ctx.from_dataset("air", ds, chunks={"time": 50})
+    ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
     plan = _explain(ctx, "SELECT lat, lon, air FROM air")
     total = 100 * 4 * 5
     assert f"Rows=Exact({total})" in plan
@@ -47,7 +51,7 @@ def test_exact_byte_size_in_scan_statistics():
         },
     )
     ctx = XarrayContext()
-    ctx.from_dataset("air", ds, chunks={"time": 50})
+    ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
     plan = _explain(ctx, "SELECT lat, lon, air FROM air")
     # 2000 rows x (lat int64 + lon int64 + air float64) = 2000 x 24 bytes.
     assert f"Bytes=Exact({100 * 4 * 5 * 24})" in plan
@@ -69,7 +73,7 @@ def test_dimension_column_min_max_in_scan_statistics():
         },
     )
     ctx = XarrayContext()
-    ctx.from_dataset("air", ds, chunks={"time": 50})
+    ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
     plan = _explain(ctx, "SELECT lat, lon, air FROM air")
     # lat spans 0..3, lon spans 0..4, both never null.
     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):
         coords={"x": coord},
     )
     ctx = XarrayContext()
-    ctx.from_dataset("values", ds, chunks={"x": 2})
+    ctx.from_dataset("values", ds.chunk(), chunks={"x": 2})
 
     plan = _explain(ctx, "SELECT x, value FROM values")
 
@@ -114,7 +118,7 @@ def test_missing_partition_bound_removes_global_coordinate_statistics():
         coords={"x": np.asarray([0.0, 1.0, np.nan, 3.0])},
     )
     ctx = XarrayContext()
-    ctx.from_dataset("values", ds, chunks={"x": 2})
+    ctx.from_dataset("values", ds.chunk(), chunks={"x": 2})
 
     plan = _explain(ctx, "SELECT x, value FROM values")
 
@@ -134,7 +138,7 @@ def test_count_star_answered_from_statistics():
         },
     )
     ctx = XarrayContext()
-    ctx.from_dataset("air", ds, chunks={"time": 50})
+    ctx.from_dataset("air", ds.chunk(), chunks={"time": 50})
     n = ctx.sql("SELECT COUNT(*) AS n FROM air").to_pandas()["n"][0]
     assert int(n) == 100 * 4 * 5
 
@@ -159,8 +163,8 @@ def test_join_picks_small_build_side():
         coords={"lat": np.arange(8), "lon": np.arange(8)},
     )
     ctx = XarrayContext()
-    ctx.from_dataset("big", big, chunks={"time": 50})
-    ctx.from_dataset("small", small, chunks={"lat": 8})
+    ctx.from_dataset("big", big.chunk(), chunks={"time": 50})
+    ctx.from_dataset("small", small.chunk(), chunks={"lat": 8})
 
     plan = _explain(
         ctx,

xarray_sql/reader.py

@@ -10,7 +10,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterator
+import itertools
+import os
+import sys
+from collections.abc import Callable, Iterable, Iterator
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -34,6 +37,17 @@
 if TYPE_CHECKING:
     from ._native import LazyArrowStreamTable
 
+if sys.version_info >= (3, 12):
+    _batched = itertools.batched
+else:
+
+    def _batched(iterable: Iterable, n: int) -> Iterator[tuple]:
+        """``itertools.batched``, which Python 3.12 added: ``iterable`` in
+        tuples of ``n``, the last one shorter."""
+        it = iter(iterable)
+        while batch := tuple(itertools.islice(it, n)):
+            yield batch
+
 
 class XarrayRecordBatchReader:
     """A lazy Arrow stream reader for xarray Datasets.
@@ -193,6 +207,114 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader:
     return pa.RecordBatchReader.from_stream(reader)
 
 
+# The largest Dataset, in bytes of long-form Arrow, that Context.from_dataset
+# registers as a native DataFusion in-memory table rather than a lazy one.
+_IN_MEMORY_TABLE_LIMIT: int = 64 * 2**20
+
+
+def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
+    """An estimate of a table's size in bytes, counting 16 bytes for each
+    value of a variable-width column."""
+    widths: list[int] = [
+        field.type.bit_width // 8 if pa.types.is_primitive(field.type) else 16
+        for field in schema
+    ]
+    return num_rows * sum(widths)
+
+
+def _shares_memory_with_arrow(values: np.ndarray) -> bool:
+    """Whether ``values``'s Arrow column, its rows in C order, can be its
+    own memory rather than a copy.
+
+    Arrow lays out a number or a datetime as NumPy does, in the machine's
+    byte order. It packs booleans into bits, and keeps strings in buffers
+    of its own, so those are copies; so is an array that isn't
+    C-contiguous, such as one in Fortran order or a transposed or strided
+    view, whose rows in C order aren't one run of memory.
+    """
+    return (
+        values.dtype.kind in "iufmM"
+        and values.dtype.isnative
+        and values.flags.c_contiguous
+    )
+
+
+def _in_memory_partitions(
+    ds: xr.Dataset,
+    chunks: Chunks = None,
+    *,
+    limit: int = _IN_MEMORY_TABLE_LIMIT,
+    batch_size: int = DEFAULT_BATCH_SIZE,
+) -> list[list[pa.RecordBatch]] | None:
+    """A small in-memory Dataset's rows as Arrow partitions, or ``None``.
+
+    A lazy table pays a fixed cost on every scan, through DataFusion's FFI,
+    of about a millisecond and a half. That is noise for a large Dataset,
+    but a query that reads small tables many times pays it each time:
+    DataFusion inlines a CTE or a view at every place that reads it. A
+    Dataset that is already in memory gains nothing from being read lazily,
+    so when it is also at most ``limit`` bytes as long-form Arrow, its rows
+    are converted once, for a native in-memory table.
+
+    The rows are in C order, and each partition is a contiguous range of
+    them, as many partitions as ``chunks`` makes chunks, up to the number
+    of CPUs. So no data variable is copied, whatever the chunks: its Arrow
+    column shares the array's memory, as does DataFusion's table, through
+    Arrow's C data interface, and a write into the array in place shows in
+    the next query, as it does for a lazy table. A NaN adds a validity
+    bitmap beside it, one bit per value, made here: a NaN written into the
+    array later isn't read as NULL. What the table adds is its dimension
+    columns, one value per row for each dimension, which long-form rows
+    need.
+
+    Returns ``None``, for the Dataset to be read lazily instead, when it
+    isn't in memory (a Dask array or a lazily loaded file), is larger than
+    ``limit``, or has no rows; or when a data variable's Arrow column
+    would have to be a copy of its array, which a write in place wouldn't
+    reach (see `_shares_memory_with_arrow`).
+    """
+    # `_in_memory` is xarray's own test for a variable whose values are
+    # loaded: a NumPy array, not a Dask array or a lazily indexed file.
+    if not all(var._in_memory for var in ds.variables.values()):
+        return None
+    if not all(
+        _shares_memory_with_arrow(var.values) for var in ds.data_vars.values()
+    ):
+        return None
+    # TODO(#286): register a dimension that no data variable has (only a
+    # coordinate does) as the lazy reader should, then drop this fallback.
+    data_vars = list(ds.data_vars.values())
+    if data_vars and set(ds.dims) - set(data_vars[0].dims):
+        return None
+    ds = _ensure_default_indexes(ds)
+    schema = _parse_schema(ds)
+    num_rows = int(np.prod([ds.sizes[d] for d in ds.dims]))
+    if num_rows == 0 or _arrow_size(schema, num_rows) > limit:
+        return None
+    # A partition per chunk, for the parallelism the lazy table has, but cut
+    # from the rows rather than by `isel`: a chunk of a trailing dimension
+    # isn't contiguous in memory, and flattening it would copy. More
+    # partitions than cores only adds overhead, since an in-memory table has
+    # no partitions to prune.
+    count = min(
+        sum(
+            1
+            for _ in _block_slices_from_resolved(ds, resolve_chunks(ds, chunks))
+        ),
+        os.cpu_count() or 1,
+    )
+    # `per_partition` batches of `size` rows each make `count` partitions,
+    # with no batch over `batch_size` rows.
+    per_partition = -(-num_rows // (count * batch_size))
+    size = -(-num_rows // (count * per_partition))
+    batches = list(iter_record_batches(ds, schema, size))
+    # The estimate undercounts booleans and long strings, so check the
+    # batches themselves.
+    if sum(batch.nbytes for batch in batches) > limit:
+        return None
+    return [list(p) for p in _batched(batches, per_partition)]
+
+
 def read_xarray_table(
     ds: xr.Dataset,
     chunks: Chunks = None,