Databricks Vector Search Client

0.67 · active · verified Fri Apr 10

The `databricks-vectorsearch` Python client library provides programmatic access to Databricks Vector Search, a serverless similarity search engine. It enables users to manage vector search endpoints and indexes, facilitating the creation of Retrieval Augmented Generation (RAG) applications. The library integrates seamlessly with the Databricks Data Intelligence Platform and Unity Catalog for data governance. Currently at version 0.67, it is actively developed with regular updates.

Warnings

breaking In Databricks Runtime 15.3 and above, the `vector_search()` SQL function's `query` argument was replaced by `query_text` or `query_vector`. Using `query` in newer runtimes will result in an error.
Fix: Update your SQL queries to use `query_text` for text-based searches or `query_vector` for vector-based searches, depending on your index configuration.
gotcha The `VectorSearchIndex` class should not be instantiated directly. Instead, obtain an index instance using methods provided by the `VectorSearchClient` (e.g., `client.get_index()`).
Fix: Always interact with index objects obtained via `VectorSearchClient` methods.
gotcha For production deployments, Databricks strongly recommends using Service Principals for authentication instead of Personal Access Tokens (PATs). Service Principals offer better security and can improve query performance.
Fix: Configure your client with `service_principal_client_id` and `service_principal_client_secret` along with `azure_tenant_id` for Azure, or equivalent for other clouds, and enable service principal usage for vector search index ownership.
gotcha Removing the owner of a Vector Search Index can break associated Databricks Liquid (DLT) pipelines and other dependencies, requiring index recreation. Assigning a Service Principal as the index owner prevents this issue.
Fix: When creating or modifying Vector Search Indexes, ensure a Service Principal is assigned as the owner.
deprecated The `budget_policy_id` parameter for creating Vector Search endpoints and indexes has been deprecated. It has been replaced by `usage_policy_id`.
Fix: Use the `usage_policy_id` parameter instead of `budget_policy_id` when creating or updating endpoints and indexes.

Install

pip install databricks-vectorsearch Install latest version

Imports

VectorSearchClient

from databricks.vector_search.client import VectorSearchClient

Quickstart

Initializes the VectorSearchClient and performs a similarity search on an existing index. Ensure `DATABRICKS_WORKSPACE_URL` and `DATABRICKS_TOKEN` environment variables are set or replaced with actual values.

import os
from databricks.vector_search.client import VectorSearchClient

# Replace with your Databricks workspace URL and Personal Access Token
workspace_url = os.environ.get("DATABRICKS_WORKSPACE_URL", "https://<your-workspace-url>.databricks.com")
pat = os.environ.get("DATABRICKS_TOKEN", "")

if not pat:
    print("Warning: DATABRICKS_TOKEN environment variable not set. This example may not work without authentication.")

# Initialize the Vector Search Client
client = VectorSearchClient(workspace_url=workspace_url, personal_access_token=pat)

# Replace with your endpoint and index names
endpoint_name = "<your-vector-search-endpoint>"
index_name = "<your-catalog>.<your-schema>.<your-index-name>"

try:
    # Get an existing index instance
    index = client.get_index(index_name=index_name)

    # Perform a similarity search
    query_text = "What is Databricks Vector Search?"
    results = index.similarity_search(
        query_text=query_text,
        num_results=2,
        columns=["text", "id"] # specify columns to return
    )

    print(f"Similarity search results for '{query_text}':")
    for result in results.get('result', {}).get('data_array', []):
        print(f"  ID: {result}, Text: {result[:100]}...")

except Exception as e:
    print(f"An error occurred: {e}")
    print("Please ensure the Vector Search endpoint and index exist and your credentials are correct.")

view raw JSON →