From 7e9fea2bb7bad02030a51d5e4679c7e163244a07 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Tue, 20 Oct 2026 23:32:11 -0700
Subject: [PATCH 1/6] 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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_sql.py    | 65 +++++++++++++++++++++++++++++++++++++++++++-
 tests/test_stats.py  | 20 ++++++++------
 xarray_sql/reader.py | 53 ++++++++++++++++++++++++++++++++++++
 xarray_sql/sql.py    | 34 +++++++++++++++++++----
 4 files changed, 158 insertions(+), 14 deletions(-)

diff --git a/tests/test_sql.py b/tests/test_sql.py
index 57d69b9b..cffa47d0 100644
--- a/tests/test_sql.py
+++ b/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,65 @@ 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, 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")
+    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")
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 619c5994..3c3acb45 100644
--- a/tests/test_stats.py
+++ b/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,
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index d2c78c4a..2b27776b 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -193,6 +193,59 @@ 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 _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, one partition per chunk, for a native in-memory
+    table.
+
+    Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
+    lazily loaded file), is larger than ``limit``, or has no rows; it is
+    read lazily instead.
+    """
+    # `_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
+    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
+    resolved = resolve_chunks(ds, chunks)
+    return [
+        list(iter_record_batches(ds.isel(block), schema, batch_size))
+        for block in _block_slices_from_resolved(ds, resolved)
+    ]
+
+
 def read_xarray_table(
     ds: xr.Dataset,
     chunks: Chunks = None,
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 9bdf1eb7..cf05c0ee 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,3 +1,4 @@
+import pyarrow as pa
 import xarray as xr
 from datafusion import SessionContext
 from datafusion.catalog import Schema
@@ -12,7 +13,7 @@
     shared_coord_arrays,
 )
 from .ds import XarrayDataFrame
-from .reader import read_xarray_table
+from .reader import _in_memory_partitions, read_xarray_table
 
 _proj: ModuleType | None
 try:  # pyproj is an optional dependency (`pip install xarray-sql[geo]`).
@@ -74,6 +75,14 @@ def from_dataset(
             )
             ctx.sql('SELECT * FROM era5.surface')
 
+        A Dataset is read lazily, chunk by chunk, as queries run, unless it
+        is already in memory (NumPy arrays, not Dask or a lazily opened
+        file) and at most 64 MiB as a table. Such a Dataset is converted
+        to Arrow once, here, and registered as a native DataFusion
+        in-memory table, which is faster to scan many times, as a query
+        does when it reads a table in several places. It is held twice,
+        as arrays and as Arrow.
+
         For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
         julian), a ``cftime()`` scalar UDF is automatically registered so
         you can write ergonomic SQL filters::
@@ -158,13 +167,28 @@ def _from_dataset(
         register = (
             self.register_table if schema is None else schema.register_table
         )
-        register(
-            table_name,
-            read_xarray_table(input_table, chunks, coord_arrays=coord_arrays),
-        )
+        partitions = _in_memory_partitions(input_table, chunks)
+        if partitions is None:
+            table = read_xarray_table(
+                input_table, chunks, coord_arrays=coord_arrays
+            )
+        else:
+            table = self._in_memory_table(partitions)
+        register(table_name, table)
         self._maybe_register_cftime_udf(input_table)
         return self
 
+    def _in_memory_table(self, partitions: list[list[pa.RecordBatch]]):
+        """A native DataFusion in-memory table of ``partitions``."""
+        before = set(self.catalog().schema().names())
+        view = self.create_dataframe(partitions).into_view()
+        # create_dataframe registers its in-memory table under a generated
+        # name. The view holds the table itself, not the name, so the name
+        # can go.
+        for name in set(self.catalog().schema().names()) - before:
+            self.deregister_table(name)
+        return view
+
     def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None:
         """Auto-register a cftime() UDF for non-Gregorian cftime coordinates."""
         for coord_name in ds.dims:

From 1bde0be391587bf881b16c07cff3d289393d787f Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Wed, 21 Oct 2026 00:03:59 -0700
Subject: [PATCH 2/6] 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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_sql.py    | 14 ++++++++++++++
 xarray_sql/reader.py | 35 ++++++++++++++++++++++++++++++-----
 xarray_sql/sql.py    | 12 +++++++-----
 3 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/tests/test_sql.py b/tests/test_sql.py
index cffa47d0..7473139b 100644
--- a/tests/test_sql.py
+++ b/tests/test_sql.py
@@ -647,3 +647,17 @@ def test_large_in_memory_dataset_is_read_lazily():
     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
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 2b27776b..646e5509 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import os
 from collections.abc import Callable, Iterator
 from typing import TYPE_CHECKING
 
@@ -223,8 +224,16 @@ def _in_memory_partitions(
     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, one partition per chunk, for a native in-memory
-    table.
+    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 a data variable whose array is C-contiguous isn't copied,
+    whatever the chunks: its Arrow column shares the array's memory, as
+    does DataFusion's table, through Arrow's C data interface. A NaN adds a validity bitmap beside it, one
+    bit per value. What the table adds is its dimension columns, one value
+    per row for each dimension, which long-form rows need. An array that
+    isn't C-contiguous (Fortran order) is copied, once.
 
     Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
     lazily loaded file), is larger than ``limit``, or has no rows; it is
@@ -239,10 +248,26 @@ def _in_memory_partitions(
     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
-    resolved = resolve_chunks(ds, chunks)
+    # 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))
     return [
-        list(iter_record_batches(ds.isel(block), schema, batch_size))
-        for block in _block_slices_from_resolved(ds, resolved)
+        batches[i : i + per_partition]
+        for i in range(0, len(batches), per_partition)
     ]
 
 
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index cf05c0ee..be674e1e 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -77,11 +77,13 @@ def from_dataset(
 
         A Dataset is read lazily, chunk by chunk, as queries run, unless it
         is already in memory (NumPy arrays, not Dask or a lazily opened
-        file) and at most 64 MiB as a table. Such a Dataset is converted
-        to Arrow once, here, and registered as a native DataFusion
-        in-memory table, which is faster to scan many times, as a query
-        does when it reads a table in several places. It is held twice,
-        as arrays and as Arrow.
+        file) and at most 64 MiB as a table. Such a Dataset is registered
+        as a native DataFusion in-memory table, which is faster to scan
+        many times, as a query does when it reads a table in several
+        places. The table shares its data variables' memory with the
+        Dataset rather than copying it, so, as with a lazy table, changing
+        the Dataset's values in place changes what queries return. It adds
+        a column per dimension, one value per row.
 
         For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
         julian), a ``cftime()`` scalar UDF is automatically registered so

From fd5cf9c5d075de762652d6e2a35e95923ad1ac3a Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Wed, 21 Oct 2026 02:14:02 -0700
Subject: [PATCH 3/6] Register the in-memory table itself, not a view of it

A view over the in-memory table made a query's plan scan the table's
generated placeholder name instead of the registered one. ddx finds the
table it differentiates by that name, so gradients with respect to a
small in-memory Dataset failed with UnknownColumn. Register the
in-memory table under a placeholder, take the table and drop the
placeholder, so the plan scans the caller's name. The table still
shares the Dataset's memory.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_sql.py |  7 +++++--
 xarray_sql/sql.py | 20 ++++++++++++--------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/tests/test_sql.py b/tests/test_sql.py
index 7473139b..cd7f74fc 100644
--- a/tests/test_sql.py
+++ b/tests/test_sql.py
@@ -592,8 +592,9 @@ def _scans(ctx, table):
 
 
 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, and
-    queries over it give what a lazily read Dataset gives."""
+    """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},
@@ -604,6 +605,8 @@ def test_small_in_memory_dataset_is_a_native_table_with_the_same_rows():
 
     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(),
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index be674e1e..b5f3a9c8 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -1,3 +1,5 @@
+import uuid
+
 import pyarrow as pa
 import xarray as xr
 from datafusion import SessionContext
@@ -182,14 +184,16 @@ def _from_dataset(
 
     def _in_memory_table(self, partitions: list[list[pa.RecordBatch]]):
         """A native DataFusion in-memory table of ``partitions``."""
-        before = set(self.catalog().schema().names())
-        view = self.create_dataframe(partitions).into_view()
-        # create_dataframe registers its in-memory table under a generated
-        # name. The view holds the table itself, not the name, so the name
-        # can go.
-        for name in set(self.catalog().schema().names()) - before:
-            self.deregister_table(name)
-        return view
+        # The in-memory table is reachable from Python only through a
+        # registered name, so register it under a placeholder, take the
+        # table, and drop the placeholder. The table, not a view of it, is
+        # what the caller registers, so a query's plan scans the caller's
+        # name, which ddx needs to find the table it differentiates.
+        placeholder = f"_xarray_sql_{uuid.uuid4().hex}"
+        self.register_record_batches(placeholder, partitions)
+        table = self.catalog().schema().table(placeholder)
+        self.deregister_table(placeholder)
+        return table
 
     def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None:
         """Auto-register a cftime() UDF for non-Gregorian cftime coordinates."""

From c3d8224c2f48e8dcd988e26a5eecd78610ae8469 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 09:46:10 -0700
Subject: [PATCH 4/6] Say that an in-place write must keep a Dataset's NaNs
 where they were

An in-memory table records which values are NaN (read as NULL) when it
is registered, so a NaN that an in-place write moves reads wrongly.
See #284.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 xarray_sql/reader.py | 10 ++++++----
 xarray_sql/sql.py    |  8 ++++++--
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 646e5509..2532855f 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -230,10 +230,12 @@ def _in_memory_partitions(
     them, as many partitions as ``chunks`` makes chunks, up to the number
     of CPUs. So a data variable whose array is C-contiguous isn't copied,
     whatever the chunks: its Arrow column shares the array's memory, as
-    does DataFusion's table, through Arrow's C data interface. A NaN adds a validity bitmap beside it, one
-    bit per value. What the table adds is its dimension columns, one value
-    per row for each dimension, which long-form rows need. An array that
-    isn't C-contiguous (Fortran order) is copied, once.
+    does DataFusion's table, through Arrow's C data interface. 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. An array that isn't C-contiguous (Fortran order)
+    is copied, once.
 
     Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
     lazily loaded file), is larger than ``limit``, or has no rows; it is
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index b5f3a9c8..46e06da5 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -84,8 +84,12 @@ def from_dataset(
         many times, as a query does when it reads a table in several
         places. The table shares its data variables' memory with the
         Dataset rather than copying it, so, as with a lazy table, changing
-        the Dataset's values in place changes what queries return. It adds
-        a column per dimension, one value per row.
+        the Dataset's values in place changes what queries return. The one
+        exception is NaN, which a query reads as NULL: the table records
+        which values are NaN when it is registered. So an in-place write
+        must keep the NaNs where they were, or the Dataset must be
+        registered again. The table adds a column per dimension, one value
+        per row.
 
         For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
         julian), a ``cftime()`` scalar UDF is automatically registered so

From 87752c5fb0a18cd5a713c7493110be5176f3c1a9 Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 11:44:08 -0700
Subject: [PATCH 5/6] Address review: batched partitions, a shorter docstring,
 two fallbacks

- Cut partitions with itertools.batched, shimmed for Python 3.10 and
  3.11.
- Say in from_dataset's docstring, in four lines, when a Dataset is an
  in-memory table and what an in-place write must keep.
- Check the Arrow batches' size against the 64 MiB limit, since the
  estimate undercounts booleans and long strings.
- Read a Dataset with a coordinate on a dimension no data variable has
  lazily, as main does; the in-memory path raised on it. Every query on
  such a table fails on main too (#286).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_sql.py    | 23 +++++++++++++++++++++++
 xarray_sql/reader.py | 29 ++++++++++++++++++++++++-----
 xarray_sql/sql.py    | 19 ++++++-------------
 3 files changed, 53 insertions(+), 18 deletions(-)

diff --git a/tests/test_sql.py b/tests/test_sql.py
index cd7f74fc..e1a6bab7 100644
--- a/tests/test_sql.py
+++ b/tests/test_sql.py
@@ -664,3 +664,26 @@ def test_in_memory_table_shares_the_dataset_s_memory(chunks):
     ds["w"].values[...] = 7.0
 
     assert ctx.sql("SELECT w FROM t").to_pandas()["w"].tolist() == [7.0] * 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()
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 2532855f..59ea1844 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -10,8 +10,10 @@
 
 from __future__ import annotations
 
+import itertools
 import os
-from collections.abc import Callable, Iterator
+import sys
+from collections.abc import Callable, Iterable, Iterator
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -35,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.
@@ -245,6 +258,11 @@ def _in_memory_partitions(
     # 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
+    # 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]))
@@ -267,10 +285,11 @@ def _in_memory_partitions(
     per_partition = -(-num_rows // (count * batch_size))
     size = -(-num_rows // (count * per_partition))
     batches = list(iter_record_batches(ds, schema, size))
-    return [
-        batches[i : i + per_partition]
-        for i in range(0, len(batches), per_partition)
-    ]
+    # 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(
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index 46e06da5..ab8ac6ed 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -77,19 +77,12 @@ def from_dataset(
             )
             ctx.sql('SELECT * FROM era5.surface')
 
-        A Dataset is read lazily, chunk by chunk, as queries run, unless it
-        is already in memory (NumPy arrays, not Dask or a lazily opened
-        file) and at most 64 MiB as a table. Such a Dataset is registered
-        as a native DataFusion in-memory table, which is faster to scan
-        many times, as a query does when it reads a table in several
-        places. The table shares its data variables' memory with the
-        Dataset rather than copying it, so, as with a lazy table, changing
-        the Dataset's values in place changes what queries return. The one
-        exception is NaN, which a query reads as NULL: the table records
-        which values are NaN when it is registered. So an in-place write
-        must keep the NaNs where they were, or the Dataset must be
-        registered again. The table adds a column per dimension, one value
-        per row.
+        A Dataset is read lazily, chunk by chunk, as queries run. One
+        already in memory (NumPy, not Dask or a lazily opened file) and at
+        most 64 MiB as a table is instead a DataFusion in-memory table,
+        which is faster to scan many times. It shares the Dataset's memory,
+        so changing values in place changes what queries return, as long
+        as no value becomes or stops being NaN (which reads as NULL).
 
         For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
         julian), a ``cftime()`` scalar UDF is automatically registered so

From bd8650de15ec97c0dcc0abb3c0aed5aaee8f203b Mon Sep 17 00:00:00 2001
From: Alex Merose <al@merose.com>
Date: Sun, 25 Oct 2026 13:56:28 -0700
Subject: [PATCH 6/6] 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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
---
 tests/test_sql.py    | 25 +++++++++++++++++++++++
 xarray_sql/reader.py | 47 +++++++++++++++++++++++++++++++++-----------
 xarray_sql/sql.py    |  6 +++---
 3 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/tests/test_sql.py b/tests/test_sql.py
index e1a6bab7..f83af95d 100644
--- a/tests/test_sql.py
+++ b/tests/test_sql.py
@@ -666,6 +666,31 @@ def test_in_memory_table_shares_the_dataset_s_memory(chunks):
     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."""
diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py
index 59ea1844..8ec88bcf 100644
--- a/xarray_sql/reader.py
+++ b/xarray_sql/reader.py
@@ -222,6 +222,23 @@ def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
     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,
@@ -241,23 +258,29 @@ def _in_memory_partitions(
 
     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 a data variable whose array is C-contiguous isn't copied,
-    whatever the chunks: its Arrow column shares the array's memory, as
-    does DataFusion's table, through Arrow's C data interface. 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. An array that isn't C-contiguous (Fortran order)
-    is copied, once.
-
-    Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
-    lazily loaded file), is larger than ``limit``, or has no rows; it is
-    read lazily instead.
+    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())
diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py
index ab8ac6ed..229dbc5a 100644
--- a/xarray_sql/sql.py
+++ b/xarray_sql/sql.py
@@ -80,9 +80,9 @@ def from_dataset(
         A Dataset is read lazily, chunk by chunk, as queries run. One
         already in memory (NumPy, not Dask or a lazily opened file) and at
         most 64 MiB as a table is instead a DataFusion in-memory table,
-        which is faster to scan many times. It shares the Dataset's memory,
-        so changing values in place changes what queries return, as long
-        as no value becomes or stops being NaN (which reads as NULL).
+        which is faster to scan many times. Either way, changing values in
+        place changes what queries return, as long as no value becomes or
+        stops being NaN (which reads as NULL) in an in-memory table.
 
         For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
         julian), a ``cftime()`` scalar UDF is automatically registered so