ai_query

Skill: databricks-ai-functions

What You Can Build

ai_query is the general-purpose AI function for everything the task-specific functions can’t handle — nested JSON extraction, custom model endpoints, multimodal image analysis, and fine-grained sampling control. Use it when you need structured output schemas, want to target a specific model, or need to process images alongside text.

In Action

“Write a SQL query that extracts structured invoice data as JSON from parsed document text, using a specific model endpoint.”

SELECT ai_query(
    'databricks-claude-sonnet-4',
    CONCAT('Extract invoice as JSON: ', text_blocks),
    responseFormat => '{"type":"json_object"}',
    failOnError => false
) AS ai_response
FROM catalog.schema.parsed_documents;

Key decisions:

• responseFormat => '{"type":"json_object"}' forces the model to return valid JSON instead of free-form text
• failOnError => false returns a {response, error} struct instead of killing the entire batch on a single row failure
• The endpoint name (databricks-claude-sonnet-4) targets a specific Foundation Model — swap it for any serving endpoint in your workspace
• CONCAT builds the prompt from a column value, giving the model per-row context

More Patterns

Parse structured output in PySpark

“Write PySpark code that parses the JSON response from ai_query into a strongly typed struct column.”

from pyspark.sql.functions import from_json, col

df = df.withColumn("invoice", from_json(
    col("ai_response.response"),
    "STRUCT<number:STRING, total:DOUBLE, items:ARRAY<STRUCT<code:STRING, description:STRING>>>"
))

The response field contains the raw JSON string. from_json maps it to a Spark struct so you can access nested fields like invoice.items[0].code in downstream queries.

Analyze images with multimodal models

“Write a SQL query that reads images from a Unity Catalog volume and generates descriptions using a vision model.”

SELECT path, ai_query(
    'databricks-llama-4-maverick',
    'Describe what is in this image.',
    files => content
) AS description
FROM read_files('/Volumes/catalog/schema/images/', format => 'binaryFile');

The files parameter accepts binary content from read_files. Currently supports JPEG and PNG formats. The model receives both the text prompt and the image, so you can ask specific questions about the visual content.

Handle errors gracefully in batch pipelines

“Write a SQL query that classifies documents in batch while capturing individual row errors in a separate column.”

SELECT id,
  ai_response.response AS classification,
  ai_response.error AS error_message
FROM (
    SELECT id, ai_query(
        'databricks-claude-sonnet-4',
        CONCAT('Classify this document: ', text),
        failOnError => false
    ) AS ai_response
    FROM catalog.schema.documents
);

With failOnError => false, failed rows get a NULL response and a populated error field. Write these to a sidecar error table for retry or manual review instead of losing the entire batch.

Wrap reusable logic in a SQL UDF

“Write a SQL UDF that encapsulates invoice extraction logic so any query can call it by name.”

CREATE FUNCTION catalog.schema.extract_invoice(text STRING)
RETURNS STRING
RETURN ai_query(
    'databricks-claude-sonnet-4',
    CONCAT('Extract invoice JSON from: ', text),
    responseFormat => '{"type":"json_object"}'
);

Now any query can call catalog.schema.extract_invoice(some_column) without repeating the prompt or model configuration. When the model or prompt changes, you update one function definition.

Watch Out For

• Skipping failOnError => false in batch jobs — a single malformed row kills the entire query. Always use false for production batch inference and route errors to a sidecar table.
• Passing raw binary to ai_query for text extraction — parse documents with ai_parse_document first, then feed the extracted text to ai_query. Raw binary content produces garbage output.
• Exceeding context window limits — long text columns silently truncate or fail. Use LEFT(text, 6000) to cap input length before passing to ai_query.
• Using ai_query when a task-specific function exists — ai_classify, ai_extract, and the other built-in functions are faster and cheaper because they use optimized endpoints. Reach for ai_query only when the task-specific functions don’t fit.