databricks-ai-functions

Files (7)

• SKILL.md
• openai.yaml · agents/
• databricks.png · assets/
• databricks.svg · assets/
• 1-task-functions.md · references/
• 2-ai-query.md · references/
• 3-ai-forecast.md · references/

SKILL.md (first lines)

# Databricks AI Functions

> **Official Docs:** https://docs.databricks.com/large-language-models/ai-functions
> Individual function reference: https://docs.databricks.com/sql/language-manual/functions/

## Overview

Databricks AI Functions are built-in SQL and PySpark functions that call Foundation Model APIs directly from your data pipelines — no model endpoint setup, no API keys, no boilerplate. They operate on table columns as naturally as `UPPER()` or `LENGTH()`, and are optimized for batch inference at scale.

**Always prefer a task-specific function over `ai_query`.** Reach for `ai_query` only when no task function fits (custom/external endpoints, multimodal, or JSON beyond `ai_extract`'s limits). Every function below shares a baseline: **DBR 15.1+** (notebooks) / **15.4 ML LTS** (batch), **not on SQL Warehouse Classic**, and region must support AI Functions — the Prereqs column lists only what's *additional*.

**Cost & speed — each call is an LLM inference (slow and billed per token).** Run a function **once** per row and persist the result to a Delta table; never re-invoke it on every downstream query. In demos, **avoid generating tables with millions of rows** — sample the input when needed so the demo runs quickly. Materialize once, then query the cheap Delta output.

The **Function** column links to the in-repo deep reference (full options, schemas, examples); **Docs** links to the official page.

| Function | Task | Input | Output | Extra prereqs | Docs |
|---|---|---|---|---|---|
| [`ai_analyze_sentiment`](references/1-task-functions.md#ai_analyze_sentiment) | Sentiment scoring | `content STRING` | `STRING` — `positive`/`negative`/`neutral`/`mixed`, or `NULL` | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_analyze_sentiment) |
| [`ai_classify`](references/1-task-functions.md#ai_classify) | Fixed-label routing | `content STRING\|VARIANT`, `labels` (2–500), `[options MAP]` | `VARIANT` — `{response:[label], error_message}` | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_classify) |
| [`ai_extract`](references/1-task-functions.md#ai_extract) | Entity / field extraction | `content STRING\|VARIANT`, `schema STRING` (JSON), `[options MAP]` | `VARIANT` — `{response:{…}, error_message, metadata}` | ≤256 fields, ≤12 nesting levels | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_extract) |
| [`ai_fix_grammar`](references/1-task-functions.md#ai_fix_grammar) | Grammar correction | `content STRING` | `STRING` (corrected) | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_fix_grammar) |
| [`ai_gen`](references/1-task-functions.md#ai_gen) | Free-form generation | `prompt STRING` | `STRING` | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_gen) |
| [`ai_mask`](references/1-task-functions.md#ai_mask) | PII redaction | `content STRING`, `labels ARRAY<STRING>` | `STRING` (entities → `[MASKED]`) | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_mask) |
| [`ai_similarity`](references/1-task-functions.md#ai_similarity) | Semantic similarity | `expr1 STRING`, `expr2 STRING` | `FLOAT` (0.0–1.0) | — | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_similarity) |
| [`ai_summarize`](references/1-task-functions.md#ai_summarize) | Summarization | `content STRING`, `[max_words INT]` (0 = uncapped) | `STRING` | Public Preview; English-tuned | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_summarize) |
| [`ai_translate`](references/1-task-functions.md#ai_translate) | Translation | `content STRING`, `to_lang STRING` | `STRING` | Langs: en, fr, de, hi, it, pt, es, th | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_translate) |
| [`ai_parse_document`](references/1-task-functions.md#ai_parse_document) | Parse PDF / Office / images | `content BINARY`, `[Map('version','2.0', …)]` | `VARIANT` — pages, elements, error_status | **DBR 17.3+**; ≤500 pages / 100 MB | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_parse_document) |
| [`ai_prep_search`](references/1-task-functions.md#ai_prep_search) | RAG chunking from parsed docs | `parsed VARIANT`, `[options MAP]` | `VARIANT` — `{document:{contents, pages, source_uri}, error_status}` | **DBR 18.2+** (serverless env v3+) | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_prep_search) |
| [`ai_query`](references/2-ai-query.md) | Any serving endpoint (built-in foundation or custom), multimodal, complex JSON (**last resort**) | `endpoint STRING`, `request STRING\|STRUCT`, `[returnType]`, `[failOnError BOOL]`, `[modelParameters STRUCT]`, `[responseFormat STRING]`, `[files]` | Parsed response; with `failOnError => false` a `STRUCT{response, errorMessage}` | **Pro/Serverless** warehouse; `CAN QUERY` on endpoint | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_query) |
| [`ai_forecast`](references/3-ai-forecast.md) | Time series forecasting (table-valued) | `observed TABLE`, `horizon`, `time_col`, `value_col`, `[group_col]`, `[prediction_interval_width]`, `[frequency]`, `[seed]`, `[parameters]` | Rows: time/group cols + per value `{v}_forecast`, `{v}_upper`, `{v}_lower` (DOUBLE) | **Pro/Serverless** warehouse; Public Preview | [↗](https://docs.databricks.com/sql/language-manual/functions/ai_forecast) |

Models run under Apache 2.0 or LLAMA 3.3 Community License — you are responsible for compliance.

## Patterns

**Chain task functions to enrich a column in one pass.** `ai_classify`/`ai_extract` return a VARIANT — read it with the colon operator (`:response`):

```sql
SELECT id,
  ai_analyze_sentiment(content)                                                   AS sentiment,
  ai_summarize(content, 30)                                                       AS summary,
  ai_classify(content, '["technical","billing","other"]', map('version','2.0')):response[0]::STRING AS category,
  ai_extract(content, '["product","error_code","date"]', map('version','2.0')):response:product::STRING AS product,
  ai_fix_grammar(content)                                                         AS content_clean
FROM raw_feedback;
```

In **PySpark**, call any of these inside `expr(...)`: `df.withColumn("category", expr("ai_classify(content, '[\"a\",\"b\"]', map('version','2.0')):response[0]::STRING"))` — and read VARIANT fields via `selectExpr("col:response:field::STRING AS field")`.

**PII redaction before storage** — `ai_mask(content, ARRAY(entity_types))` returns text with entities → `[MASKED]`.

```sql
SELECT ai_mask(message, array('person','email','phone','address')) AS message_safe FROM raw_messages;
```

**Semantic matching / dedup** — `ai_similarity` returns 0–1; self-join and threshold:

```sql
SELECT a.id, b.id, ai_similarity(a.name, b.name) AS score
FROM companies a JOIN companies b ON a.id < b.id
WHERE ai_similarity(a.name, b.name) > 0.85;
```

**Forecasting** — table-valued; one row per future period (+ per group). Full param/group/interval forms → [3-ai-forecast.md](references/3-ai-forecast.md):

```sql
SELECT * FROM ai_forecast(
    observed => TABLE(SELECT date, sales FROM daily_sales),
    horizon => '2026-12-31', time_col => 'date', value_col => 'sales');
-- Returns: date, sales_forecast, sales_upper, sales_lower
```

**Nested JSON via `ai_query`** (last resort — only past `ai_extract`'s limits) — parse the response with `from_json`. Model names, multimodal `files =>`, `modelParameters`, SQL UDF → [2-ai-query.md](references/2-ai-query.md):

```sql
SELECT from_json(
    ai_query('databricks-claude-sonnet-4',
        concat('Extract invoice as JSON with nested line_items array: ', text_blocks),
        responseFormat => '{"type":"json_object"}', failOnError => false).response,
    'STRUCT<numero:STRING, total:DOUBLE, line_items:ARRAY<STRUCT<code:STRING, qty:DOUBLE>>>'
) AS invoice
FROM parsed_documents;
```

Document parsing (`ai_parse_document`) and RAG chunking (`ai_prep_search`) get their own staged pipeline below.

## Document Processing Pipeline

Chain AI Functions stage-by-stage into Delta tables for batch document processing. The example is written as a **Spark Declarative Pipeline (SDP / Lakeflow / DLT)** — `CREATE OR REFRESH STREAMING TABLE` with `STREAM(...)` sources. To run the same logic standalone in a **notebook / SQL warehouse**, swap each `CREATE OR REFRESH STREAMING TABLE x AS` for `CREATE OR REPLACE TABLE x AS` and drop the `STREAM(...)` wrappers. In SDP Python it's `@dp.table` with `from pyspark import pipelines as dp`.

```sql
-- Stage 1 — parse binary docs (any type), filter parse errors
CREATE OR REFRESH STREAMING TABLE raw_parsed AS
SELECT path,
  concat_ws('\n', transform(parsed:document:elements, e -> e:content::STRING)) AS text_blocks,
  parsed:error_status AS parse_error
FROM (
  SELECT path, ai_parse_document(content, map('version','2.0')) AS parsed
  FROM STREAM read_files('/Volumes/my_catalog/doc_processing/landing/', format => 'binaryFile')
)
WHERE parsed:error_status IS NULL;

-- Stage 2 — classify document type (cheap, no endpoint selection)
CREATE OR REFRESH STREAMING TABLE classified_docs AS
SELECT *,
  ai_classify(text_blocks, '["invoice","purchase_order","receipt","contract","other"]', map('version','2.0')):response[0]::STRING AS doc_type
FROM STREAM raw_parsed;

-- Stage 3 — extract fields; ai_extract returns a VARIANT, read fields with `:`
CREATE OR REFRESH STREAMING TABLE extracted AS
SELECT path, doc_type,
  result:response:invoice_number::STRING AS invoice_number,
  result:response:vendor_name::STRING    AS vendor_name,
  result:response:total_amount::DOUBLE   AS total_amount,
  result:error_message::STRING           AS extract_error
FROM (
  SELECT *, ai_extract(text_blocks,
    '{"invoice_number":{"type":"string"},"vendor_name":{"type":"string"},"total_amount":{"type":"number"}}',
    map('version','2.0')) AS result

From ThomazRossito/ai-data-agents, MIT.

More from ThomazRossito/ai-data-agents