Batch inference with ai_query()

If a workload is a table of rows rather than a person waiting, it is a batch job. ai_query is the general-purpose AI Function that sends each row to a model from SQL or Python, with no endpoint for you to manage when the model is Databricks-hosted.

Databricks recommends starting with a task-specific function such as ai_classify, ai_extract or ai_summarize when one matches, and turning to ai_query for precise control over prompt, parameters or output, or for custom, fine-tuned or external models.

SELECT text, ai_query(
  "system.ai.meta-llama-3-3-70b-instruct",
  "Summarise this article: " || text,
  modelParameters => named_struct(
    'max_tokens', 100, 'temperature', 0.2),
  failOnError => false
) AS summary
FROM main.docs.articles;

Requirements and return behaviour

ai_query is not available on Classic SQL warehouses. It needs Databricks Runtime 15.4 LTS or above, with 18.2 or above recommended. These arguments decide its behaviour:

Which models and what governance

Requests to system.ai services route through Unity Gateway automatically, but only part of it applies.

The caller needs EXECUTE on the model service in system.ai. For an endpoint, the function reference says the definer must have CAN QUERY.

Common trap: Relying on Unity Gateway guardrails to screen a batch job - guardrails, inference tables and rate limits are not applied to ai_query requests. Filter and validate in the pipeline instead, and use budgets for spend control.

Common trap: Splitting a large dataset into small batches to protect the service - AI Functions automatically handle parallelisation, retries and scaling, and manual splitting can reduce throughput. Submit the full dataset in one query and set failOnError to false so that failures come back row by row.

Production patterns

Batch pipelines can be deployed as Lakeflow pipelines, scheduled jobs, or Structured Streaming for near-real-time or micro-batch work. A useful sequence is ingest, infer, validate, then transform. The documented pipeline checks model output with data quality expectations before later steps use it, which matters when no person reads each row.

CREATE OR REFRESH MATERIALIZED VIEW doc_tags AS
SELECT
  inputs,
  ai_query(
    "system.ai.meta-llama-3-3-70b-instruct",
    "Extract the category: " || inputs,
    responseFormat => 'STRUCT<result:STRUCT<category:STRING>>'
  ) AS meta_data
FROM docs_raw;

Select a model from the batch-optimised list in the ai_query documentation, because a non-optimised model can reduce throughput and lengthen jobs. For batch, Databricks recommends system.ai services instead of provisioned throughput endpoints.

Get the whole book

This note is one section of Ultra Transcenders DBX-GENAI: Databricks Generative AI Engineer Associate, an independent study guide that explains every topic the exam covers by technology, with comparison tables, diagrams and the common traps, plus a glossary linked to the Databricks documentation.

About the book · DBX-GENAI terms in the glossary · All DBX-GENAI study notes