Editing physical plans

You can inspect and edit a physical plan before executing it. This example inserts a barrier between a filter and a join, then shows how to add a Python function through the query builder.

The example runs on the CPU. It plans queries and prints their structure; it does not load model weights or run inference.

The full example is in demos/plan_walkthrough.py.

Run the example

uv run python demos/plan_walkthrough.py

The query

The example uses a small version of a FEVER query from QUAIL-B. In FEVER, each claim is a factual statement and each evidence passage is a Wikipedia paragraph. The goal is to find which passages support which claims.

The input contains three claims and three evidence passages. The query first keeps claims about a person. It then checks whether the evidence passages support those claims. A URL equality limits the pairs evaluated by the model.

The script creates a session and registers both tables. It then builds this query with the Python builder API:

from quail import col, prompt

PERSON = "Is the claim in DOCUMENT {0} about a person?"
SUPPORT = (
    "Does the passage in DOCUMENT {1} support the claim in "
    "DOCUMENT {0}?"
)

query = (
    session.docs("claims").alias("c")
    .ai_if(prompt(PERSON, col("c.claim")))
    .join(
        session.docs("evidence").alias("e"),
        on=col("c.evidence_wiki_url") == col("e.id"),
    )
    .ai_if(prompt(SUPPORT, col("c.claim"), col("e.text")))
    .select("c.id", "e.id")
)

The .join(..., on=...) restricts which combinations of claims and evidence the model evaluates.

Inspect the plan

Call query.explain() to see the plan without executing it. The physical plan identifies each operator and its estimated work:

physical: backend=quail, model=qwen3-4b-fp8, workers=1
  KV=bf16, chunk budget=110,376 tokens

                                                 est. rows  est. pass  est. time
  Project: c.id, e.id
    AiJoin: anchor=c                                                     4.29 ms
      KV: anchor=streamed from its filter, retain after join=no
      join 1 full (c, e) over hash_join:c-e            0.6       20%*
      AiFilter: c                                      0.6               2.39 ms
        survivors stream into the join with KV pinned
        predicate 1                                      3       20%*
        Scan claims as c
          tokens=306, mean_doc_tokens=102
      Scan evidence as e
        tokens=4,506, mean_doc_tokens=1,502
      HashJoin: c.evidence_wiki_url = e.id

The line "survivors stream into the join with KV pinned" means a claim can enter the join as soon as it passes the filter. KV is the model's saved key and value vectors for the claim's tokens. The engine keeps that KV in GPU memory and reuses it when evaluating evidence passages.

Insert a barrier

Insert a Barrier to require the filter to finish over all claims before the join starts. The barrier collects the passing claim IDs, then supplies them to the join.

from quail.physical import Barrier

plan = query.plan()
edited = plan.insert(
    Barrier(node_id="barrier:c", next_anchor="c", aliases=("c",)),
    between=("ai_filter:c", "ai_join:c"),
)

The barrier separates the filter and join into two pipelines. The edited plan sets keep_kv=True on the filter so passing claims' KV can be retained for the join. Retention is limited by available KV memory:

      Barrier: next_anchor=c
        AiFilter: c
          retain KV for later joins

We can remove the barrier and get back the original plan:

assert edited.remove("barrier:c") == plan

Only Barrier, Foreign, and Exchange nodes can be removed. A Foreign node holds a Python function, so its removal can change the results. The removal of an LLM node raises an error:

PlanEditError: 'ai_filter:c' is a AiFilter; removing it would change
what the query means

Plan editing API

Start with query.plan() to inspect the plan without running inference. If planning returns a Refusal, resolve the reported constraint before editing it.

Each edit returns a new plan, validates port connections and cycles, and re-estimates the graph. An invalid edit raises PlanEditError.

Only Foreign, Barrier, and Exchange nodes can be removed. Quail rejects removal of a scan, model operator, or filter. Removing a Foreign node can change results if its Python function filters rows.

Insert a user function

With .apply(), we can run our own Python function between operators. The full signature is:

query.apply(fn, columns=(), *, name=None, ids=None, kind="per_batch")

• fn receives a dict[str, pa.Table] keyed by alias. Each table has the alias's row indices plus the columns you listed in columns.
• columns is a list of col(...) references that your function needs to read.
• After a .join(), fn returns a table with both alias columns (e.g. c and e) selecting which combinations the next AI predicate evaluates.
• Without a preceding .join(), fn returns the ids to keep.
• kind="per_batch" runs on each batch of rows in a pipeline. Outside a pipeline, it gets all input rows in one call. kind="barrier" waits for all survivors, then runs once.

For example, instead of an ON clause, we can match claims to evidence with an Arrow join inside our function:

def match_by_url(tables):
    """Match each claim to the evidence page its URL names."""
    return tables["c"].join(
        tables["e"],
        keys=["evidence_wiki_url"],
        right_keys=["id"],
        join_type="inner",
    ).select(["c", "e"])

query = (
    session.docs("claims").alias("c")
    .ai_if(prompt(PERSON, col("c.claim")))
    .join(session.docs("evidence").alias("e"))
    .apply(match_by_url, columns=[col("c.evidence_wiki_url"),
                                  col("e.id")])
    .ai_if(prompt(SUPPORT, col("c.claim"), col("e.text")))
    .select("c.id", "e.id")
)

The function determines which row pairs reach the model predicate. Quail uses the returned c and e index columns to identify those pairs; it does not inspect the function's internal logic.

In the physical plan, match_by_url appears as a Foreign node:

      Foreign: match_by_url (per_batch, pairs) on c, e
        columns: c.evidence_wiki_url, e.id