ai_query() in Databricks: Running an LLM From Plain SQL

How to classify free text with one SQL function, what happens under the hood, and three mistakes that cost real money.

ai_query() in Databricks: Running an LLM From Plain SQL

Every data team has a column nobody can group by.

It is usually free text. Support tickets. Product reviews. Notes typed by a sales rep at the end of a long day. Thousands of rows, all useful, none of them sorted.

For years, the fix was a Python script that called a model API one row at a time. It worked on a laptop. It broke at scale.

Databricks now gives you a simpler path. A SQL function called ai_query() sends each row to a model and returns the answer as a new column. This article explains what it does, how it behaves under the hood, and the three mistakes that make it slow or expensive.

The old way, and why it hurt

The usual pattern looked like this. Write a Python UDF. Inside it, call an HTTP endpoint. Apply the UDF to a DataFrame.

On ten rows, it is fine. On a million rows, three things go wrong.

First, every row becomes its own network call, and Spark has no idea those calls are happening. It cannot batch them or retry them sensibly.

Second, the model provider has rate limits. Fifty Spark tasks hitting the same API at once trigger errors, and one failed task can fail the whole job.

Third, the API key has to live somewhere in your code or job settings. That is a security problem waiting to happen.

What ai_query() does instead

ai_query() is a built-in function. You give it the name of a model serving endpoint and a prompt. It gives you back text.

SELECT
  ticket_id,
  ai_query(
    'databricks-meta-llama-3-3-70b-instruct',
    CONCAT('Classify this support ticket as billing, bug, or other. Reply with one word. Ticket: ', ticket_text)
  ) AS category
FROM support_tickets

Because it is part of the engine, Databricks handles the hard parts for you. Requests are sent in batches. Retries and rate limits are managed by the platform. Access goes through workspace permissions, so there is no API key in your query.

flowchart LR
  A[Delta table rows] --> B[Spark tasks]
  B --> C[Batched requests]
  C --> D[Model serving endpoint]
  D --> E[Answers as a new column]
  E --> F[Silver Delta table]

The mental model is simple. It behaves like any other column expression, except the work happens on a model instead of a CPU.

A small example you can try

Create a tiny table first. Run this in a SQL cell.

CREATE OR REPLACE TABLE workspace.default.support_tickets AS
SELECT * FROM VALUES
  (1, 'I was charged twice for my plan this month'),
  (2, 'The export button does nothing when I click it'),
  (3, 'Can I change the email on my account?'),
  (4, 'App crashes when I upload a large file')
AS t(ticket_id, ticket_text)

Now ask a model to classify each ticket.

SELECT
  ticket_id,
  ticket_text,
  ai_query(
    'databricks-meta-llama-3-3-70b-instruct',
    CONCAT('Classify this support ticket as billing, bug, or other. Reply with one lowercase word only. Ticket: ', ticket_text)
  ) AS category
FROM workspace.default.support_tickets

You should see something like this:

ticket_idcategory
1billing
2bug
3other
4bug

A note on availability. ai_query() needs a model serving endpoint. Open Serving in your workspace sidebar to see which endpoints you can use, and swap the name in the query if yours is different. Model availability on Free Edition can change and comes with usage limits, so if no endpoint is listed, read this section as a concept for now.

Getting structured answers back

Free text answers are hard to trust. A model might reply "Billing." with a capital letter and a full stop, and your GROUP BY now has two buckets.

You can ask for a fixed shape instead with the responseFormat argument.

SELECT
  ticket_id,
  ai_query(
    'databricks-meta-llama-3-3-70b-instruct',
    CONCAT('Classify this ticket and rate urgency from 1 to 5. Ticket: ', ticket_text),
    responseFormat => 'STRUCT<result: STRUCT<category: STRING, urgency: INT>>'
  ) AS triage
FROM workspace.default.support_tickets

Now the answer is a struct, not a sentence. You can select triage.result.category and filter on triage.result.urgency like any normal column.

Three mistakes that cost real money

Calling the model before filtering. Every row you send costs tokens and time. Drop spam, empty text, and duplicates first. A WHERE clause is free. A model call is not.

Running it again and again on the same rows. A query that calls ai_query() will call the model every time it runs. Run it once, save the result to a table, and query that table afterwards.

CREATE OR REPLACE TABLE workspace.default.support_tickets_triaged AS
SELECT
  ticket_id,
  ticket_text,
  ai_query(
    'databricks-meta-llama-3-3-70b-instruct',
    CONCAT('Classify this support ticket as billing, bug, or other. Reply with one lowercase word only. Ticket: ', ticket_text)
  ) AS category
FROM workspace.default.support_tickets
WHERE ticket_text IS NOT NULL

For a growing table, only send new rows. Process the rows that are not already in the triaged table, then append them. The same incremental thinking you use for any pipeline applies here.

Trusting every answer. Models are usually right, not always right. Check the output with a simple query, such as counting categories that are not in your allowed list. Treat odd values like any other data quality issue.

Where it fits in a lakehouse

Keep raw text untouched in Bronze. That is your record of what really arrived.

Add AI enrichment on the way to Silver. That is where data is cleaned and shaped, and a category or urgency score is just another derived column.

Gold tables then stay simple. A dashboard of tickets by category no longer needs to know a model was ever involved.

The takeaway

ai_query() turns a model into a column expression. That is the whole idea.

It removes the fragile parts of the old approach: per-row API calls, unmanaged rate limits, and keys in code. What it does not remove is good engineering. Filter first, save results, process only new rows, and check what comes back.

Continue learning