datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream #283

🌀 Claude (xarray-sql)

A draft to review, then re-file upstream in apache/datafusion-python.

Summary

datafusion-python can build a physical plan once (DataFrame.execution_plan()) and run it again (SessionContext.execute(plan, partition)). But a plan that has run once keeps state from that run. A hash join keeps its build side, so a second run reads its left input as it was in the first. DataFusion's Rust API resets that state with reset_plan_states, which calls ExecutionPlan::reset_state on every node. datafusion-python has no binding for it, and nothing else in Python can reach it:

• with_new_children keeps a hash join's build side (HashJoinExecBuilder clones left_fut), and FFI_ExecutionPlan forwards with_new_children but not reset_state. So an extension can't reset a plan through datafusion-ffi either.
• Decoding the plan again (ExecutionPlan.from_proto(ctx, plan.to_proto())) gives fresh nodes, but it copies every in-memory table into the plan. On a deep plan it also fails with "recursion limit reached".

Why it matters

Planning dominates small, repeated queries. Running a planned query again is how a prepared statement or a training loop would avoid it. In xarray-sql's ML benchmarks, which train a small GPT in SQL with autodiff from ddx, one training step runs 76 small queries (a median of 16 operators each). Planning each from scratch takes 3.3 ms, and running a reused physical plan takes 0.65 ms (DataFusion 55.0.0, target_partitions=1, Intel i7-8700). The GPT's forward pass has a 13,756-operator plan that takes 2.1 s to plan and 2.4 ms to run again. All 2,988 of its joins are HashJoinExec with mode=CollectLeft, which keeps its build side. With one partition, DataFusion uses CollectLeft even when hash_join_single_partition_threshold is 0.

Reproduction

datafusion-python 55.0.0. pyarrow shares the NumPy arrays' memory, so writing into them changes the tables in place:

import numpy as np
import pyarrow as pa
from datafusion import SessionConfig, SessionContext

ctx = SessionContext(SessionConfig().with_target_partitions(1))
a = np.arange(4.0)  # the build side; pyarrow shares its memory
b = np.full(4, 10.0)
ctx.register_record_batches("a", [[pa.record_batch({"i": np.arange(4), "v": a})]])
ctx.register_record_batches("b", [[pa.record_batch({"i": np.arange(4), "w": b})]])
query = "SELECT sum(t.v2 * b.w) AS s FROM (SELECT i, 2 * v AS v2 FROM a) t JOIN b USING (i)"
plan = ctx.sql(query).execution_plan()

def run(plan):
    return [batch.to_pyarrow()["s"][0].as_py() for batch in ctx.execute(plan, 0)]

print(run(plan))                    # [120.0]
a[:] = 100.0
print(run(plan))                    # [120.0]: the hash join's build side, from the first run
print(ctx.sql(query).to_pydict())   # {'s': [8000.0]}
b[:] = 1.0
print(run(plan))                    # [12.0]: the probe side is read again

Proposal

fresh = plan.reset_state()  # a copy of the plan, with every node's state reset

ExecutionPlan.reset_state() would return a new ExecutionPlan from reset_plan_states(self.plan), and leave self unchanged. It's a few lines of Rust. Its docstring should carry reset_plan_states's limits: it doesn't reset dynamic filters or a recursive query's work table (apache/datafusion#19905). A caller reusing plans should turn off datafusion.optimizer.enable_dynamic_filter_pushdown until that's fixed.

A hash join isn't the only node that keeps state. A plan with a RepartitionExec runs once, and running it again panics: "partition not used yet", a PanicException in Python. The plan's partitions run one after another work the first time. reset_plan_states covers both, and so does the proto round trip, on a plan shallow enough to decode.

One smaller gap for the same use: the Python ExecutionPlan has no schema(), so a caller that gets no batches back can't build an empty table of the right type.

Related

• Stateless execution plans for plan caching apache/datafusion#19351: stateless execution plans for plan caching, which would make the reset unnecessary.
• MaterializedCteExec::reset_state clears its shared buffer in place apache/datafusion#26229: MaterializedCteExec::reset_state (new in DataFusion 56) clears its buffer in place.