Metadata Filtering for Databricks RAG Using ai_extract and ai_prep_search

A Databricks pipeline built on the SQL AI function ai_prep_search already does a lot. Parsed document output goes in, semantic and context-enriched chunks come out, and they're ready for Vector Search. The flow is tidy: parse the documents, prepare chunks, flatten them into a Delta table and index it.

Yet something is missing when that pipeline runs over a document corpus, here a set of Orion Robotics manuals. The chunks themselves are good. The chunk_to_embed field carries proper document context, and chunk_to_retrieve is clean enough to hand to an LLM. The trouble starts at query time, because there's no dependable way to filter. A request like "find me the firmware guide for the Orion Motion Controller" has to be satisfied by semantic similarity alone, since the index has no structured columns such as document_type, product_name or target_audience.

For agent-driven retrieval that's a genuine gap. Pure semantic search holds up in a demo. In production, an agent has to shrink the search space before vector similarity runs, and metadata filters do exactly that. The fix is adding ai_extract v2.1 to the pipeline.

Why metadata filters matter

Hybrid search in RAG usually means vector search combined with keyword matching such as BM25. Useful, but for structured document collections another kind of hybrid matters more: vector search plus metadata filtering.

Picture an agent that already knows the user needs a manual for the Orion Arm. Ranking 10,000 chunks from every kind of document would be wasteful. A better plan is to trim the data to the relevant documents first and let vector similarity order the best matches within that smaller pool. Databricks Vector Search allows this via filter columns declared at index creation.

The catch is that the metadata must be clean, consistent and correctly typed. If document_type reads "User Manual" in one row, "manual" in another and nothing in a third, filters misfire. Retrieval quality then looks bad, and hours go into debugging it while the real culprit is inconsistent metadata created earlier.

ai_extract v2.1 addresses this. It consumes the structured VARIANT produced by ai_parse_document and returns typed metadata checked against a schema. Enum fields pin down fixed values, string fields steer free-text extraction and arrays cover multi-valued fields. Since extraction follows a schema, output stays uniform across documents even when their layouts differ.

Three functions, three jobs

The roles of the three functions don't overlap, so it helps to separate them before looking at the notebook.

1. ai_parse_document v2.0 turns raw binaries (PDF, DOCX, images) into a structured VARIANT, identifying text blocks, headings, tables, figures, bounding boxes and the section hierarchy. It knows what's inside a document structurally and has no idea what the document is about.
2. ai_extract v2.1 takes the parsed VARIANT and, following a schema the user supplies, pulls out document-level metadata. It understands the subject only because the schema says what to look for. Its result is a new VARIANT with typed fields such as document_title, document_type, product_name and target_audience, plus any others defined. Version 2.1 introduced confidence scores and citations and is the release to use.
3. ai_prep_search splits the parsed VARIANT into search-ready semantic chunks, using the structural cues from parsing (bounding boxes, element types, section hierarchy). It leaves document-level metadata alone, since that belongs to ai_extract.

Every function has one responsibility, and the pipeline composes them.

Step by step in the notebook

The notebook works on the Orion Robotics corpus, which sits in a Unity Catalog Volume and mixes PDFs, DOCX files and images. Parsing is first:

SELECT
  path,
  ai_parse_document(
    content,
    MAP('version', '2.0')
  ) AS parsed
FROM READ_FILES(
  '/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/',
  format => 'binaryFile'
)

Each file in the volume is read as binary and converted to a structured VARIANT. Pinning version 2.0 explicitly is wise, because that schema includes the bounding box coordinates and figure descriptions ai_prep_search depends on for chunking decisions. Staying on the default would let Databricks change behavior underneath the pipeline.

The sample run produced this result:

The next query extracts metadata from the parsed output:

WITH parsed_documents AS (
  SELECT path, ai_parse_document(content, MAP('version', '2.0')) AS parsed
  FROM READ_FILES('/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/', format => 'binaryFile')
)
SELECT
  path,
  ai_extract(
    parsed,
    '{
      "document_title":  {"type": "string", "description": "Main title or heading of the document"},
      "document_type":   {"type": "enum", "labels": ["manual", "datasheet", "guide", "report", "specification", "policy", "other"]},
      "product_name":    {"type": "string", "description": "Primary product or system referenced"},
      "version":         {"type": "string", "description": "Document or product version if mentioned"},
      "key_topics":      {"type": "array", "items": {"type": "string"}},
      "target_audience": {"type": "enum", "labels": ["engineer", "operator", "manager", "customer", "general"]}
    }',
    MAP(
      'version', '2.1',
      'instructions', 'These are technical documents from Orion Robotics. Extract metadata that would help users filter and find relevant documents.'
    )
  ) AS metadata
FROM parsed_documents

Some choices need explaining. Both document_type and target_audience are enums, since filters need a closed set of values. As free text they might come back as "manual", "User Manual" or "technical manual", which breaks matching. The instructions option supplies background, here that the files are technical documents from Orion Robotics, so extraction improves. Also, ai_extract accepts the VARIANT from ai_parse_document as is, with no plain-text conversion.

Output from the sample run:

Chunking comes next:

WITH parsed_documents AS ( ... )
SELECT
  path,
  ai_prep_search(parsed) AS chunked
FROM parsed_documents

ai_prep_search consumes the identical parsed VARIANT. Documents are split into chunks stored in a document.contents array, and each chunk exposes chunk_id, chunk_position, chunk_to_retrieve and chunk_to_embed. The last two are intentionally not the same text.

chunk_to_embed carries extra context like the document title, section header and page details so the embedding model can place the chunk. chunk_to_retrieve is cleaner text, normally the version shown to users or passed to the LLM.

Splitting them is a sound design. Many RAG pipelines embed and answer from the same text, yet context-rich text suits embeddings while clean text suits reading, so keeping both is better.

Sample output:

Joining it all together

The combined query is where the pieces meet:

WITH parsed_documents AS (
  SELECT path, ai_parse_document(content, MAP('version', '2.0')) AS parsed
  FROM READ_FILES('/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/', format => 'binaryFile')
),
doc_metadata AS (
  SELECT
    path,
    parsed,
    ai_extract(
      parsed,
      '{ ... }',
      MAP('version', '2.1', 'instructions', 'These are technical documents from Orion Robotics. Extract metadata for search filtering.')
    ) AS metadata
  FROM parsed_documents
),
doc_chunks AS (
  SELECT path, ai_prep_search(parsed) AS chunked
  FROM doc_metadata
)
SELECT
  chunk.value:chunk_id::STRING          AS chunk_id,
  chunk.value:chunk_position::INT       AS chunk_position,
  chunk.value:chunk_to_retrieve::STRING AS chunk_to_retrieve,
  chunk.value:chunk_to_embed::STRING    AS chunk_to_embed,
  doc_chunks.chunked:document.source_uri::STRING AS source_uri,
  dm.metadata:response:document_title::STRING    AS document_title,
  dm.metadata:response:document_type::STRING     AS document_type,
  dm.metadata:response:product_name::STRING      AS product_name,
  dm.metadata:response:version::STRING           AS doc_version,
  dm.metadata:response:target_audience::STRING   AS target_audience
FROM
  doc_chunks
  JOIN doc_metadata dm ON doc_chunks.path = dm.path,
  LATERAL variant_explode(doc_chunks.chunked:document.contents) AS chunk

The LATERAL variant_explode step expands the nested chunk array into rows. The JOIN on path attaches document-level metadata to each chunk of that document. The outcome is one row per chunk, carrying both the retrieval fields and the filterable metadata columns. That's the shape stored in the Delta table and used to feed the Vector Search index.

Result from the sample run:

The table gets materialized like this:

CREATE OR REPLACE TABLE <YOUR_CATALOG>.<YOUR_SCHEMA>.orion_docs_rag_chunks AS
-- [full CTE chain from Step 4]

With the table in place, a Mosaic AI Vector Search Delta Sync index is created with these settings:

1. Embedding source column: chunk_to_embed
2. Primary key: chunk_id
3. Filterable fields: doc_version plus document_type, product_name and target_audience

The chunk_id is stable from run to run, composed of the document ID and the chunk's position. Because of that, incremental MERGE updates to the Delta table leave unchanged chunks alone instead of re-indexing them, a real saving when the corpus is refreshed on a schedule.

The payoff at query time

The benefit appears when a user asks something, and not at load time. Suppose the request is "Find the firmware guide for the Orion Motion Controller". Retrieval can now proceed in two moves:

1. Work out the matching metadata filters, such as document_type = "guide" and product_name = "Orion Motion Controller".
2. Embed the query and search only among documents that pass those filters.

Lacking ai_extract, an agent depends on semantic similarity alone. For a few similar documents that may be fine. Across thousands of files covering many products, types and audiences it turns unreliable. Enum-based metadata fixes it: the agent needn't guess between a manual and a guide and simply filters on document_type = "manual" for exact hits.

Limits and costs

ai_prep_search gives no control over chunking. Chunk size, overlap and the algorithm can't be set. For many document RAG cases that's acceptable, since semantic chunking with added context often beats plain fixed-size splits. Parent-child chunks or detailed overlap control still call for custom logic.

ai_extract adds cost and latency, because each document needs one more LLM call. Delta Sync can skip unchanged documents, but the first full ingestion may still be costly, so testing on a sample before the full corpus is smart.

Use version 2.0 of ai_parse_document when its VARIANT output goes into ai_prep_search or ai_extract. That version supplies bounding boxes, element types and figure descriptions. Raw text can be passed to ai_extract, though the structural benefits of parsing are lost.

Wrapping the design up

Adding ai_extract v2.1 takes little extra code but changes the architecture meaningfully, because chunks now carry clean, typed, filterable metadata and hybrid retrieval gets more dependable. The pipeline stays at three clear stages:

1. ai_parse_document captures document structure.
2. ai_extract captures useful metadata.
3. ai_prep_search builds semantic chunks.

The separation pays off across documents from different products, versions, audiences and types. With no metadata filters, retrieval leans mainly on similarity search. With them, the agent narrows first and ranks afterwards.

The approach won't fit everything. Chunk-level metadata, parent-child retrieval or custom chunking still need custom code. For many Databricks document RAG pipelines, though, it's a cleaner and simpler place to start. And chunks should still be tested and evaluated and never trusted blindly.