Skip to main content

Databricks Vector Search

Databricks Vector Search is a serverless similarity search engine that allows you to store a vector representation of your data, including metadata, in a vector database. With Vector Search, you can create auto-updating vector search indexes from Delta tables managed by Unity Catalog and query them with a simple API to return the most similar vectors.

This notebook shows how to use LangChain with Databricks Vector Search.

Install databricks-vectorsearch and related Python packages used in this notebook.

%pip install --upgrade --quiet  langchain-core databricks-vectorsearch langchain-openai tiktoken

Use OpenAIEmbeddings for the embeddings.

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

Split documents and get embeddings.

from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("../../how_to/state_of_the_union.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()
emb_dim = len(embeddings.embed_query("hello"))

Setup Databricks Vector Search clientโ€‹

from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

Create a Vector Search Endpointโ€‹

This endpoint is used to create and access vector search indexes.

vsc.create_endpoint(name="vector_search_demo_endpoint", endpoint_type="STANDARD")

Create Direct Vector Access Indexโ€‹

Direct Vector Access Index supports direct read and write of embedding vectors and metadata through a REST API or an SDK. For this index, you manage embedding vectors and index updates yourself.

vector_search_endpoint_name = "vector_search_demo_endpoint"
index_name = "vector_search_demo.vector_search.state_of_the_union_index"

index = vsc.create_direct_access_index(
endpoint_name=vector_search_endpoint_name,
index_name=index_name,
primary_key="id",
embedding_dimension=emb_dim,
embedding_vector_column="text_vector",
schema={
"id": "string",
"text": "string",
"text_vector": "array<float>",
"source": "string",
},
)

index.describe()
from langchain_community.vectorstores import DatabricksVectorSearch

dvs = DatabricksVectorSearch(
index, text_column="text", embedding=embeddings, columns=["source"]
)

Add docs to the indexโ€‹

dvs.add_documents(docs)

Optional keyword arguments to similarity_search include specifying k number of documents to retrive, a filters dictionary for metadata filtering based on this syntax, as well as the query_type which can be ANN or HYBRID

query = "What did the president say about Ketanji Brown Jackson"
dvs.similarity_search(query)
print(docs[0].page_content)

Work with Delta Sync Indexโ€‹

You can also use DatabricksVectorSearch to search in a Delta Sync Index. Delta Sync Index automatically syncs from a Delta table. You don't need to call add_text/add_documents manually. See Databricks documentation page for more details.

delta_sync_index = vsc.create_delta_sync_index(
endpoint_name=vector_search_endpoint_name,
source_table_name="vector_search_demo.vector_search.state_of_the_union",
index_name="vector_search_demo.vector_search.state_of_the_union_index",
pipeline_type="TRIGGERED",
primary_key="id",
embedding_source_column="text",
embedding_model_endpoint_name="e5-small-v2",
)
dvs_delta_sync = DatabricksVectorSearch(delta_sync_index)
dvs_delta_sync.similarity_search(query)

Was this page helpful?


You can also leave detailed feedback on GitHub.