Tracking: ML in SQL as fast as JAX on CPU

Ouverte
  • Quelque chose ne va pas ?
#285 25 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 3 jours

Personne n'a encore pris cette issue.

Évaluation

Difficulté
5/5
Temps estimé
Plus d'une semaine
Accessibilité débutants
38/100
Type d'issue
Fonctionnalité
Clarté
Plutôt claire
Activité
Active
Stack technique
python, sql
Domaine
backend, machine-learning

Piste de recherche

Start with benchmarks/ml/CONTRIBUTING.md and run the two benchmark commands in the issue to collect CPU and timing results. For the DuckDB join-order workstream, create the minimal repro the issue calls for; for a new demo, follow the contribution guide. Done means posting reproducible measurements or adding a model demo with benchmark results.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

ml

🌀 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.

JAX (CPU) DataFusion DuckDB
attention forward, 512 tokens 0.92 ms 0.28 s (window) / 0.37 s (GROUP BY) 0.48 s / 0.27 s
attention gradient, 512 tokens 2.4 ms 0.94 s 2.4 s (GROUP BY)
GPT loss, 1 layer, batch 4 0.15 ms 3.3 s 0.067 s
GPT SGD step 0.37 ms 0.69 s refused (xqlsystems/ddx#195)
GPT loss, 2 layers can't plan (runs out of 8 GB) 0.12 s

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 (#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 (xqlsystems/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. #267. DataFusion 56 (due Oct/Nov 2026, apache/datafusion#24461) adds WITH … AS MATERIALIZED (apache/datafusion#25580), and 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: #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: xqlsystems/ddx#202, the same for ddx's program steps.
  3. ddx's per-step overhead. xqlsystems/ddx#186 (fewer, larger steps), xqlsystems/ddx#184, xqlsystems/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 (#273). Small in-memory Datasets avoid it (#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 (xqlsystems/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).

Langage dominant
Python
Étoiles
142
Forks
19
Merge moyen
3 j 47 min
PR mergées (30 j)
8

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de xqlsystems/xarray-sql

Toutes les issues de xqlsystems/xarray-sql

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.