Batch Inference with ai_query(): Every Argument, Structured Output and Error Handling

What this deep dive adds

Chapter 37 of the course covers when batch inference beats a real-time endpoint, the shaping arguments by name, and the failOnError default that sinks large jobs. This post is the reference behind it: every argument with its runtime requirement, the request shape for each endpoint type, structured output in both styles, a job that keeps its failed rows for a retry, and how to run and cost it in production.

Batch inference with ai_query() is a Section 4 (Assembling and Deploying Applications) objective in the 18 March 2026 exam guide. The serving endpoints it calls are covered in the Foundation Model APIs and AI Gateway post.

Every argument

Two details catch people. When modelParameters leaves temperature unset, ai_query() uses 0.0, not the endpoint's usual default, so batch output is deterministic unless you ask otherwise. And the function needs Databricks Runtime 15.4 LTS or above (the docs announce a move to 17.3 LTS) and does not run on Databricks SQL Classic warehouses.

From the Databricks SQL reference page "ai_query function" (last updated 5 November 2026). Newer examples there write models as system.ai.<model>; this post uses the databricks- endpoint names that the AI Functions overview recommends for batch work. Use the name your workspace's Serving page lists.

The request for each endpoint type

A chat or completion model on Foundation Model APIs takes a string, so build the prompt with ||, CONCAT() or format_string():

SELECT ticket_id,
       ai_query(
         'databricks-meta-llama-3-3-70b-instruct',
         'Summarise this support ticket in one sentence: ' || body
       ) AS summary
FROM support.tickets;

A custom model takes the features it was trained on, as a struct, and a returnType when its signature does not supply one:

SELECT order_id,
       ai_query(
         endpoint   => 'churn-classifier',
         request    => named_struct('tenure_months', tenure_months,
                                    'plan', plan,
                                    'monthly_spend', monthly_spend),
         returnType => 'DOUBLE'
       ) AS churn_score
FROM crm.customers;

A custom chat model or agent that expects a messages array gets one built with named_struct and ARRAY, with returnType describing the response shape. Wrapping a call in a SQL function (CREATE FUNCTION ... RETURN ai_query(...)) and granting EXECUTE on it lets analysts reuse a vetted prompt without seeing or editing it.

Structured output: parameters and schemas

For extraction and classification, ask for a schema instead of parsing free text. A JSON schema with "strict": true constrains the model to it, and an enum limits a field to fixed labels:

SELECT ticket_id,
       ai_query(
         'databricks-meta-llama-3-3-70b-instruct',
         'Classify this support ticket: ' || body,
         modelParameters => named_struct('max_tokens', 200),
         responseFormat  => '{
           "type": "json_schema",
           "json_schema": {
             "name": "ticket_triage",
             "schema": {
               "type": "object",
               "properties": {
                 "category": {"type": "string",
                              "enum": ["billing", "bug", "account", "other"]},
                 "urgent":   {"type": "boolean"}
               }
             },
             "strict": true
           }
         }'
       ) AS triage
FROM support.tickets;

With the json_schema type the column holds a JSON string, so read fields with get_json_object(triage, '$.category') or convert it with from_json. The DDL style ('STRUCT<category:STRING, urgent:BOOLEAN>') is a shorter way to state the same schema. responseFormat works only against chat foundation models; for a custom model, the shape comes from returnType.

failOnError: a job that keeps its failures

With the default failOnError => true, one row that the endpoint rejects (an input over the model's limit, for example) fails the whole query, and a ten-million-row run loses all its work. Set it to false and each row returns a struct with two fields:

That turns a fragile job into a two-step pattern: write everything, then retry only what failed.

CREATE OR REPLACE TABLE support.ticket_summaries AS
SELECT ticket_id,
       ai_query(
         'databricks-meta-llama-3-3-70b-instruct',
         'Summarise this support ticket in one sentence: ' || body,
         failOnError => false
       ) AS result
FROM support.tickets;

-- Rows to fix and rerun
SELECT ticket_id, result.errorMessage
FROM support.ticket_summaries
WHERE result.errorMessage IS NOT NULL;

Successful rows have a null errorMessage; failed rows have a null response. The flag covers errors from the model endpoint only; per the reference, a row that fails for any other reason still fails the whole query.

Running it in production

• One query, full dataset. Databricks manages parallelism, retries and scaling, and its guidance is to submit the whole table in a single query rather than splitting it into small batches yourself.
• Databricks-hosted endpoints. For batch work the docs recommend the provisionless, Databricks-hosted foundation model endpoints (prefixed databricks-) over provisioned throughput.
• Where it runs. The same call works in a scheduled job, a Lakeflow Spark Declarative Pipelines materialized view (with expectations to validate the output), or Structured Streaming for incremental inference.
• What it costs. ai_query() is billed for the endpoint it calls, on top of the query's own compute. Hosted foundation model usage appears in the billing system table under the MODEL_SERVING product with the BATCH_INFERENCE offering type, which is how to attribute it to a job or warehouse.

For common tasks, the task-specific functions (ai_classify, ai_extract, ai_summarize, ai_parse_document and others) need no endpoint or prompt, and Databricks recommends them for getting started. ai_query() is the general-purpose option when you choose the model, write the prompt or call a custom endpoint.

From the Databricks pages "Transform unstructured data using AI Functions" and "Deploy batch inference pipelines".

What this means for the exam

Distractors fix the wrong layer: a lower temperature for errors, a Python loop for scale, or a classic warehouse, where AI Functions do not run.

FAQ