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

Offen
  • Stimmt etwas nicht?
#283 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 3 Tagen

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Anfängerfreundlichkeit
25/100
Issue-Typ
Feature
Klarheit
Klar beschrieben
Aktivitätsstatus
Aktiv
Tech-Stack
python, rust
Bereich
api, backend

Rechercherichtung

The issue proposes adding ExecutionPlan.reset_state() in datafusion-python, but this is a draft to review and re-file in apache/datafusion-python. Start by checking the Python ExecutionPlan binding and the Rust reset_plan_states API; the issue names no repository files or tests. Done means returning a reset copy while leaving the original unchanged, with the documented limits around dynamic filters and recursive query work tables.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

ml

🌀 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

  • apache/datafusion#19351: stateless execution plans for plan caching, which would make the reset unnecessary.
  • apache/datafusion#26229: MaterializedCteExec::reset_state (new in DataFusion 56) clears its buffer in place.
Vorherrschende Sprache
Python
Sterne
142
Forks
19
Ø Merge
3 T. 47 Min.
Gemergte PRs (30 T.)
8

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus xqlsystems/xarray-sql

Alle Issues in xqlsystems/xarray-sql

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.