Tracking: ML in SQL as fast as JAX on CPU #285

🌀 Claude (xarray-sql ML demos)

This issue coordinates the work to make machine learning in SQL fast. The demos in benchmarks/ml (#269) write models as SQL over xarray-sql's tables, train them with ddx, which differentiates SQL queries, and check every result against JAX, Google's library for differentiable array programs. The goal is to match JAX on CPU with a SQL engine. Anyone can pick up a workstream below. Comment here to claim one, or to post a measurement.

Where we are

Intel i7-8700 (6 cores), DataFusion 55.0.0, DuckDB 1.5.6, ddx at its development branch after ddxdb 0.4.1, measured on 2026-10-24 and 2026-10-25. JAX is compiled with jax.jit.

The two demos are slow for different reasons:

• The GPT spends its time planning, not running. DataFusion copies a CTE into the plan at every place that reads it, so the 1-layer loss is 13,756 operators, and a second layer multiplies that by about 100 (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267). Planning the 1-layer loss takes 2.1 s, and the same physical plan runs in 2.4 ms. A ddx training step is 76 small queries, each planned in 3.3 ms and run in 0.65 ms (prepare: plan each program step once, and run its physical plan on every training step ddx#202).
• Attention spends its time running joins. It stores all N×N scores, and its plan takes 3–4 ms of a 0.1–0.4 s query. Running faster is the job of einfold, a query optimizer that rewrites SQL plans into array algorithms; it would find FlashAttention's tiled plan, for one.

Workstreams

1. Plan each shared CTE once. DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267. DataFusion 56 (due Oct/Nov 2026, Release DataFusion 56.0.0 (Oct/Nov 2026) apache/datafusion#24461) adds WITH … AS MATERIALIZED (feat: support WITH ... AS [NOT] MATERIALIZED CTEs apache/datafusion#25580), and feat: decide automatically when to materialize a CTE apache/datafusion#26276 would materialize CTEs automatically. Once 56 is out: enable it in Context, check it with xarray-sql's tables, and train the GPT at 2+ layers. ddx must read the new plan nodes too.
2. Plan once, run many times. Like jax.jit, which compiles once.
   • xarray-sql: Context.prepare is on the prepared-queries branch. It plans a query once and runs it on new values. It needs a reset of the plan's state between runs that datafusion-python doesn't expose: datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream #283, a draft to send upstream. Until then it decodes a fresh copy of the plan, which works only for plans shallow enough to decode.
   • ddx: prepare: plan each program step once, and run its physical plan on every training step ddx#202, the same for ddx's program steps.
3. ddx's per-step overhead. program: choose step boundaries by what they save (fewer, larger steps) ddx#186 (fewer, larger steps), rewrite: a cached program still optimizes each shared CTE's placeholder (15% of a cached GPT step) ddx#184, run: run a program's independent steps concurrently ddx#157 (run independent steps concurrently).
4. The FFI cost per scan. A lazy table costs about 1.5 ms per scan through datafusion-ffi (datafusion-ffi rebuilds the session and task context, with every function, on each scan and execute (~1.5 ms per scan) #273). Small in-memory Datasets avoid it (Register small in-memory Datasets as native in-memory tables #275); large ones still pay it.
5. DuckDB.
   • ddx's DuckDB backend can't yet plan a CTE that DuckDB computes once for several reads (DuckDB: a loss that reads a view with a multiply-read CTE fails (Not implemented Error: CTE) ddx#195).
   • DuckDB orders joins about 3× worse over an Arrow scan, which is how xql.register gives it a table, than over a native table: 0.64 s against 0.20 s for one gradient step. Needs an issue with a minimal repro.
6. Faster execution: einfold. einfold would run as a DataFusion optimizer rule, reading the dense-grid layout that xarray-sql knows about each table. That contract still needs writing down.
7. Bigger models. 02_nanogpt.py takes --layers, --embed, --heads and --context. nanoGPT's Shakespeare model is 6, 384, 6 and 256. Every workstream above moves this one.

Good places to start

• Run the demos on your machine, and post the tables here (see below). Other CPUs and OSes are welcome.
• Peak memory per host, from a fresh process per measurement.
• The DuckDB join-order repro in workstream 5.
• A new model as a demo. See benchmarks/ml/CONTRIBUTING.md. A new demo is also a new benchmark.

How to measure

uv run benchmarks/ml/01_attention.py --sizes 64 256 512 --reps 3 --csv attention.csv
uv run benchmarks/ml/02_nanogpt.py --steps 5 --reps 3 --csv nanogpt.csv

With a post, include:

• the CPU;
• python -c "import datafusion, duckdb, ddxdb; print(datafusion.__version__, duckdb.__version__, ddxdb.__version__)";
• the ddx commit, if you built ddx from a checkout.

To see whether a query's time goes to planning or running, time ctx.sql(q).execution_plan() (planning) against ctx.sql(q).to_arrow_table() (planning and running).