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

Abierto
  • ¿Algo no está bien?
#283 1 comentario 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
25/100
Tipo de issue
Nueva funcionalidad
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, rust
Área
api, backend

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.
Lenguaje dominante
Python
Estrellas
142
Forks
19
Merge medio
3 d 47 min
PR fusionados (30 d)
8

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de xqlsystems/xarray-sql

Todos los issues de xqlsystems/xarray-sql

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.