Repository navigation
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
SQL autodifferentiation integration xarray_sql/ad.py, xarray_sql/sql.py, tests/test_ad.py, pyproject.toml |
Adds an optional ddx interface for grad, vjp, and jvp queries. Context.sql routes matching queries through the extension. Tests compare gradient operations with NumPy results. |
Benchmark hosts and reporting benchmarks/ml/_hosts.py, benchmarks/ml/_bench.py, tests/test_ml_demos.py, .github/workflows/ci.yml, benchmarks/ml/CONTRIBUTING.md, benchmarks/ml/README.md, pyproject.toml |
Adds DataFusion and DuckDB hosts, shared timing and reporting utilities, demo subprocess tests, a separate CI job, and dependency configuration. Documentation describes the demo workflow and host behavior. |
Causal attention benchmark benchmarks/ml/01_attention.py, benchmarks/ml/README.md |
Adds two SQL spellings of causal attention, JAX forward and gradient comparisons, an infinity-based masking check, and a Nadaraya–Watson comparison with NumPy. |
nanoGPT benchmark benchmarks/ml/02_nanogpt.py, benchmarks/ml/README.md |
Adds a SQL character-level GPT and training updates using SGD or AdamW. The benchmark compares losses and parameter updates with JAX. |
Priority:
Estimated code review effort: 4 (Complex) | ~60 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant AttentionDemo
participant Host
participant SQLEngine
participant JAX
participant Report
AttentionDemo->>JAX: compute reference results
AttentionDemo->>Host: load inputs and execute SQL
Host->>SQLEngine: run SQL and return results
AttentionDemo->>Report: record timings and compare results
Possibly related issues
- Tracking: ML in SQL as fast as JAX on CPU #285: Tracks performance work for the ML demos added here.
- xarray-sql's facts, and an einfold host for xarray-sql's ML demos einfold#39: Requests an einfold contract and host for the two demos; this pull request adds the demos but not that contract.
- Python path: step plans' column names grow with depth and slow DataFusion's planning ddx#211: Reports computed-column name growth in ddx plans during ML training steps added here.
- program: choose step boundaries by what they save (fewer, larger steps) ddx#186: Proposes larger ddx steps to reduce the nanoGPT training cost of workflows added here.
- prepare: plan each program step once, and run its physical plan on every training step ddx#202: Proposes reusing ddx program-step plans in the nanoGPT training loop added here.
- Exit: xarray-sql's ML demos differentiate on DuckDB ddx#150: Defines DuckDB gradient support and a ddxdb release requirement for the demos and
mlextra added here. - DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267: Documents DataFusion CTE expansion in the ML queries added here.
- ad.prepare: reuse each step's physical plan with reset_plan_states instead of planning it again ddx#213: Proposes physical-plan reuse for ddx training steps in the nanoGPT demo.
- datafusion-ffi rebuilds the session and task context, with every function, on each scan and execute (~1.5 ms per scan) #273: Measures DataFusion FFI scan overhead in the demo data path.
- A view keeps reading a table after it is replaced: DataFusion binds views at creation #271: Describes stale views when the GPT demo replaces parameter tables between training steps.
- M2: the scaled accumulator as a monoid action, and fold pipelines as fusion (ideas from Haskell) einfold#50: Proposes fold fusion for the causal-attention workload added here.
Possibly related PRs
- xqlsystems/xarray-sql#275: Changes
Context._from_dataset, which the DataFusion benchmark host uses to register datasets. - xqlsystems/xarray-sql#290: Addresses repeated CTE and view work for the NanoGPT loss query used by this pull request.
- xqlsystems/xarray-sql#287: Defines prepared execution behavior for differentiation queries that this pull request routes through
Context.sql.
Merge Risk: 
1361a
-
The PR adds useful autodiff and ML demos, but parameterized SQL can regress, the benchmark documentation remains incomplete, and default development setup still fails on Intel macOS. These issues should be addressed before merging.
Security architecture risk: 
1361a
-
The integration is optional and uses existing local database sessions. The main unresolved risk is whether statements routed through differentiation preserve table ownership and prepared-query invalidation. No security vulnerability was established, but incomplete dependency and failure-recovery coverage prevents a minimal-risk assessment.
Retained concerns
- Low · reliability · inferred: The differentiation route returns before table replacement/drop invalidation. If ddxdb executes DDL selected by the textual detector, including a match inside a comment, Dataset ownership metadata and prepared-query invalidation counters can remain stale. DDL acceptance is unverified, so this is a conditional lifecycle-contract concern rather than a demonstrated security vulnerability.
Security review details
-
Security Blast Radius
- inferred — The demonstrated exposure is the Python process and database sessions using the optional integration, including their registered datasets. Installing ddxdb affects all Context instances in that process. The inspected consumers do not establish a tenant-facing service or cross-tenant attack path; external dependency internals remain outside verified coverage.
Trust Boundaries and Controls
- observed — Caller-supplied SQL can select a new third-party planning path through textual matching. The adapter passes the same session and statement to ddxdb; it is not an authorization boundary. Extra query arguments are rejected on this path rather than silently discarded.
Resilience and Maintainability Implications
- observed — The existing table-lifecycle control removes Dataset metadata and increments deregistration counters after replacement or drop. Prepared queries check those counters before writing or executing. Unmatched SQL and explicit deregistration retain these controls; the new AD early return does not pass through them.
Hardening Proposals
- proposed — Keep successful table replacement/drop finalization common to both planning routes, or explicitly reject mutating statements on the differentiation route. Specify session reuse, temporary-state cleanup and concurrency guarantees for sql_all after failure before relying on it in shared long-lived contexts.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title accurately and concisely summarizes the main changes: adding the ML dependency extra, SQL gradients, and two ML demos. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
Looking good so far! :)
| * Dense tensor operations are slower in SQL than Xarray. It seems like the vanilla HashJoin implementation in standard DB engines are not as performant as array ops (they don't take advantage of contiguous array layouts), though this gives us some idea for how it could [one day be possible](https://github.com/xqlsystems/einfold). | ||
| * Sparse tensor operations are just as fast (or faster!) than Xarray today. Specifically, SQL is an excellent way to perform regridding and reprojecting between different spatial scales (see also: we include a `proj` udf via the `xarray-sql[geo]` extra). | ||
|
|
||
| ## Gradients in SQL |
There was a problem hiding this comment.
Let's omit this section for now. I want to publicize it when it is really ready.
There was a problem hiding this comment.
🌀 Claude: Removed in 21f99d6. The section's example still lives in benchmarks/ml/README.md and in xarray_sql.ad's docstring. I listed it under "Not done yet" in the PR description so it isn't forgotten.
| # [tool.uv.sources] | ||
| # xarray-sql = { path = "../../", editable = true } | ||
| # /// | ||
| """Demo 1: attention in SQL, the query einfold should turn into FlashAttention. |
There was a problem hiding this comment.
In general in all documentation, if you make a reference to something, please introduce it first.
Let's put this note of feedback in the Agents.md file so future agents don't make the same mistake.
There was a problem hiding this comment.
🌀 Claude: Done in 21f99d6. This docstring now says what einfold, FlashAttention (with the paper), JAX, ddx and a host are before using them. I went through every document in the PR for the same problem: the other two demos, benchmarks/ml/README.md, _hosts.py, _bench.py and xarray_sql/ad.py (Substrait, JAX), plus the GraphCast docs on #270 (GraphCast, ERA5, forcings, swish).
AGENTS.md now has the rule: introduce every project, paper, algorithm, dataset or term of art the first time a document mentions it, and link it where there is a link.
| # One statement per input, sharing one backward pass. | ||
| gradients = [ | ||
| f"WITH {ctes}, {LOSS} " | ||
| f"SELECT t, d, grad(loss, x) AS x FROM {name} ORDER BY t, d" | ||
| for name in "qkv" | ||
| ] |
There was a problem hiding this comment.
This is compact code, but I'd rather make it legible / instructive SQL code, even if it is a bit more verbose. If it's all the same perf anyway, I'd rather demo what attention looks like as a long SQL statement.
There was a problem hiding this comment.
🌀 Claude: Done in 21f99d6. Each spelling is now two complete SQL statements, written out in full: the forward pass, and its gradient. The gradient statement is the forward query plus a loss CTE and a different final SELECT, so the two can be read side by side.
While doing this, I found that ddx can take the gradient with respect to all three inputs in one query, from a join of the three tables:
SELECT q.t, q.d,
grad(loss, q.x) AS dq,
grad(loss, k.x) AS dk,
grad(loss, v.x) AS dv
FROM q
JOIN k ON q.t = k.t AND q.d = k.d
JOIN v ON q.t = v.t AND q.d = v.d
ORDER BY q.t, q.dAt 256 tokens it takes 0.59–0.67 s, against 0.61–0.62 s for three statements sharing one backward pass through sql_all, and it matches jax.grad to 3e-15. So the comprehension is gone, and so is sql_all in this demo. The cost is that the four CTEs appear twice per spelling.
| LAYERS = 1 # transformer blocks | ||
| EMBED = 16 # e: channels of the residual stream | ||
| HEADS = 2 # h: attention heads | ||
| HEAD = EMBED // HEADS # c: channels per head | ||
| HIDDEN = 4 * EMBED # f: channels of the MLP | ||
| CONTEXT = 16 # pos: characters per sequence | ||
| LEARNING_RATE = 0.5 |
There was a problem hiding this comment.
Just want to double check: is this the same param count (and hparams) as NanoGPT? Can you add a link in a comment above this so we can double check?
There was a problem hiding this comment.
In general, will you double check that our model matches the reference?
There was a problem hiding this comment.
🌀 Claude: Thanks for asking me to check. The hyperparameters are not nanoGPT's, on purpose. The model wasn't quite nanoGPT's either, which wasn't on purpose, and 21f99d6 fixes that.
Size. nanoGPT's Shakespeare model (config/train_shakespeare_char.py) has 6 layers, 6 heads, 384 channels and 256 characters of context: 10.65M parameters. This one has 1 layer, 2 heads, 16 channels and 16 characters: 4,416 parameters. Even at this size, the 1-layer loss is 12,794 operators and 9.4 s on DataFusion, so nanoGPT's size isn't reachable yet. Each constant now names its GPTConfig counterpart and that config's value, with links to model.py and the config, and the script prints its parameter count. The vocabulary matches: 65 characters, from sorted(set(text)), as in nanoGPT's prepare.py.
What I compared against model.py, and fixed:
- Layer-norm gains:
LayerNormhas a learnedweight, and the SQL had none. There are now tablesln_1,ln_2andln_f, with gains initialized to 1. Biases stay off, matchingtrain.py's defaultbias = False. - Activation: nanoGPT's MLP uses GELU, and mine used ReLU² (from modded-nanogpt). It's now GELU in its tanh approximation, as GPT-2 computed it. nanoGPT's
nn.GELU()is the exacterfform, which DataFusion doesn't have. The two differ by less than 5e-4. ReLU² can come back as a one-line diff in Part B, where einfold's zero-skipping needs it. - Initialization: nanoGPT scales
c_proj(woandw_projhere) to 0.02/√(2·n_layer). Mine used 0.02 for every weight; it now uses nanoGPT's scaling.
What already matched: tied wte and lm_head, learned position embeddings, pre-norm residual blocks, eps 1e-5, the 1/√head_size scale, the causal mask, and mean cross-entropy. c_attn holds the same parameters as wq, wk and wv, in one matrix.
Still different, now listed in the docstring:
- SGD at a fixed learning rate. nanoGPT uses AdamW with warmup, cosine decay, weight decay and gradient clipping.
- No dropout, so that a run can be checked against JAX. nanoGPT's Shakespeare config trains with dropout 0.2.
AdamW in SQL means keeping the optimizer's state in versioned tables, which einfold's Demo 2 calls out as an open question. I'd do it with Part B unless you want it here.
After the changes, five SGD steps still match JAX on every step: the loss to 2e-16, and all 11 parameter tables to 1e-15.
| * **GELU** is its tanh approximation, as GPT-2 computed it. nanoGPT's | ||
| ``nn.GELU()`` uses the exact form, through the error function ``erf``, | ||
| which DataFusion doesn't have. The two differ by less than 5e-4. | ||
| * **The optimizer** is SGD at a fixed learning rate. nanoGPT trains with |
There was a problem hiding this comment.
Can we devise a way to match NanoGPT in the optimizer? For example, could we use a table of learning rates to define a LR schedule?
There was a problem hiding this comment.
Maybe we should keep simple SDG and have the more complex thing behind a flag.
There was a problem hiding this comment.
🌀 Claude: Done in 5f92e69, along the lines you suggested in the follow-up: SGD stays the default, and --optimizer adamw trains as nanoGPT does. The learning rate is a table, schedule(step, lr), built from nanoGPT's get_lr (linear warmup, then cosine decay to min_lr). Each step's update joins it on s.step = {step}.
An AdamW step is two phases:
- The gradient of every parameter table,
SELECT …, grad(loss, w) AS g FROM wteand so on, throughsql_all, so all eleven share one backward pass. Each result is registered as the table{name}_grad. - One plain-SQL statement per parameter table, with no ddx in it:
grad_normandclip: the global norm of all eleven gradients, through aUNION ALLof the gradient tables, and the clip factorLEAST(1.0, 1.0 / (norm + 1e-6)), asclip_grad_norm_computes it;moments: Adam's two moments, which live as columnsmandvbesidewin each parameter table, so the optimizer's state is versioned with the weights;- the update: decoupled weight decay (0 for the layer-norm gains, as in nanoGPT's
configure_optimizers), bias correction froms.step + 1, and eps.
Settings: nanoGPT's Shakespeare model's: lr 1e-3 to 1e-4, β = (0.9, 0.99), weight decay 0.1, clip 1.0. One difference: the schedule is stretched to --steps (warmup is 2% of the steps, as 100 of 5,000 is), because nanoGPT's spans 5,000 steps.
Checks:
- The JAX reference step matches
torch.optim.AdamWafterclip_grad_norm_exactly (difference 0.0 over four steps, with clipping active). A control that changes torch's weight decay shows a difference of 5e-4, so the comparison isn't vacuous. - The SQL matches that reference to 6e-15 after each of five steps: the weights and both moments of all eleven tables.
Because the update is plain SQL, it would also run on DuckDB once ddx can take gradients there.
There was a problem hiding this comment.
🌀 Claude: Agreed, and that's how 5f92e69 does it: --optimizer sgd (the default) or --optimizer adamw. The SGD step is still one statement per table, w - 0.5 * grad(loss, w). Details in the reply above.
3baf981 to
fd5cf9c
Compare
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
ddxdb 0.4.3 differentiates a view whose CTE computes a window on DuckDB (xqlsystems/ddx#204), so the demos no longer expect DuckDB to refuse attention's window spelling or the GPT's training step. Record the GPT on DuckDB: its weights match JAX's, its loss takes 0.062 s against DataFusion's 2.7 s, and an SGD step takes 0.87 s against 0.41 s, most of it DuckDB planning ddx's 77 steps afresh and the adapter's reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
DataFusion plans a CTE again at every place that reads it, so a model whose CTEs read each other more than once has a plan that grows exponentially with its depth (#267). The 1-layer loss took 2.7 s at the demo's size and 108 s at nanoGPT's CPU size, and the 2-layer loss ran out of memory while planning. - The GPT's SQL builders return their CTEs as (name, sql) pairs, and every statement starts with the model's and the loss's CTEs, through _hosts.with_query. - Hosts gain query_with(ctes, select) and operators_with. DataFusion's runs each CTE on its own into an in-memory table that later CTEs read, under TODO(apache/datafusion#26237). It uses collect_partitioned and register_record_batches rather than DataFrame.cache(), which einfold can't yet plan a join of (xqlsystems/einfold#93). - The model is no longer the view logits. Creating a DataFusion view writes its CTEs out, so the 3-layer view didn't fit in 8 GB; ddx, given the CTEs inline, computes each shared one once. apache/datafusion#24780's streaming fan-out example doesn't fit these queries: by its own notes it can deadlock when both sides of a hash join read the shared plan, and doesn't nest. DataFusion 55 parses `AS MATERIALIZED` only in its PostgreSQL dialect, and ignores it. With ddxdb 0.4.4, at batch 4, on an Intel i7-8700: the 2-layer loss takes 0.07 s on DataFusion, and a 6-layer SGD step 2.4 s, against 0.30 s and 5.5 s on DuckDB; every step matches JAX. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
…reader ddxdb 0.4.5 releases ddx's fixes for attention's backward (xqlsystems/ddx#220, #225), drops intermediate tables once nothing reads them (#223), and halves DuckDB's per-call overhead (#224). The GPT demo's steps match JAX on it, on DataFusion and DuckDB. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
) 🪭 Claude (einfold) For the maintainer's request to integrate einfold with xarray-sql's demos. `python/` is a Python package, `einfold`, that installs einfold in any session of datafusion-python (Apache DataFusion's Python bindings), and so in xarray-sql's `Context` and ddxdb's: ```python from datafusion import SessionContext from einfold.datafusion import EinFold einfold = EinFold() ctx = SessionContext().with_extensions(einfold) # every later query is planned by einfold einfold.plans_with_einfold(), einfold.plans_with_pipelines() ``` The package is organized by engine: `einfold.datafusion` now, with `datafusion` an optional extra (`pip install "einfold[datafusion]"`), so that `einfold.duckdb` can sit beside it. The README is written for someone who first meets einfold on PyPI. It's spike S29's bridge (`docs/spikes/s29-python-planner`), made fit to keep: - **Detection includes pipelines.** The planner calls `DetectFolds::detect`, so demo 1's attention runs fused. - **The codecs no longer leak.** - A table is held only from its encoding, on the session's side, to its decoding on einfold's, within one planning call. The spike kept every table forever, and holding them weakly doesn't work: the codec is handed a fresh FFI wrapper per encode. - Plans are held the same way, from encoding to decoding. A plan that has run keeps what its operators hold, such as a hash join's table, so keeping plans any longer kept every query's tables: in xarray-sql's GPT demo, about 2 GB per query. - **`MAX` is relinked as well as `SUM`,** in aggregates and windows. Built-ins cross the FFI as foreign wrappers, and detection matches them by type, so the attention pipeline wasn't found until `MAX` was relinked too. - **The returned plan can be reset,** for recursive queries, by resetting the plan inside it. - **No debug switches.** The extension counts what it planned instead. - **Plans of any depth decode** (fixes #77). The session hands each logical plan over as protobuf, and the decoder, prost, stopped at 100 nested messages, which a chain of 15 joins exceeds. The decoder is this library's copy of prost, inside DataFusion's FFI, so the crate turns on prost's `no-recursion-limit` feature. xarray-sql's GPT demo (in xqlsystems/xarray-sql#269) now plans and trains with einfold installed, matching DataFusion's loss to `2e-16`. Deeper plans then overflowed the stack, from about 250 joins, where DataFusion alone plans about 1,350. datafusion-proto decodes a plan, on the thread that asks for it, by recursing without growing its stack (#95). So installing einfold raises the process's soft stack limit to 64 MB, as `ulimit -s` would: Linux grows the main thread's stack on demand up to that limit. A chain of 1,300 joins now plans. **From Python** (datafusion-python 55.0.0 from PyPI, i7-8700): | query | without einfold | with einfold | |---|---:|---:| | demo 1's causal attention, 1,024 tokens | 2.07 s | 17.5 ms | | demo 1's causal attention, 4,096 tokens | 109 s | 46 ms | | matrix product, `50,000 × 16 · 16 × 8` | 37.9 ms | 20.9 ms | | `SELECT 1` | 0.23 ms | 1.69 ms | The last row is the fixed cost of crossing the FFI, 1.5–3 ms per query. That's for datafusion-python and its FFI to fix (#58). **Packaging.** The crate is outside the Cargo workspace, like the benchmarks. It needs `unsafe` to read datafusion-python's FFI capsules, which the workspace forbids, and pyo3 and DataFusion's FFI crates. Each `unsafe` block carries a `SAFETY` comment, and `unsafe_op_in_unsafe_fn` is denied. A new CI job builds the wheel with maturin, runs clippy on the crate, and runs the tests against datafusion-python 55.0.0 from PyPI. **Tests** (`python/tests/test_einfold.py`): each query gives the same rows with einfold installed as without, to `1e-12`, and einfold runs what it should: - a matrix product (a fold), and demo 1's attention (a pipeline); - a query with no fold, left alone; - installing in one context and querying through the shared session; - a query run again, with nothing left in the codec after each run, and a table replaced between runs; - xarray-sql's `Context` with `from_dataset`, and ddxdb's `grad`, both skipped where the packages aren't installed; - plans that are chains of 15 and 60 joins (#77), and of 350 in a subprocess, past the default stack. Another test pins the limit that remains: 350 joins planned on a thread with an 8 MB stack still crash, until datafusion-proto's decoding grows its stack (#95). All 9 pass locally against xarray-sql's environment. CI has neither xarray-sql nor ddxdb, so it runs 7. **Next, in xarray-sql:** an `EinFold` host in `benchmarks/ml/_hosts.py`, a `DataFusion` host whose context is `xql.Context().with_extensions(EinFold())`. That's a PR there, once xarray-sql#269 merges. <!-- - [x] cargo fmt / clippy -D warnings (python crate) / reuse lint; pytest 9 passed locally --> ## After CodeRabbit's review - **Each foreign FFI struct's version is checked before it's cloned:** the session's query planner, and both codecs. - **The README has a Limits section.** Very deep plans can overflow a thread's stack outside Linux's main thread. A function registered as `sum` or `max` is replaced by DataFusion's, since nothing crossing the FFI tells them apart (#111). - The deep-plan test sets its core limit in the child, on Linux only, and is skipped elsewhere. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
🪭 Claude (einfold) With #77 fixed, the GPT in xarray-sql's ML demos plans with einfold installed (`benchmarks/ml/02_nanogpt.py`, in xqlsystems/xarray-sql#269: a small GPT trained in SQL, with gradients from ddx, an XQL Systems project that differentiates SQL queries). But it was much slower than DataFusion alone. At the demo's default size, its loss took 6.2 s to plan and 2.4 s to run with einfold, against 2.3 s and 0.35 s without. Most of that was scans. Every table a plan reads reaches einfold's Python extension as a foreign table, even an in-memory one, so each scan crossed DataFusion's FFI, its stable C interface between libraries. A foreign scan converts the session at planning, and the task context at each partition it runs. Each conversion copies the whole configuration as strings and wraps each of several hundred functions, both ways. The GPT's plan expands its views at each read, so it has about 1,500 scans. This change does two things: - **Every table is wrapped** (`Lean`), so that its scans get the session's configuration only, with no functions or catalogs, and remember their statistics, which DataFusion's join planner asks for at every join. A scan given filters to evaluate still gets the whole session, since filters can name functions. - **Small in-memory tables are copied once per plan** (`localize`). That means DataFusion's own data source, with exact statistics of at most 64 MiB, the bound xarray-sql uses for the datasets it keeps in memory. The copy shares the table's buffers, and every scan of that name in the plan reads it, so none crosses the FFI. xarray-sql's lazily loaded datasets also report exact statistics, but through their own scan, so they're left to stream. Results, from xarray-sql's environment, on an i7-8700: | GPT loss, default size | plan | run | |---|---:|---:| | DataFusion alone | 2.3 s | 0.35 s | | with einfold, before | 6.2 s | 2.4 s | | with einfold, wrapped scans | 4.0 s | 1.5 s | | with einfold, wrapped and copied | 3.9 s | 1.1 s | At a mid size (batch 8, 64 channels, context 32), the loss is now about even with DataFusion alone (7.8 s against 7.3 s), and a training step is faster (1.6–2.1 s against 2.5–3.1 s). Two costs remain, and I'll take them up next: - each fold's fallback is planned again with everything below it; - `EinFoldExec` reports no statistics, so DataFusion can't tell the tables above it are small, and repartitions them. Peak memory at that size is 5.9 GB with einfold, against 2.8 GB without. That's #79. I tried relinking every built-in scalar function, as einfold already relinks `SUM` and `MAX`, since a foreign function also crosses the FFI at each call. It made no measurable difference here, so it isn't in this change. **Unnamed tables.** A DataFrame's cached result (`DataFrame.cache()`) is a table without a name: every scan of one is called `?table?`. So a scan of an unnamed table gets a copy of its own, rather than sharing the copy of its name. Without that, two cached tables in one plan read the same copy: the plan failed to plan where their columns differed, and gave wrong answers where they matched (fixes #93). **Tests.** The Python suite passes (12 tests). The deep-plan tests read one table up to 60 times in a plan, so they check the copy against DataFusion's answers. `test_tells_unnamed_tables_apart` joins two cached DataFrames, one with different columns and one with matching ones; both fail without the fix. A new test changes an xarray-sql in-memory dataset's values in place between two queries, and checks that einfold sees the change: xarray-sql promises this, and a copy kept across plans would break it. ## After CodeRabbit's review - **Copies are shared by data, not by name.** A name can name two tables in one plan, as when a DataFrame made before the name was registered again is joined with one made after. Each scan's table is a fresh foreign wrapper, so nothing else identifies it. A foreign scan of an in-memory table hands over the table's own buffers. So `localize` keys each copy by the scan's name, schema and row count, and by where every column of its first batch lies. Only scans that match share a copy. Planning now opens each scan once to read those addresses, which the release candidates' timings will show. `test_tells_a_table_registered_again_apart` gave a wrong answer before. - **A scan with filters runs with the session's functions,** as it's planned with them. If copying the tables fails, the bridge plans the query as the session would. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…ython does (#83) 🪭 Claude (einfold) For #79 (peak memory as a goal). With einfold installed, a session that ran xarray-sql's GPT loss again and again kept growing. (The GPT is `benchmarks/ml/02_nanogpt.py`, in xqlsystems/xarray-sql#269.) Two causes: - **The codec kept executed plans.** A plan that has run keeps what its operators hold, such as a hash join's table, and the codec kept the last 1,024 plans. That cost about 2.2 GB per query. It's fixed on #75 ("Hold a plan from its encoding to its decoding"), since the code is there. - **The allocator** is this PR. This library allocated with the system allocator, glibc's, which keeps memory per thread, and einfold's kernels run on blocking threads. A session still grew by about 150 MB per query. With `MALLOC_ARENA_MAX=2` the growth stopped, so this is fragmentation, not retention. datafusion-python links mimalloc, and its sessions don't grow, so this library now allocates with mimalloc too. A buffer one library allocates is freed by its own release callback, through Arrow's C data interface, so the two allocators never free each other's memory. The GPT's loss at batch 8, 64 channels, context 32, run ten times in one session, from xarray-sql's environment: | | after the 1st | after the 5th | after the 10th | |---|---:|---:|---:| | plans kept | 2.9 GB | 9.9 GB (killed at 10 GB) | | | plans freed, glibc | 2.8 GB | 3.4 GB | | | plans freed, mimalloc | 3.3 GB | 3.1 GB | 3.1 GB | DataFusion alone: 2.7 GB after the 1st, 2.3 GB after the 5th. The whole demo at that size, three steps, now peaks at 3.6 GB with einfold against 2.7 GB without; before both fixes it reached 5.9 GB in two. With einfold, a training step takes 1.3 s against 2.5 s. The loss still peaks about 600 MB higher with einfold, which #79 is about. **Tests.** The Python suite passes (12 tests). #75's test of running queries again now checks that the codec holds no tables or plans after each query. ## After CodeRabbit's review - **mimalloc 2, not 3:** mimalloc 3 crashes on macOS arm64 when another copy of it runs in the process, as PyArrow 24's does (apache/datafusion-python#1607). #110's release workflow tests the macOS arm64 wheel. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
🪭 Claude (einfold) For #79's first ask: measure peak memory next to time, einfold off and on. `python/benchmarks/memory.py` runs each case in a process of its own, so each peak is that case's alone. The peak is Linux's VmHWM, the process's peak resident memory, reset just before the query, so it counts what the query adds above its loaded tables. Each case runs three times, and the fastest time and largest peak are reported. The cases: - #79's contraction `x(s, p, k) · w(k, e)`, at s = 16, 64 and 256; - demo 1's causal attention, at 1,024 and 4,096 tokens; - with `XARRAY_SQL` pointing at an xarray-sql checkout, the loss and the gradients of its GPT demo (`benchmarks/ml/02_nanogpt.py`, in xqlsystems/xarray-sql#269), at batch 8, 64 channels, context 32. From xarray-sql's environment, on this stack (#82, #83), on an i7-8700: | case | einfold off | einfold on | |---|---:|---:| | contraction-16 | 104 ms, 106 MB | 8 ms, 71 MB | | contraction-64 | 302 ms, 171 MB | 14 ms, 103 MB | | contraction-256 | 1,359 ms, 400 MB | 42 ms, 175 MB | | attention-1024 | 1,912 ms, 630 MB | 14 ms, 99 MB | | attention-4096 | 48,905 ms, 6,059 MB | 39 ms, 117 MB | | gpt-loss | 7,048 ms, 2,320 MB | 7,272 ms, 3,079 MB | | gpt-gradients | 2,510 ms, 1,184 MB | 1,240 ms, 587 MB | einfold takes less memory than DataFusion in every case but the GPT's loss. That's the case #79's second ask is about: `EinFoldExec` reads its inputs whole, where the host's joins stream their larger side. It's a script to run by hand, not a test: it needs datafusion-python, and the GPT cases need xarray-sql and ddxdb (ddx's Python package; ddx, an XQL Systems project, differentiates SQL queries). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
einfold, a query optimizer for DataFusion that runs tensor computations in SQL with array algorithms, released its first candidate. The demos gain an einfold host: xarray-sql's Context with einfold as its query planner. Demo 1 now checks that einfold runs both spellings of attention, and the Nadaraya-Watson control, as one fused pipeline: FlashAttention's tiled plan, found from the SQL. The README gives einfold's numbers. einfold has no Windows wheel yet (xqlsystems/einfold#118), so the dev group skips it on Windows rather than build it from source. ddxdb 0.4.6 differentiates through UNION ALL and computes a step read once inside its reader, which takes the GPT's program from 77 steps to 64. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
Context.sql sent a query to ddx when "grad(" appeared anywhere in its
text, so a parameterized query with the word in a string literal raised
TypeError. calls_ad now ignores string literals and comments.
jaxlib publishes no wheels for Intel macOS, so the dev group installs
JAX everywhere else, and the demos' test skips without it. The demos'
README says what the Nadaraya-Watson estimator is.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
The nanogpt script is really confusing. I suspect it changed bc of the caching CTE issue in datafusion. I think we should either leave a TODO or make it legible now and find long term fix later. I can make an announcement for the first demo, and will showcase the second demo later when the long term problems have been addressed and when we can do it right.
|
|
||
| def layer_norm( | ||
| x: str, out: str, gain: str, layer: int | None = None | ||
| ) -> list[tuple[str, str]]: |
There was a problem hiding this comment.
I find returning lists of tuples really confusing. Why do we do this? Can we just have pure functions? Please explain what's going on.
There was a problem hiding this comment.
🌀 Claude: You're right that it was confusing, and the workaround caused it. SQL names an intermediate result that's read more than once with a CTE, and a layer norm reads its input three times, so the model has to be one query with a long WITH clause. Before the workaround, each builder returned its CTEs as text. The workaround runs each CTE on its own on DataFusion, so the builders started returning them as (name, sql) pairs, and every function had to be told what to name its output.
714ba63 replaces that with values. Each function takes relations and returns one, a single SELECT that reads its inputs by name in an f-string, so block now reads as nanoGPT's Block.forward:
x = residual(x, attention(layer_norm(x, "ln_1", layer), layer))
return residual(x, mlp(layer_norm(x, "ln_2", layer), layer))_relation.py (about 60 lines) turns the loss into the CTE list, naming each relation and never with a table's name. The functions are pure, no names are threaded through, and only loss_ctes() sees the pairs. The SQL is the same 27 CTEs as before, and the demo matches JAX on all three hosts, with SGD at 1 layer and with AdamW at 2 layers.
There was a problem hiding this comment.
🌀 Claude: Following up: 6debe24 goes further, as you asked. Relation and the shared-CTE workaround are gone. The model is plain SQL again: a view, built from layer_norm(), attention(), mlp() and residual(), each returning its CTEs as text, and a loss that reads the view. Training steps are as fast as before. DataFusion's forward loss takes 2.8 s again, and a deeper model waits on DataFusion computing a shared CTE once.
einfold 0.1.0rc2 allocates large blocks with the system's allocator, so fused attention at 32,768 tokens peaks 0.1 GB lower than on rc1 (xqlsystems/einfold#119). The ml demos job installs the test extra too, as review asked. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
Each piece of the model took the names of its input and output tables
and returned the CTEs it added, as (name, sql) pairs: a shape the
DataFusion workaround for shared CTEs needed, which made the model hard
to read. Now each piece takes relations and returns one, a SELECT that
reads its inputs by name, so a block reads as nanoGPT's Block.forward:
x = residual(x, attention(layer_norm(x, "ln_1", layer), layer))
return residual(x, mlp(layer_norm(x, "ln_2", layer), layer))
_relation.py lists every relation a loss reads, in order, as the CTEs of
one query, and names each one, never as a table it would hide. The SQL
is the same, CTE for CTE, and the demo matches JAX as before.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
|
🌀 Claude: On the review: I made it legible now rather than leave a TODO. You were right that it changed because of the CTE workaround: the builders began returning Also new since the last review:
The description is updated, and leads with demo 1. |
The model is a view again, built by functions that each return a few of its CTEs as text, and the loss is a query that reads it. The DataFusion host no longer runs each CTE on its own, and _relation.py is gone: the demo shows the SQL a reader would write, and the work to make it fast on DataFusion belongs upstream. ddx computes each CTE once on its own, so training steps are as fast as before; DataFusion's forward loss takes 2.8 s rather than 0.04 s, and a model of more than one layer doesn't plan on DataFusion until it computes a shared CTE once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
This adds an
mlextra that brings ddx's gradients intoxarray_sql.Context; ddx differentiates SQL queries. It also adds two ML demos inbenchmarks/ml, and a guide for contributing more. Each demo:The headline: demo 1 writes causal attention as plain SQL, as two sums over joins, and einfold runs it as FlashAttention's algorithm:
einfold has no rule for attention. Its rule fuses any softmax-weighted average of scores computed from a join, and the demo checks that the same rule fuses a control that isn't attention. At 1,024 tokens, the forward pass takes 17–19 ms with einfold, against 0.9–3.0 s on DataFusion and DuckDB.
The extra
pip install "xarray-sql[ml]"installs ddxdb (>=0.4.6,<0.5; ddxdb exchanges Substrait plans with DataFusion 55).xarray_sql/ad.pyfollows the pattern of thegeoextension:Contextregistersddx_stop_gradient.Context.sql()computesgrad,vjpandjvpcalls in a select list. A query that doesn't call them, outside its string literals and comments, never reaches ddx's parser.xql.ad.sql_all(ctx, statements)runs several statements over one backward pass. A training step uses it to update each parameter table.tests/test_ad.pycovers:grad,vjpandjvpagainst their closed forms;sql_allsharing one loss;ddx_stop_gradient;grad(in a string literal.The
testextra now includesml.The demos
01_attention.pyGROUP BYspellings, and a loss that reads it; plus a control that isn't attention, the Nadaraya–Watson estimator (a Gaussian-kernel weighted average)02_nanogpt.pylayer_norm(),attention(),mlp()andresidual()as nanoGPT'sBlockis, as a view, with a loss that reads it. It trains on Tiny Shakespeare with SGD or nanoGPT's AdamW, and is sized by--layers,--embed,--headsand--context.Each script runs the same SQL on every host in
_hosts.py:datafusion:xarray_sql.Context;einfold: the sameContext, with einfold installed as its query planner;duckdb: DuckDB, throughxql.register.All three take ddx's gradients.
_bench.pyholds the shared command line, timing, checks and results table: each step's time, its error against JAX, and its plan's operator count. A gradient that a host refuses is recorded asrefused:with the error. A refusal the demo doesn't list as expected fails the run.For contributors
benchmarks/ml/CONTRIBUTING.md: what a demo is made of, how to write its SQL, what to do when ddx or an engine can't run it, and how to measure.tests/test_ml_demos.py: runs every demo at its smallest size, including an AdamW run, and fails if a demo isn't listed. Theml demosCI job runs it.AGENTS.md: two rules from review:TODO(#NN), and nowhere else in the source.mllabel.Findings
They're written up in full in
benchmarks/ml/README.md, and #285 has the latest numbers. These were measured on an i7-8700 with DataFusion 55.0.0, DuckDB 1.5.6, ddxdb 0.4.6 and einfold 0.1.0rc2.Correctness: ddx's gradients match JAX to about 1e-15 on every host, for every parameter, after each SGD or AdamW step.
einfold finds FlashAttention's plan:
At 32,768 tokens, where the scores alone would take 4.3 GB, the forward pass takes 1.8 s, and peak memory grows by 0.42 GB. The gradient isn't fused yet: that's einfold's next milestone.
The GPT trains on every host:
DataFusion plans a CTE again at every read (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267), so a deep model's plan grows exponentially, and the GPT doesn't grow past one layer on DataFusion yet: its 1-layer loss is 9,882 operators and takes 2.8 s, where DuckDB takes 0.06 s. The demo keeps the SQL as a reader would write it, and the fix belongs upstream: DataFusion's coming
WITH … AS MATERIALIZED(feat: supportWITH ... AS [NOT] MATERIALIZEDCTEs apache/datafusion#25580), and deciding on its own when to compute a CTE once (Decide automatically when to materialize a CTE apache/datafusion#26237).Not done yet
#270 (the GraphCast demo) is stacked on this one.
🤖 Generated with Claude Code
https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h