Tracking: ML in SQL as fast as JAX on CPU

Abierto
  • ¿Algo no está bien?
#285 25 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
38/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, sql
Área
backend, machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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

Lenguaje dominante
Python
Estrellas
142
Forks
19
Merge medio
3 d 47 min
PR fusionados (30 d)
8

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de xqlsystems/xarray-sql

Todos los issues de xqlsystems/xarray-sql

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.