Mosaic AI Vector Search in Python: Create, Query and Filter an Index

What this deep dive adds

Chapter 33 of the course builds one Delta Sync index and runs one query with a single equality filter. This post is the reference behind that chapter: complete create code for all three index flavors, keeping an index current, every similarity_search argument, the full filter syntax on both endpoint types, the reranker, and the retriever tool an agent calls.

"Create and query a Vector Search index" is an objective in Section 4, Assembling and Deploying Applications, of the 18 March 2026 exam guide. The concepts around it (endpoints, index types, sizing) are in the Vector Search and RAG post.

Create each index flavor

A Standard endpoint and a Delta Sync index whose embeddings Databricks computes:

%pip install databricks-ai-search
dbutils.library.restartPython()

from databricks.ai_search.client import AISearchClient

client = AISearchClient()
client.create_endpoint(name="rag_endpoint", endpoint_type="STANDARD")

index = client.create_delta_sync_index(
    endpoint_name="rag_endpoint",
    source_table_name="main.rag.doc_chunks",
    index_name="main.rag.doc_chunks_index",
    pipeline_type="TRIGGERED",
    primary_key="chunk_id",
    embedding_source_column="chunk_text",
    embedding_model_endpoint_name="databricks-qwen3-embedding-0-6b",
    columns_to_sync=["chunk_text", "source_url", "region", "doc_year"],
)

Databricks embeds chunk_text through the named endpoint at sync time, and each query_text with the same model. columns_to_sync matters more than it looks: by default every source column is synced, but when set, only those columns plus the primary key and embedding column are indexed, and only indexed columns can be returned or filtered on. A source_url left out here never comes back.

When an upstream pipeline already writes vectors, point the index at that column instead:

index = client.create_delta_sync_index(
    endpoint_name="rag_endpoint",
    source_table_name="main.rag.doc_chunks_embedded",
    index_name="main.rag.doc_chunks_embedded_index",
    pipeline_type="TRIGGERED",
    primary_key="chunk_id",
    embedding_dimension=1024,
    embedding_vector_column="chunk_vector",
)

The vector column is array<float> and embedding_dimension must match it (on Storage Optimized endpoints the dimension must divide evenly by 16). Queries now pass query_vector, embedded by the same model you used upstream.

A Direct Vector Access index has no source table, so you declare its schema yourself. It cannot be created in the UI, only through the SDK or REST API:

index = client.create_direct_access_index(
    endpoint_name="rag_endpoint",
    index_name="main.rag.live_notes_index",
    primary_key="note_id",
    embedding_dimension=1024,
    embedding_vector_column="note_vector",
    schema={
        "note_id": "int",
        "note_text": "string",
        "region": "string",
        "note_vector": "array<float>",
    },
)

Keeping an index current

# Delta Sync on TRIGGERED: start a sync after the source table changes
index = client.get_index(index_name="main.rag.doc_chunks_index")
index.sync()

# Direct Vector Access: your code writes every change
live = client.get_index(index_name="main.rag.live_notes_index")
live.upsert([
    {"note_id": 1, "note_text": "Refunds post in 5 days",
     "region": "US", "note_vector": [0.12] * 1024},
])
live.delete(primary_keys=[1])

A CONTINUOUS Delta Sync index updates itself, so sync() is for TRIGGERED pipelines, typically the last task of the job refreshing the chunk table. upsert() takes a list of dicts matching the declared schema; delete() takes a list of primary keys. On a Delta Sync index you change the source table instead.

The index schema is fixed at creation: returning a new source column means building a new index and switching traffic to it.

Every similarity_search argument

Sourced from the Databricks documentation pages "Create AI Search endpoints and indexes" and "Query an AI Search index", and the similarity_search reference in the databricks-ai-search Python SDK (version 0.82).

Filter syntax on each endpoint type

Standard endpoints take a dict whose key holds column and operator; Storage Optimized endpoints take one SQL-like WHERE string.

Operators and examples from the "Use filters on queries" section of the Databricks page "Query an AI Search index".

On Standard, several keys in one dict must all match, as in {"region": ["US", "CA"], "doc_year >=": 2024}. A Standard LIKE matches whitespace-separated tokens, so "apple" matches "apple pear" but not "pineapple"; the Storage Optimized LIKE is the SQL one.

The bigger difference is when the filter runs. Storage Optimized endpoints overfetch and then filter what was fetched, so a selective filter can return fewer rows than num_results, or none, even when matching rows exist lower down.

ANN, HYBRID and the reranker

from databricks.ai_search.reranker import DatabricksReranker

results = index.similarity_search(
    query_text="error E-4031 at checkout",
    columns=["chunk_id", "chunk_text", "source_url"],
    filters={"region": "US"},
    num_results=5,
    query_type="HYBRID",
    reranker=DatabricksReranker(columns_to_rerank=["chunk_text"]),
    debug_level=1,
)
rows = results["result"]["data_array"]
print(results["debug_info"])

ANN compares embeddings only, which suits paraphrased questions. HYBRID runs ANN and a keyword search in parallel and merges them with Reciprocal Rank Fusion, which catches exact codes like the one above; it returns at most 200 results. The reranker is a cross-encoder second pass over the first-pass results. columns_to_rerank is read in order and only its first 2,000 characters count, so put short, informative columns first. num_results is still the final count.

The index as an agent's retriever

from databricks_langchain import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name="main.rag.doc_chunks_index",
    tool_name="policy_search",
    tool_description="Search refund and shipping policies.",
    num_results=5,
    columns=["chunk_id", "chunk_text", "source_url"],
    filters={"region": "US"},
    query_type="HYBRID",
)
docs = vs_tool.invoke("How long do refunds take?")

For new Python agents the docs recommend VectorSearchRetrieverTool from databricks-langchain (or databricks-openai): the description tells the model when to search, and columns, filters and query_type work as above. On a self-managed or Direct Vector Access index it also needs text_column and an embedding model to embed the query. The DatabricksVectorSearch vector store, wrapped with as_retriever(), still serves a fixed chain, with LangChain's names k and filter.

From the Databricks page "Connect agents to unstructured data" and the databricks-langchain 0.22 package source.

What this means for the exam

Distractors change the wrong argument: more num_results for more fields, HYBRID for metadata, or a join back to the source table.

FAQ