Skip to main content

AI Function tasks for jobs

Beta

This feature is in Beta. To use it, a workspace admin must turn on Lakeflow AI Function Task from the Previews page. See Manage Databricks previews.

Use an AI Function task to run an AI function on every file in a Unity Catalog volume or every row of a table as part of a job. The task writes one result per file or row to a Unity Catalog table that it creates and manages. Each run processes only the files or rows that were added, changed, or deleted since the previous run. In the Jobs API and Declarative Automation Bundles, an AI Function task is an ai_function_task, and its function setting determines its type: ai_parse_document, ai_extract, or ai_classify.

AI Function task types​

Each task type runs one AI function and is named for it, for example an ai_parse_document task. In the Jobs UI, the task types appear under Document intelligence. To add a task, see the page for its task type.

Intelligent document processing​

Task type

Jobs UI name

Input

Description

ai_parse_document

Document Parsing

Files in a volume, or a FILE or BINARY table column

Parse structured content (text, tables, figure descriptions) and layout from unstructured documents.

ai_extract

Information Extraction

A text or VARIANT table column, such as the response column of an ai_parse_document task

Extract structured fields from documents or text using a schema you define.

ai_classify

Classification

A text or VARIANT table column, such as the response column of an ai_parse_document task

Classify input text according to labels you provide.

Task type

Jobs UI name

Input

Description

ai_parse_document

Document Parsing

Files in a volume, or a FILE or BINARY table column

Parse structured content (text, tables, figure descriptions) and layout from unstructured documents.

ai_extract

Information Extraction

A text or VARIANT table column, such as the response column of an ai_parse_document task

Extract structured fields from documents or text using a schema you define.

ai_classify

Classification

A text or VARIANT table column, such as the response column of an ai_parse_document task

Classify input text according to labels you provide.

For example, you can use AI Function tasks to:

  • Parse new contracts that are added to a volume each night.
  • Extract invoice fields from documents that an ai_parse_document task parsed.
  • Classify incoming support tickets by urgency before a later task routes them.

Requirements​

  • If Document intelligence doesn't appear in the task list, a workspace admin must turn on the Lakeflow AI Function Task preview.

  • The AI functions must be available in your workspace region. See Features with limited regional availability.

  • The output schema must exist before the task runs.

  • The job's Run as user must have the following privileges:

    • USE CATALOG and USE SCHEMA on the catalogs and schemas of the input and output.
    • SELECT on the input table, or READ VOLUME on the input volume.
    • CREATE TABLE on the output schema, for the run that creates the output table.
    • MODIFY on the output table, for later runs.
    • EXECUTE on the function in system.ai, if your organization restricts access to AI Functions. See AI Functions Unity Catalog permissions.
  • An input table must have change data feed enabled. To enable it, run the following:

    SQL
    ALTER TABLE <table-name> SET TBLPROPERTIES (delta.enableChangeDataFeed = true);
  • If you don't add an ID column and the input column isn't FILE type, an input table must also have row tracking enabled. To enable it, run the following:

    SQL
    ALTER TABLE <table-name> SET TBLPROPERTIES (delta.enableRowTracking = true);

Input​

Under Input, select a Unity Catalog volume or a folder in a volume to run the function on each file. Or select a table and the Column that holds the content to run the function on each row. The supported sources and column types depend on the task type. See AI Function task types.

For document inputs, Databricks recommends the following:

  • Store documents in a table as a FILE (Beta) column rather than a BINARY column.
  • For a volume input, turn on file events for the volume's external location. File events aren't required, but the task finds new and changed files faster with them.

ID column​

For a table input, you can add an ID column, such as a primary key. To add it, click Add next to ID column. The task uses the ID column to match results to source rows across runs and copies it to the output table. Without an ID column, the task identifies rows by their row tracking ID. For a FILE column, the task identifies rows by the file's URI instead, so row tracking isn't required, and rows that reference the same file share one result.

Output​

Under Output, select the Location: the Unity Catalog catalog and schema for the output table. If you don't enter a Table name, the task generates one that starts with parsed, extracted, or classified and stays the same across runs.

The task creates and manages its output table. A run fails if a table with that name already exists and the task didn't create it.

Use the output in a later task​

A later task can read an AI Function task's output table with the dynamic value reference {{tasks.<task_name>.output.table_name}}.

When you add an AI Function task after another AI Function task, its input defaults to the earlier task's output table and response column. For example, an ai_parse_document task followed by an ai_extract task extracts fields from the parsed documents. See Use with other AI functions.

Define AI Function tasks in code​

You can define AI Function tasks in Declarative Automation Bundles, the Databricks CLI, or the Jobs API. Add an ai_function_task to the task, with one input (volume or delta_table), the output location, and one function (ai_parse_document, ai_extract, or ai_classify). If you leave out output.table, the task generates the table name. For an example, see the page for each task type.

How AI Function tasks work​

Incremental processing​

On each run, the task processes files and rows that are new or changed since the previous run, and deletes the results of files and rows that were deleted. Changes to the task's settings, such as the fields to extract, the labels, or instructions, apply only to inputs processed after the change. To apply them to all inputs, enter a new Table name. If you change the task's input or output table, the next run processes all inputs again, which adds cost. See View costs for AI Function workloads.

Output table​

The output table is a streaming table that a later task can read incrementally. It contains the following columns:

  • ID column or id: Identifies the input, such as the path of a file in a volume.
  • response: The function's full result as a VARIANT, in the format that the function returns. Optional values, such as confidence scores, rationales, and citations, are part of this result. The task doesn't add a column for each value. To get a value as a column, select it from response in a query.
  • error: The error for this input, or NULL when the request succeeds.

To combine results with source data, join on the ID column.

Monitor a run​

While the task runs, the task graph shows the percentage of inputs processed. The run output shows the number of Rows, Succeeded, and Failed inputs, and the estimated time remaining (ETA). After the run, Sample Data shows rows from the output table.

Limitations​

  • A task can have only one active run. If you start a run while another run of the same task is in progress, the new run fails.
  • AI Function task runs in a workspace share the workspace's AI Functions capacity, so runs take longer when many run at once. For best performance, keep about 20 or fewer runs active at a time.

Additional resources​