Each scan of a registered Dataset rebuilds its batches in Python: 73% of the GPT demo's forward pass #272

🌀 Claude (ML demos)

Summary

Scanning a table that xarray-sql registered from a Dataset costs about 1.5–2.3 ms more than scanning the same rows from an in-memory DataFusion table. Each scan calls the partition's Python stream factory under the GIL (PyArrowStreamPartition::execute in src/lib.rs), which builds the record batches from the Dataset again. That doesn't matter for one scan of a large dataset. But a query that reads small tables many times pays it on every read: DataFusion inlines a CTE or a view at each reference (#267), so the plan scans the same tables thousands of times.

For the GPT in benchmarks/ml/02_nanogpt.py (#269), registering the same eleven parameter tables and the batch as in-memory tables makes:

So 6.7 s of the 9.2 s forward pass, 73%, is the per-scan cost, and the loss's plan size (#267) multiplies it.

Reproduction

import time
import numpy as np
import xarray as xr
import xarray_sql as xql

ctx = xql.Context()
ds = xr.Dataset({"w": (("tok", "e"), np.random.rand(65, 16))})
ctx.from_dataset("wte", ds)
batches = ctx.sql("SELECT * FROM wte").to_arrow_table().to_batches()
ctx.register_record_batches("wte_mem", [batches])

for name in ["wte", "wte_mem"]:
    ctx.sql(f"SELECT * FROM {name}").to_arrow_table()
    start = time.perf_counter()
    for _ in range(100):
        ctx.sql(f"SELECT * FROM {name}").to_arrow_table()
    print(name, (time.perf_counter() - start) * 10, "ms per scan")

On an Intel i7-8700 with DataFusion 54.0.0, SELECT * FROM wte takes 1.69 ms through xarray-sql against 0.25 ms in memory. A join and GROUP BY of two such tables takes 5.6 ms against 1.8 ms.

Possible fixes

• For a Dataset whose variables are already in memory (NumPy-backed, not Dask or a lazy backend), build each partition's batches once, at registration or on the first scan, and keep them. Scans then cost what a MemTable does. Memory is the trade-off: a cached Dataset is held twice, as arrays and as Arrow.
• Or cache a partition's batches by projection for the life of the table, so repeated scans in one plan share them.
• Either way, the cache must be invalidated when a table is registered again under the same name.