Python API

You can build queries with Python method calls instead of writing SQL. Import col and prompt to get started:

from quail import col, prompt

Example

query = (
    session.docs("reviews").alias("r")
    .ai_if(prompt('Is this review positive?\n\n{0}', col("r.body")))
    .ai_join(
        session.docs("aspects").alias("a"),
        prompt('Does the review discuss this aspect?\n\n{0}\n\n{1}',
               col("r.body"), col("a.aspect")))
    .select("r.id", "a.aspect")
)
result = query.collect()

Operators

ai_if

ai_if(prompt, selectivity=None)

Filter rows with an AI predicate. The model evaluates the prompt on each row and keeps those that return TRUE.

query = (
    session.docs("reviews").alias("r")
    .ai_if(
        prompt('Is this review positive?\n\n{0}', col("r.body")),
        selectivity=0.6)
    .select("r.id")
)
result = query.collect()

selectivity is optional; a predicate without one is planned as if 20% of its rows pass. If you want to apply several predicates, add multiple ai_if calls. Quail reorders them by estimated cost to reduce total work. Pass order="as_written" to select() to keep the written order.

ai_join

ai_join(others, prompt, selectivity=None, anchor=None, semantics="full")

Join two or more tables with an AI predicate. The model evaluates the prompt on each pair of rows.

query = (
    session.docs("reviews").alias("r")
    .ai_join(
        session.docs("aspects").alias("a"),
        prompt('Does the review discuss this aspect?\n\n{0}\n\n{1}',
               col("r.body"), col("a.aspect")),
        selectivity=0.15)
    .select("r.id", "a.aspect")
)
result = query.collect()

selectivity and anchor are optional. Pass semantics="exists" to keep outer rows that match at least one inner row, or semantics="anti" to keep outer rows that match none.

join

join(other, on=None)

Add relational predicates to restrict which pairs the model sees. Follow join() with ai_if() over both tables:

query = (
    session.docs("claims").alias("c")
    .join(
        session.docs("evidence").alias("e"),
        on=col("c.wiki_url") == col("e.id"))
    .ai_if(
        prompt('Does the passage support the claim?\n\n{0}\n\n{1}',
               col("c.claim"), col("e.text")))
    .select("c.id", "e.id")
)
result = query.collect()

limit

limit(n)

Return at most n rows.

query = (
    session.docs("reviews").alias("r")
    .ai_if(prompt('Is this review positive?\n\n{0}', col("r.body")))
    .limit(100)
    .select("r.id", "r.body")
)

select

select(*cols)

Specify which columns to return. This must be the last operator and returns a Query.

.select("r.id")               # one column
.select("r.id", "a.aspect")   # multiple columns
.select("*")                   # all columns from all tables

collect

collect(limit=None)

Execute the query and return a pyarrow.Table. This is when the model runs. See Explain plans and runs for run(), explain(), explain(analyze=True), and the execution report.

User-defined functions

User-defined functions run on the CPU between AI operators. They do not run on the GPU. Quail may reorder AI filters and joins around each other, but it will not move an AI operator across a user-defined function.

apply

apply(fn, columns=[], name=None, ids=None, kind="per_batch")

Insert a Python function into the query plan. The function receives a dict of Arrow tables keyed by alias.

The main use case is pairing rows from two tables before an AI join. For example, a claims table has an evidence_wiki_url column that names the evidence page for each claim. Instead of evaluating the AI predicate on every claim-evidence pair, you pair them by URL first:

def match_by_url(tables):
    """Pair each claim with the evidence row its URL names.

    tables["c"] is a pyarrow.Table with columns "c" (the document ids),
    "evidence_wiki_url", and any other columns requested.
    tables["e"] has columns "e" (the document ids) and "id".
    The returned table must have columns "c" and "e" (the id columns
    from each side) to identify the pairs.
    """
    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")
    .join(session.docs("evidence").alias("e"))
    .apply(match_by_url,
           columns=[col("c.evidence_wiki_url"), col("e.id")])
    .ai_if(
        prompt('Does the passage support the claim?\n\n{0}\n\n{1}',
               col("c.claim"), col("e.text")))
    .select("c.id", "e.id")
)
result = query.collect()

You can also use apply after a filter to drop rows with custom Python logic before the next AI operator runs.

kind="per_batch" (default) runs the function on each batch as they stream through.

apply_table

apply_table(fn, columns=[], name=None, ids=None)

Like apply, but waits for the preceding stage to finish and calls fn once with all surviving rows. Equivalent to apply(..., kind="barrier").

Use this when your function needs to see every surviving row at once. With apply, rows arrive in batches as they stream through. With apply_table, Quail finishes the preceding stage first, then calls your function once with all the results.

import pyarrow.compute as pc

def deduplicate(tables):
    reviews = tables["r"]
    unique = pc.unique(reviews.column("body"))
    mask = pc.is_in(reviews.column("body"), unique)
    return reviews.filter(mask).column("r")

query = (
    session.docs("reviews").alias("r")
    .apply_table(deduplicate, columns=[col("r.body")])
    .ai_if(prompt('Is this review positive?\n\n{0}', col("r.body")))
    .select("r.id")
)
result = query.collect()

Note that apply_table fixes the execution order at that point in the query. The planner will not reorder operators across it.

Errors

Mistakes raise quail.logical.CompileError at the call that caused them.

Python API

On this page