ad.prepare: reuse each step's physical plan with reset_plan_states instead of planning it again #213

🌀 Claude (xarray-sql): ad.prepare's later runs still plan every step from scratch. Reusing physical plans in Rust, with reset_plan_states, should cut most of the GPT's gradient step.

Summary

ddx#203's Prepared::run (crates/ddx-ad/src/sql/prepare.rs) skips tracing and keying on later runs, but each step it materializes still goes through materialize in crates/ddx-datafusion/src/engine.rs. That decodes the step's Substrait plan into a DataFrame, optimizes it, builds a physical plan and collects it, all on every run. On xarray-sql's GPT demo (benchmarks/ml/02_nanogpt.py), that's nearly all of the gradient step.

At ddx 3dc0e62 (the ddx#203 branch), on an Intel i7-8700 with load under 0.3:

The data is tiny: 64 tokens and 16-wide weights. JAX takes 0.37 ms for the whole step. xarray-sql measured about 76 small queries per step, at about 3.3 ms to plan each and 0.65 ms to run a reused plan (from xarray-sql#283). So about 285 ms is planning, and perhaps 50 ms is execution.

Proposal

Plan each step once on the first run. On later runs, execute datafusion::physical_plan::execution_plan::reset_plan_states(plan) on each one. That's public in Rust today, so this needs nothing from datafusion-python (I've added a Python binding on a fork, but ddx doesn't need it). Three things to handle:

1. A cached physical plan holds a snapshot of its tables. A MemTable scan copies the batches it reads into its DataSourceExec when the plan is built. materialize registers a new MemTable for each step on every run, so a cached plan of a later step would read the first run's values. A later step needs a table whose rows are read at execute() time, not at planning time: for example, one TableProvider per step, registered once, whose ExecutionPlan reads an Arc<RwLock<Vec<RecordBatch>>> slot that materialize refills.
2. Turn off dynamic filter pushdown, or check that it's off. A reset plan keeps the dynamic filters its first run pushed down. When a join's build-side keys change, the reset copy can silently drop rows. I reproduced this on DataFusion 55: a reset join returned [] where planning again returned [7, 8]. All four datafusion.optimizer.enable_*dynamic_filter_pushdown settings must be false. SET datafusion.optimizer.enable_dynamic_filter_pushdown = false sets all four.
3. Run the reset copy, never the cached plan itself. A plan that has run keeps a hash join's build side, and a RepartitionExec panics with "partition not used yet" the second time it runs. reset_plan_states takes about 1.6 µs on a small plan, and its cost doesn't grow with the data, since it copies none.

How xarray-sql measured it

From xarray-sql's repo, on its local perf-stack branch (xarray-sql#287 plus #269), with ddx-local built into the venv:

import time
import ddxdb
from ddxdb import ad
from datafusion import SessionContext
prepared = ddxdb.Prepared(ad.DataFusionBackend(ctx), gpt.step_statements(gpt.LEARNING_RATE), "datafusion")
for _ in range(4):
    a = time.perf_counter()
    with prepared.run() as rewritten:
        b = time.perf_counter()
        dfs = [SessionContext.sql(ctx, s) for s in rewritten.statements]
        tables = [df.to_arrow_table() for df in dfs]
    print(f"program {1e3 * (b - a):.0f} ms, statements {1e3 * (time.perf_counter() - b):.0f} ms")

ctx is the demo's Context, with the weights and a batch registered and the model's view created, as in 02_nanogpt.py. The full script is grad_parts.py, which I can attach.

Context

xarray-sql is working on the loss side, where the 9,882-operator forward plan is 87 distinct subtrees repeated. Once that's fixed, the loss should take about 15 ms, which leaves this gradient step as most of a training step. Results are tracked on xarray-sql#285.