Skip to content

feat(vector-search): migrate to databricks-ai-search SDK, add filtering reference #197

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jiayi-wu-3150 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(vector-search): migrate to databricks-ai-search SDK, add filtering reference #197

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 73 additions & 110 deletions skills/databricks-vector-search/SKILL.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
---
name: databricks-vector-search
description: "Databricks Vector Search endpoints and indexes for RAG and semantic search; covers index types, search modes, end-to-end RAG patterns"
description: "Databricks AI Search (formerly Vector Search) endpoints and indexes for RAG and semantic search; covers index types, search modes, filtering, end-to-end RAG patterns"
metadata:
version: "0.1.0"
parent: databricks-core
---

# Databricks Vector Search
# Databricks AI Search (formerly Vector Search)

**FIRST**: Use the parent `databricks-core` skill for CLI basics, authentication, and profile selection.

Patterns for creating, managing, and querying vector search indexes for RAG and semantic search applications.
Patterns for creating, managing, and querying AI Search indexes for RAG and semantic search applications.

## When to Use

Expand All @@ -23,7 +23,7 @@ Use this skill when:

## Overview

Databricks Vector Search provides managed vector similarity search with automatic embedding generation and Delta Lake integration.
Databricks AI Search provides managed vector similarity search with automatic embedding generation and Delta Lake integration.

| Component | Description |
|-----------|-------------|
Expand All @@ -47,17 +47,22 @@ Databricks Vector Search provides managed vector similarity search with automati
| **Delta Sync (self-managed)** | You provide | Auto from Delta | Custom embeddings |
| **Direct Access** | You provide | Manual CRUD | Real-time updates |

## Installation

```bash
%pip install databricks-ai-search
```

## Quick Start

### Create Endpoint

```python
from databricks.sdk import WorkspaceClient
from databricks.ai_search.client import AISearchClient

w = WorkspaceClient()
client = AISearchClient()

# Create a standard endpoint
endpoint = w.vector_search_endpoints.create_endpoint(
client.create_endpoint(
name="my-vs-endpoint",
endpoint_type="STANDARD" # or "STORAGE_OPTIMIZED"
)
Expand All @@ -68,37 +73,30 @@ endpoint = w.vector_search_endpoints.create_endpoint(

```python
# Source table must have: primary key column + text column
index = w.vector_search_indexes.create_index(
name="catalog.schema.my_index",
index = client.create_delta_sync_index(
endpoint_name="my-vs-endpoint",
source_table_name="catalog.schema.documents",
index_name="catalog.schema.my_index",
pipeline_type="TRIGGERED", # or "CONTINUOUS"
primary_key="id",
index_type="DELTA_SYNC",
delta_sync_index_spec={
"source_table": "catalog.schema.documents",
"embedding_source_columns": [
{
"name": "content", # Text column to embed
"embedding_model_endpoint_name": "databricks-gte-large-en"
}
],
"pipeline_type": "TRIGGERED" # or "CONTINUOUS"
}
embedding_source_column="content",
embedding_model_endpoint_name="databricks-gte-large-en"
)
```

### Query Index

```python
results = w.vector_search_indexes.query_index(
index_name="catalog.schema.my_index",
columns=["id", "content", "metadata"],
index = client.get_index(
endpoint_name="my-vs-endpoint",
index_name="catalog.schema.my_index"
)

results = index.similarity_search(
query_text="What is machine learning?",
columns=["id", "content", "metadata"],
num_results=5
)

for doc in results.result.data_array:
score = doc[-1] # Similarity score is last column
print(f"Score: {score}, Content: {doc[1][:100]}...")
```

## Common Patterns
Expand All @@ -107,7 +105,7 @@ for doc in results.result.data_array:

```python
# For large-scale, cost-effective deployments
endpoint = w.vector_search_endpoints.create_endpoint(
client.create_endpoint(
name="my-storage-endpoint",
endpoint_type="STORAGE_OPTIMIZED"
)
Expand All @@ -117,72 +115,52 @@ endpoint = w.vector_search_endpoints.create_endpoint(

```python
# Source table must have: primary key + embedding vector column
index = w.vector_search_indexes.create_index(
name="catalog.schema.my_index",
index = client.create_delta_sync_index(
endpoint_name="my-vs-endpoint",
source_table_name="catalog.schema.documents",
index_name="catalog.schema.my_index",
pipeline_type="TRIGGERED",
primary_key="id",
index_type="DELTA_SYNC",
delta_sync_index_spec={
"source_table": "catalog.schema.documents",
"embedding_vector_columns": [
{
"name": "embedding", # Pre-computed embedding column
"embedding_dimension": 768
}
],
"pipeline_type": "TRIGGERED"
}
embedding_dimension=768,
embedding_vector_column="embedding"
)
```

### Direct Access Index

```python
import json

# Create index for manual CRUD
index = w.vector_search_indexes.create_index(
name="catalog.schema.direct_index",
index = client.create_direct_access_index(
endpoint_name="my-vs-endpoint",
index_name="catalog.schema.direct_index",
primary_key="id",
index_type="DIRECT_ACCESS",
direct_access_index_spec={
"embedding_vector_columns": [
{"name": "embedding", "embedding_dimension": 768}
],
"schema_json": json.dumps({
"id": "string",
"text": "string",
"embedding": "array<float>",
"metadata": "string"
})
embedding_dimension=768,
embedding_vector_column="embedding",
schema={
"id": "string",
"text": "string",
"embedding": "array<float>",
"metadata": "string"
}
)

# Upsert data
w.vector_search_indexes.upsert_data_vector_index(
index_name="catalog.schema.direct_index",
inputs_json=json.dumps([
{"id": "1", "text": "Hello", "embedding": [0.1, 0.2, ...], "metadata": "doc1"},
{"id": "2", "text": "World", "embedding": [0.3, 0.4, ...], "metadata": "doc2"},
])
)
index.upsert([
{"id": "1", "text": "Hello", "embedding": [0.1, 0.2, ...], "metadata": "doc1"},
{"id": "2", "text": "World", "embedding": [0.3, 0.4, ...], "metadata": "doc2"},
])

# Delete data
w.vector_search_indexes.delete_data_vector_index(
index_name="catalog.schema.direct_index",
primary_keys=["1", "2"]
)
index.delete(primary_keys=["1", "2"])
```

### Query with Embedding Vector

```python
# When you have pre-computed query embedding
results = w.vector_search_indexes.query_index(
index_name="catalog.schema.my_index",
columns=["id", "text"],
results = index.similarity_search(
query_vector=[0.1, 0.2, 0.3, ...], # Your 768-dim vector
columns=["id", "text"],
num_results=10
)
```
Expand All @@ -193,71 +171,56 @@ Hybrid search combines vector similarity (ANN) with BM25 keyword scoring. Use it

```python
# Combines vector similarity with keyword matching
results = w.vector_search_indexes.query_index(
index_name="catalog.schema.my_index",
columns=["id", "content"],
results = index.similarity_search(
query_text="SPARK-12345 executor memory error",
query_type="HYBRID",
query_type="hybrid",
columns=["id", "content"],
num_results=10
)
```

## Filtering

Filter syntax differs by endpoint type. See [references/filtering.md](references/filtering.md) for the full operator reference.

### Standard Endpoint Filters (Dictionary)

```python
# filters_json uses dictionary format
results = w.vector_search_indexes.query_index(
index_name="catalog.schema.my_index",
columns=["id", "content"],
results = index.similarity_search(
query_text="machine learning",
columns=["id", "content"],
num_results=10,
filters_json='{"category": "ai", "status": ["active", "pending"]}'
filters={"category": "ai", "status": ["active", "pending"]}
)
```

### Storage-Optimized Filters (SQL-like)

Storage-Optimized endpoints use SQL-like filter syntax via the `databricks-vectorsearch` package's `filters` parameter (accepts a string):
### Storage-Optimized Filters (SQL-like string)

```python
from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()
index = vsc.get_index(endpoint_name="my-storage-endpoint", index_name="catalog.schema.my_index")

# SQL-like filter syntax for storage-optimized endpoints
results = index.similarity_search(
query_text="machine learning",
columns=["id", "content"],
num_results=10,
filters="category = 'ai' AND status IN ('active', 'pending')"
)

# More filter examples
# filters="price > 100 AND price < 500"
# filters="department LIKE 'eng%'"
# filters="created_at >= '2024-01-01'"
```

### Trigger Index Sync

```python
# For TRIGGERED pipeline type, manually sync
w.vector_search_indexes.sync_index(
index = client.get_index(
endpoint_name="my-vs-endpoint",
index_name="catalog.schema.my_index"
)
index.sync()
```

### Scan All Index Entries

```python
# Retrieve all vectors (for debugging/export)
scan_result = w.vector_search_indexes.scan_index(
index_name="catalog.schema.my_index",
num_results=100
)
index.scan(num_results=100)
```

## Reference Files
Expand All @@ -266,8 +229,9 @@ scan_result = w.vector_search_indexes.scan_index(
|-------|------|-------------|
| Index Types | [references/index-types.md](references/index-types.md) | Detailed comparison of Delta Sync (managed/self-managed) vs Direct Access |
| End-to-End RAG | [references/end-to-end-rag.md](references/end-to-end-rag.md) | Complete walkthrough: source table → endpoint → index → query → agent integration |
| Search Modes | [references/search-modes.md](references/search-modes.md) | When to use semantic (ANN) vs hybrid search, decision guide |
| Operations | [references/troubleshooting-and-operations.md](references/troubleshooting-and-operations.md) | Monitoring, cost optimization, capacity planning, migration |
| Search Modes | [references/search-modes.md](references/search-modes.md) | When to use semantic (ANN) vs hybrid search, reranker, decision guide |
| Filtering | [references/filtering.md](references/filtering.md) | Full filter operator reference for Standard (dict) and Storage-Optimized (SQL string) endpoints |
| Operations | [references/troubleshooting-and-operations.md](references/troubleshooting-and-operations.md) | Monitoring, cost optimization, capacity planning, migration, performance targets |

## CLI Quick Reference

Expand Down Expand Up @@ -297,9 +261,9 @@ databricks vector-search-indexes delete-index catalog.schema.my_index
|-------|----------|
| **Index sync slow** | Use Storage-Optimized endpoints (20x faster indexing) |
| **Query latency high** | Use Standard endpoint for <100ms latency |
| **filters_json not working** | Storage-Optimized uses SQL-like string filters via `databricks-vectorsearch` package's `filters` parameter |
| **Filters not working** | Standard endpoints use a dict (`filters={"col": "val"}`); Storage-Optimized use a SQL string (`filters="col = 'val'"`). See [references/filtering.md](references/filtering.md) |
| **Embedding dimension mismatch** | Ensure query and index dimensions match |
| **Index not updating** | Check pipeline_type; use sync_index() for TRIGGERED |
| **Index not updating** | Check pipeline_type; use `index.sync()` for TRIGGERED |
| **Out of capacity** | Upgrade to Storage-Optimized (1B+ vectors) |
| **`query_vector` truncated** | Large vectors (e.g. 1024-dim) can be truncated when serialized as JSON. Use `query_text` instead (for managed embedding indexes), or use the Databricks SDK to pass raw vectors |

Expand All @@ -314,12 +278,11 @@ Databricks provides built-in embedding models:

```python
# Use with managed embeddings
embedding_source_columns=[
{
"name": "content",
"embedding_model_endpoint_name": "databricks-gte-large-en"
}
]
index = client.create_delta_sync_index(
...
embedding_source_column="content",
embedding_model_endpoint_name="databricks-gte-large-en"
)
```

## Notes
Expand All @@ -328,7 +291,7 @@ embedding_source_columns=[
- **Delta Sync recommended** — easier than Direct Access for most scenarios
- **Hybrid search** — available for both Delta Sync and Direct Access indexes
- **`columns_to_sync` matters** — only synced columns are available in query results; include all columns you need
- **Filter syntax differs by endpoint** — Standard uses dict-format filters, Storage-Optimized uses SQL-like string filters. Use the `databricks-vectorsearch` package's `filters` parameter which accepts both formats
- **Filter syntax differs by endpoint** — Standard uses a dict, Storage-Optimized uses a SQL-like string. See [references/filtering.md](references/filtering.md)
- **Management vs runtime** — CLI and SDK handle lifecycle management; for agent tool-calling at runtime, use `VectorSearchRetrieverTool`

## Related Skills
Expand Down
Loading