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

Đang mở
  • Có gì không ổn?
#283 1 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 3 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
25/100
Loại issue
Tính năng
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, rust
Lĩnh vực
api, backend

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.
Ngôn ngữ chính
Python
Star
142
Fork
19
Merge trung bình
3 ngày 47 phút
Pull request đã merge (30 ngày)
8

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của xqlsystems/xarray-sql

Tất cả issue của xqlsystems/xarray-sql

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.