Repository navigation
DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267
Description
DataFusion plans a CTE once for every place that reads it, so a query whose CTEs read earlier CTEs more than once grows exponentially with its depth. DuckDB computes such a CTE once. ML models written in SQL hit this constantly: a layer norm reads its input three times, a residual connection twice, and a message-passing round reads the node states as sender, receiver and residual.
Reproduction
Each CTE joins the previous one with itself:
from datafusion import SessionContext
def query(depth):
ctes = ["c0 AS (SELECT i, x FROM t)"] + [
f"c{d} AS (SELECT a.i, a.x + b.x AS x FROM c{d-1} a JOIN c{d-1} b ON a.i = b.i)"
for d in range(1, depth + 1)
]
return f"WITH {', '.join(ctes)} SELECT SUM(x) FROM c{depth}"
ctx = SessionContext()
ctx.sql("CREATE TABLE t AS SELECT * FROM (VALUES (1, 1.0), (2, 2.0)) v(i, x)").collect()
plan = ctx.sql(query(8)).execution_plan().display_indent()| depth | DataFusion operators | DataFusion table scans | DuckDB 1.5.5 operators |
|---|---|---|---|
| 2 | 16 | 4 | 13 |
| 4 | 64 | 16 | 23 |
| 6 | 256 | 64 | 33 |
| 8 | 1,024 | 256 | 43 |
In the ML benchmarks
The same queries, on DataFusion 54.0.0 through xarray_sql.Context and on DuckDB through xql.register, Intel i7-8700:
| Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB |
|---|---|---|---|---|
| 1-layer GPT loss, batch 4 (112 lines of SQL) | 12,166 | 120 | 8.4 s | 0.055 s |
| GraphCast-shaped forecast, 1 processor round | 1,816 | 243 | 43 s | 17 s |
| GraphCast-shaped forecast, 2 processor rounds | 5,948 | 282 | 159 s | 13 s |
A table read through xarray-sql is read again for every inlined copy. With ERA5 registered lazily from GCS, the 1-round forecast spent most of its time waiting on the network, so the benchmark loads its three time steps into memory first.
A training step through ddx (forward and backward) takes 2.7 s on DataFusion, less than the 8.4 s forward query alone, because ddx stores each step of its program as a table.
Upstream
DataFusion is working on this. apache/datafusion#25580 added WITH ... AS MATERIALIZED, in the PostgreSQL dialect only (apache/datafusion#17737), and apache/datafusion#26237 proposes deciding automatically, as PostgreSQL and DuckDB do. Neither is in DataFusion 54.
Options for xarray-sql
- Once a DataFusion release has
AS MATERIALIZED, check that it works with xarray-sql's table providers and thedatafusiondialectContextuses, and document it. - Follow Decide automatically when to materialize a CTE apache/datafusion#26237, which would fix this without changing queries.
- Until then, document the behavior in
docs/limitations.md, with the workaround: register an intermediate result as a table.
Activity
alxmrs commentedon Oct 25, 2026
On DataFusion 55.0.0 (2026-10-24, Intel i7-8700), the GPT benchmark now takes a model size from the command line (
02_nanogpt.py --layers N), which shows the growth per layer:Building the 2-layer logical plan alone (
ctx.sql(...).logical_plan()) takes 8.6 s.Nearly all of the 1-layer loss's 2.5 s is planning: the same physical plan, built once and run again with
SessionContext.execute, takes 1.8 ms. So once each CTE is planned once, planning a query once per training loop rather than once per step matters next. Today's API can't do that safely: a physical plan reads the batches its table scans held when it was planned, and a hash join keeps its build side across runs untilreset_state, which datafusion-python doesn't expose.Upstream status: apache/datafusion#25580 (
WITH … AS [NOT] MATERIALIZED) merged on 2026-10-15, after 55 branched, so it ships in DataFusion 56. apache/datafusion#26269, which makes the nodesLogicalPlanvariants, merged today. Open: apache/datafusion#26276 (decide automatically, behinddatafusion.optimizer.enable_automatic_cte_materialization, in every dialect), apache/datafusion#26270 (keep a materialized CTE's output per run, not in the plan: what reusing a plan needs), and the statistics, ordering and pushdown follow-ups (apache/datafusion#26271, apache/datafusion#26428, apache/datafusion#26274, apache/datafusion#26275).Plan for xarray-sql, on DataFusion 56: turn on automatic materialization in
Context(or writeAS MATERIALIZEDin the benchmarks' SQL if the automatic rule hasn't landed), check that it works with xarray-sql's table providers, and rerun the GPT at 2 and more layers. ddx will need to read the new plan nodes, since it differentiates DataFusion's plans.Run the GPT's CTEs once each on DataFusion, so it trains past one layer
alxmrs commentedon Oct 28, 2026
🌀 Claude: The GPT demo now works around this, in ef30a6e on
ml-demos(#269). It stays until DataFusion decides on its own when to materialize a CTE (apache/datafusion#26237).What it does:
(name, sql)pairs.query_with(ctes, select)runs each CTE on its own, into an in-memory table that later CTEs read. That's what DuckDB does for a CTE read more than once. It usescollect_partitionedandregister_record_batchesrather thanDataFrame.cache(), because einfold can't plan a join of cached tables (A join of tables from DataFrame.cache() fails to plan with einfold (Input field name e does not match with the projection expression mean) einfold#93).logitsany more. When DataFusion creates a view, it writes the view's CTEs out, so creating the 3-layer model's view alone ran out of 8 GB. Each statement now starts with the model's CTEs. Given those, ddx computes each shared CTE once.Why not the other routes:
AS MATERIALIZED. DataFusion 55 accepts it only in the PostgreSQL dialect, and ignores it there:random()read twice gives two values. The change that honours it (feat: supportWITH ... AS [NOT] MATERIALIZEDCTEs apache/datafusion#25580) isn't released yet.The results were measured with ddxdb 0.4.4, DataFusion 55.0.0 and DuckDB 1.5.6, at batch 4, on an Intel i7-8700. Every step matches JAX.
At nanoGPT's CPU size (1 layer, 4 heads, 128 channels, context 64, batch 12), the loss takes 2.4 s on DataFusion, down from 108 s and 9 GB.
alxmrs commentedon Oct 31, 2026
🌀 Claude: GraphCast hits this too, now that its demo is written as views (#270). Each processor round is a view of the round before, and it reads the mesh three times: as each edge's sender, as its receiver, and as a residual. On DataFusion 55.0.0:
mesh_0(the encoder)mesh_1mesh_2mesh_3The plan grows about 4.7× per round, and a forecast has 16. Computing each shared read once would let the demo run as written. For comparison, the earlier version of the demo stored each stage as a table, so each stage was computed once. That version ran the encoder and one round on einfold in about 10 minutes.
DuckDB has the opposite problem on the same views: its plan stays small (504 operators for
mesh_1), but planning takes 45 s for one round and more than 10 minutes for two (#292).