ai_query in Lakeflow Pipelines: Enrichment and Classification at Scale

1.Where inference sits in a batch pipeline

For large-scale batch inference, you can put AI Functions straight into production pipelines, either task-specific functions or ai_query. You can write them in SQL or Python and deploy them as Lakeflow pipelines, as scheduled Databricks workflows, or as Structured Streaming jobs. Inference is one step in a longer flow, with ingestion before it and post-processing after it.

The Databricks reference pipeline works on news articles. A streaming table, news_raw, ingests CSV files from a Unity Catalog volume. A materialized view, news_categorized, then calls ai_query on each article to extract a title and a category.

@dp.materialized_view(
  comment="Extract category and title from news articles using LLM inference."
)
def news_categorized():
  # Limit the number of rows to 2 as in the SQL version
  df_raw = spark.read.table("news_raw").limit(2)
  # Inject the JSON schema variable into the ai_query call using an f-string.
  return df_raw.withColumn(
    "meta_data",
    expr(
      f"ai_query('system.ai.meta-llama-3-3-70b-instruct', "
      f"concat('Extract the category of the following news article: ', inputs), "
      f"responseFormat => '{news_extraction_schema}')"
    )
  )

In Python, ai_query is a SQL expression wrapped in expr() inside withColumn. Nothing needs installing: the same SQL function runs from a DataFrame.

Sources12

2.Classify with a schema, validate, then enrich

The news_extraction_schema variable is a json_schema responseFormat with two properties. title is a free string. category is a string limited to an enum of six labels. That enum turns open-ended generation into classification: the model has to pick one of your labels. Because responseFormat returns a JSON string, later stages read individual fields with get_json_object(meta_data, '$.category').

Model output is still data you didn't write, so the next table checks it before anything relies on it. news_validated declares expectations on the parsed fields. The title must have at least three words and the category must be one of the six labels. In the Python version the category rule uses @dp.expect_or_fail. How expectation actions behave is covered in the data quality lesson. What matters here is the pattern: check the inference output before the next step uses it.

CREATE OR REFRESH MATERIALIZED VIEW news_validated (
  CONSTRAINT valid_title EXPECT (size(split(get_json_object(meta_data, '$.title'), ' ')) >= 3),
  CONSTRAINT valid_category EXPECT (get_json_object(meta_data, '$.category') IN ('Politics', 'Sports', 'Technology', 'Health', 'Entertainment', 'Business'))
)
COMMENT "Validated news articles ensuring the title has at least 3 words and the category is valid."
AS
SELECT *
FROM news_categorized;

The last stage, news_summarized, reads category and title out of meta_data and calls ai_query again with a plain prompt, "Summarize the following political news article in 2-3 sentences". One pipeline can chain several inference steps, and each adds columns to rows that passed the step before.

When there is no schema to enforce, a few-shot prompt is another way to classify. The scheduled-workflow example tells the model it is an opinion-mining service, lists the allowed classifications (Quality, Service, Design, Safety, Efficiency, Usability, Price) and includes one worked DOCUMENT/RESULT pair. It then appends each REVIEW_TEXT.

Sources2

3.When ai_classify is the better tool

If the job is just "put this text in one of these buckets," Databricks' advice to start with a task-specific function points to ai_classify. You supply labels, either as a JSON array of names or as a JSON object mapping each name to a description. Descriptions improve accuracy. You need at least 2 labels and no more than 500. The labels argument can be any expression that evaluates to a STRING, including a Delta table column, so the taxonomy can live in a table. The content can be raw text or a VARIANT produced by an upstream AI function such as ai_parse_document.

Databricks recommends version 2.1. It adds enableConfidenceScores and enableRationales, and it changes the output shape. Each returned classification becomes an object with a value key instead of a plain label string. That change applies to every 2.1 call, even with the new options turned off, so moving from 2.0 to 2.1 breaks downstream code that reads the label directly. Pin the version with options => map('version', '2.1').

Sources3

4.Running inference at production scale

By default one bad row can sink a large batch. failOnError is an optional boolean that defaults to true, and it needs Databricks Runtime 15.3 or above. With the default, ai_query returns the parsed response and a failed row fails the query. Set it to false and ai_query returns a STRUCT holding the response and an error message. A row that succeeds has a null errorMessage. A row that fails because of a model endpoint error has a null response. Errors of any other kind still fail the whole query.

SELECT text, ai_query(
    "system.ai.meta-llama-3-3-70b-instruct",
    "Summarize the given text comprehensively, covering key points and main ideas concisely while retaining relevant details and examples. Ensure clarity and accuracy without unnecessary repetition or omissions: " || text,
 ?  => false
) AS summary
FROM uc_catalog.schema.table;

Databricks lists further practices for large workloads. Use Databricks-hosted system.ai model services rather than provisioned throughput endpoints, because they scale automatically. Pick a model that is optimized for batch inference, since other models can cut throughput. Submit the whole dataset in one query: AI Functions already handles parallelization, retries and scaling, and splitting the data into small batches by hand can lower throughput. To see progress, open the query profile and click AI Query. It shows how many inferences completed and how many failed.

Sources41

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.With failOnError => false, ai_query never fails the query.Why is that wrong?

   Only model endpoint errors become per-row error messages. Errors of any other kind still fail the whole query.

   Covered in Running inference at production scale

2. 2.Splitting a large table into small manual batches makes ai_query faster and safer.Why is that wrong?

   AI Functions handles parallelization, retries and scaling itself. Databricks says to submit the full dataset in one query.

   Covered in Running inference at production scale

3. 3.Upgrading ai_classify from version 2.0 to 2.1 is safe as long as you don't turn on confidence scores or rationales.Why is that wrong?

   The 2.1 output shape, an object with a value key for each label, applies to every 2.1 call. The upgrade is not output-compatible.

   Covered in When ai_classify is the better tool

Practise it for real

Enrich a sample table with ai_query, then make the query robust and tune it the way a production pipeline would

1. 1.On a non-Classic SQL warehouse or Databricks Runtime 15.4 LTS+, run the documented query that calls ai_query with 'system.ai.meta-llama-3-3-70b-instruct' and a prompt joined to pickup_zip, against samples.nyctaxi.trips with LIMIT 10.

   Why: It confirms that a Databricks-hosted model works with no endpoint setup.

   You should see: Ten trip rows, each with an extra column holding the model's answer about the state for that ZIP code.

2. 2.Add failOnError => false as a third argument to the same call.

   Why: Production batch jobs should keep the rows that succeed when some inferences fail.

   You should see: The new column is now a STRUCT with a response field and an errorMessage field, and errorMessage is null for successful rows.

3. 3.Replace failOnError with modelParameters => named_struct('max_tokens', 100, 'temperature', 0.7).

   Why: It shows that model settings are constants passed through modelParameters, not per-row values.

   You should see: The query runs and the answers may vary more, since temperature is now above the 0.0 default.

4. 4.Open the query profile from the results window and click AI Query.

   Why: This is where you monitor batch inference progress.

   You should see: Counts of completed and failed inferences, plus the total time the request took.