Array algorithms for tensor computations written in SQL.
Matrix products, tensor contractions and attention layers can all be written in SQL, as joins and aggregates. A matrix product, for one:
SELECT a.i, b.k, SUM(a.v * b.v) AS v
FROM a JOIN b ON a.j = b.j
GROUP BY a.i, b.kA SQL engine runs that as SQL. It builds every row of the join, then hashes each into its group. That's correct, but slow, and for attention it needs memory that grows with the square of the sequence length. einfold finds these computations in a query's plan and runs them with array algorithms instead. Nothing in the query says it's a matrix product: einfold recognizes the shape.
einfold is part of XQL Systems, which builds SQL access to large scientific arrays, such as climate and weather data stored in Zarr, the chunked array format. In XQL's model, a dataset is a table with one row per combination of coordinates, so array math becomes joins and sums. einfold makes those fast. The name joins einsum (Einstein summation, a notation for tensor contractions) with fold, functional programming's word for a reducing pass.
einfold 0.1 is released, for Apache DataFusion, the Rust SQL engine, from Python and from Rust. Its first two milestones are done:
- M1, folds: sums of products over a join, the shape of matrix products and of the gradients that ddx (an XQL Systems project for automatic differentiation of SQL queries) writes.
- M2, fold pipelines: a weighted sum of such a sum, run tile by tile. Attention is one, so einfold rediscovers FlashAttention's algorithm from attention written in plain SQL (demo 1).
Next is M3: the backward pass of attention fused the same way, chains of three or more products, and wider detection (roadmap). DuckDB is planned as a second engine.
python -m pip install --upgrade "einfold[datafusion]"Wheels are built for Linux (x86_64, aarch64) and macOS (Apple silicon, Intel), for Python 3.10 and later. Windows isn't supported yet (#118).
from datafusion import SessionContext
from einfold.datafusion import EinFold
einfold = EinFold()
ctx = SessionContext().with_extensions(einfold)
ctx.register_parquet("a", "a.parquet")
ctx.register_parquet("b", "b.parquet")
table = ctx.sql("""
SELECT a.i, b.k, SUM(a.v * b.v) AS v
FROM a JOIN b ON a.j = b.j
GROUP BY a.i, b.k
""").to_arrow_table()
einfold.plans_with_einfold() # 1: einfold ran the matrix productEvery query in the session is now planned by einfold. Whatever it doesn't recognize runs as DataFusion would run it. xarray-sql, which queries xarray Datasets with SQL, is a DataFusion session, so with_extensions(EinFold()) works there too. The package's README has the details and its limits.
[dependencies]
datafusion = "55.2"
einfold-datafusion = "0.1"einfold 0.1 supports DataFusion 55, from 55.2.
use datafusion::execution::SessionStateBuilder;
use datafusion::prelude::SessionContext;
let state = einfold_datafusion::install(SessionStateBuilder::new().with_default_features()).build();
let ctx = SessionContext::new_with_state(state);- Sums of products over a join, grouped by columns of either side: matrix products, tensor contractions, and weighted sums such as regridding with a weight table. This includes the same sums left-joined from a table, so that every row of the table gets one, which is how ddx writes a gradient.
- A weighted sum of such a sum, fused with it. One query computes scores, and a second sums values weighted by a function of them. Attention is one: weights
exp(s - MAX(s)), normalized by their sum, written with a window or withGROUP BYjoined back. The weights can be other expressions of the score, such as a sigmoid's or a Cauchy kernel's. A plain chain of two matrix products fuses too. - Scores that are any expression of two rows' values, such as a Gaussian kernel's
exp(-power(x - y, 2) / h)over a cross join, for kernel smoothing. - Masks that compare keys: causal, a sliding window, or a band. einfold skips masked tiles rather than computing and discarding them.
einfold rewrites query plans, and doesn't parse SQL. It looks for its shapes in DataFusion's optimized logical plan, just before physical planning. It replaces each with a node that its own operator, EinFoldExec, runs. The operator picks a kernel when it runs, from the data it reads:
- dense, positional kernels when the keys are small integers, as an array's coordinates are;
- sorting-based kernels (radix ranking, then sparse products) for any other keys;
- a fused, tiled kernel for pipelines. It holds one block of scores per thread, with an online softmax as FlashAttention does, so attention's memory grows with the sequence, not its square.
Where a kernel doesn't apply, or the session's memory pool runs short, the operator runs DataFusion's own plan for that part of the query instead. So einfold changes how fast a query runs, not what it returns:
- The same rows as DataFusion, up to the rounding of floating-point sums. NULLs, NaN, infinities and signed zeros follow DataFusion's semantics.
- Deterministic: where einfold computes a sum, it gives the same bits on any number of threads and any batch size. Where DataFusion's plan ran instead, its sums can vary with the number of partitions.
EXPLAIN ANALYZEshows this ashost_planson einfold's operator.
These hold under randomized testing against DataFusion itself, with a nightly simulation soak.
Times on an Intel i7-8700 (6 cores), against DataFusion 55 and DuckDB 1.5. Attention is causal, with 64 features. It's demo 1's query in plain SQL, measured in Rust (benchmark):
| tokens | DataFusion | DuckDB | einfold on DataFusion |
|---|---|---|---|
| 1,024 | 1.50 s, 0.64 GB | 1.33 s, 0.16 GB | 14 ms, 0.07 GB |
| 4,096 | 45.1 s, 5.9 GB | 12.7 s, 0.71 GB | 44 ms, 0.10 GB |
| 8,192 | out of memory | 48.5 s, 1.7 GB | 118 ms, 0.14 GB |
| 65,536 | 5.5 s, 0.69 GB |
At 65,536 tokens, the scores below the diagonal alone would take 17 GB. einfold's results are closer to an exact reference than DataFusion's, and the same bits on 1, 2, 5 and 12 threads.
The gradient of ddx's matrix-product program (SUM(tanh(a·w)²), with a of 50,000 × 16) takes 95 ms with einfold, against 377 ms without, also in Rust (benchmark).
From Python, in xarray-sql's machine-learning demos, on einfold 0.1.0rc1 (results):
- One training step of a GPT written in SQL, at nanoGPT's CPU size (4 heads, 128 channels, context 64): 2.7 s, against 9.5 s on DataFusion and 16.4 s on DuckDB.
- Attention's gradient isn't fused yet (M3). At 1,024 tokens it takes 1.8–2.2 s, against 4.0–4.4 s on DataFusion.
From Python, each query costs about 1.5–3 ms to hand its plan across DataFusion's foreign-function interface. So for tiny queries, such as the GPT at the demo's own size, einfold is slower than DataFusion alone.
crates/einfold-ir: the description of a fold or a pipeline, independent of any engine.crates/einfold-kernels: the kernels, on Apache Arrow data, also engine-independent.crates/einfold-datafusion: detection in DataFusion's plans, andEinFoldExec. The only crate that depends on an engine.crates/einfold-testkit: the seeded random generator and shrinker the property tests share.python: theeinfoldpackage for DataFusion's Python bindings.benchmarks: the milestones' gates, M1 and M2.
- Design: the problem, principles, architecture, and roadmap. The supplement holds the details, numbered to match.
- Lessons: what review and testing have taught us that the design doesn't say.
- Target demos: FlashAttention rediscovered from SQL; NanoGPT in SQL; GraphCast weather forecasts that read only the data they need.
- Spikes: short experiments that settled design questions, each with code and a report.
- Layout convention: a draft XQL Systems convention for Zarr datasets to declare how their data is ordered and summarized, so that engines can scan less.
See CONTRIBUTING.md, and AGENTS.md for how AI agents work on einfold.
- xarray-sql: query xarray Datasets with SQL.
- ddx: automatic differentiation of SQL queries, so that models written in SQL can be trained.
- duckdb-zarr: query Zarr stores from DuckDB.
einfold is licensed under the Apache License, Version 2.0. It follows the REUSE specification: every file states its copyright and license, and reuse lint checks it.