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

Ouverte
  • Quelque chose ne va pas ?
#283 1 commentaire 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 3 jours

Personne n'a encore pris cette issue.

Évaluation

Difficulté
4/5
Temps estimé
3-5 jours
Accessibilité débutants
25/100
Type d'issue
Fonctionnalité
Clarté
Clairement spécifiée
Activité
Active
Stack technique
python, rust
Domaine
api, backend

Piste de recherche

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.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

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.
Langage dominant
Python
Étoiles
142
Forks
19
Merge moyen
3 j 47 min
PR mergées (30 j)
8

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de xqlsystems/xarray-sql

Toutes les issues de xqlsystems/xarray-sql

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.