Tracking: ML in SQL as fast as JAX on CPU

Đang mở
  • Có gì không ổn?
#285 25 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 3 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
5/5
Thời gian dự kiến
Hơn một tuần
Mức phù hợp với người mới
38/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, sql
Lĩnh vực
backend, machine-learning

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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

Ngôn ngữ chính
Python
Star
142
Fork
19
Merge trung bình
3 ngày 47 phút
Pull request đã merge (30 ngày)
8

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của xqlsystems/xarray-sql

Tất cả issue của xqlsystems/xarray-sql

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.