Quickstart: Vector index with Python in Azure DocumentDB

In this quickstart, you compare three vector index algorithms (DiskANN, HNSW, and IVF) and three similarity functions (cosine, L2, and inner product) to find the optimal configuration for your search workload. This quickstart uses a sample hotel dataset with precalculated embeddings from the text-embedding-3-small model.

Find the sample code on GitHub.

Prerequisites

  • An Azure subscription. If you don't have an Azure subscription, create a free account.

Create a Python project

  1. Create a new directory for your project and open it in Visual Studio Code:

    mkdir -p select-algorithm-python
cd select-algorithm-python
code .

  1. In the terminal, create and activate a virtual environment:

    For Windows:

    python -m venv venv
venv\Scripts\activate

    For macOS/Linux:

    python -m venv venv
source venv/bin/activate

  2. Create a requirements.txt file with the following content:

    pymongo>=4.7.0
openai>=1.0.0,<2.0.0
azure-identity>=1.15.0
tabulate>=0.9.0

  3. Install the required packages:

    pip install -r requirements.txt

    The requirements include:

    • pymongo: MongoDB driver for Python (≥4.7 required for OIDC authentication).
    • openai: OpenAI client library to create vectors.
    • azure-identity: Azure Identity library for passwordless authentication.
    • tabulate: Formatted table output for comparison results.

  4. Create the src directory:

mkdir -p src

Create data file with vectors

  1. Create a new data directory and download the hotels data file with vectors:

    mkdir -p data
curl -o data/Hotels_Vector.json https://raw.githubusercontent.com/Azure-Samples/documentdb-samples/refs/heads/main/ai/data/Hotels_Vector.json

Verify the file was downloaded:

ls data/

You should see Hotels_Vector.json in the data directory.

Configure environment variables

Set the required environment variables in your current shell session before you run the sample:

export DOCUMENTDB_CLUSTER_NAME=<your-cluster-name>
export AZURE_OPENAI_EMBEDDING_ENDPOINT=https://<your-resource>.openai.azure.com
export AZURE_OPENAI_EMBEDDING_MODEL=text-embedding-3-small
export AZURE_DOCUMENTDB_DATABASENAME=Hotels
export DATA_FILE_WITH_VECTORS=data/Hotels_Vector.json
export EMBEDDED_FIELD=DescriptionVector
export EMBEDDING_DIMENSIONS=1536

For the passwordless authentication in this article, replace the placeholder values with your own information:

  • AZURE_OPENAI_EMBEDDING_ENDPOINT: Your Azure OpenAI resource endpoint URL
  • DOCUMENTDB_CLUSTER_NAME: Your Azure DocumentDB cluster name

Prefer passwordless authentication, but it requires additional setup. For more information on setting up managed identity and the full range of your authentication options, see Authenticate Python apps to Azure services by using the Azure SDK for Python.

Create code files

Create the following project structure:

select-algorithm-python/
├── data/
│   └── Hotels_Vector.json
├── src/
│   ├── compare_all.py
│   └── utils.py
└── requirements.txt

Create the algorithm comparison code

Create the src/compare_all.py file with the following code:

"""
Compare All Algorithms — Unified comparison runner.

Executes all 9 combinations (3 algorithms × 3 similarity metrics) in a single
invocation and prints a formatted comparison table.

Algorithms: IVF, HNSW, DiskANN
Metrics: COS, L2, IP
"""
import os
import sys
import time
from typing import Dict, List, Any

from tabulate import tabulate
from utils import (
    get_clients_passwordless, get_config, read_file_return_json,
    insert_data
)

# Index definitions: (algo_label, kind, extra_params)
ALGORITHMS = [
    ("IVF", "vector-ivf", {"numLists": 1}),
    ("HNSW", "vector-hnsw", {"m": 16, "efConstruction": 64}),
    ("DiskANN", "vector-diskann", {"maxDegree": 32, "lBuild": 50}),
]

METRICS = ["COS", "L2", "IP"]


def get_compare_config() -> Dict[str, Any]:
    """Load comparison-specific configuration from environment variables."""
    config = get_config()
    config["query_text"] = os.getenv("QUERY_TEXT", "luxury hotel near the beach")
    config["top_k"] = int(os.getenv("TOP_K", "5"))
    return config


def index_name(algo: str, metric: str) -> str:
    """Generate canonical index name: vector_{algo}_{metric}."""
    return f"vector_{algo.lower()}_{metric.lower()}"


def get_existing_index_names(collection) -> List[str]:
    """Return names of existing indexes on the collection."""
    return [idx["name"] for idx in collection.list_indexes()]


def drop_vector_indexes(collection, vector_field: str) -> None:
    """Drop all existing vector indexes on *vector_field*."""
    for idx in collection.list_indexes():
        name = idx.get("name", "")
        key = idx.get("key", {})
        if vector_field in key and key[vector_field] == "cosmosSearch":
            collection.drop_index(name)


def create_vector_index(collection, name: str, kind: str, vector_field: str,
                        dimensions: int, similarity: str,
                        extra_params: Dict[str, Any]) -> None:
    """Create a single vector index."""
    cosmos_options = {
        "kind": kind,
        "dimensions": dimensions,
        "similarity": similarity,
        **extra_params,
    }

    index_command = {
        "createIndexes": collection.name,
        "indexes": [
            {
                "name": name,
                "key": {vector_field: "cosmosSearch"},
                "cosmosSearchOptions": cosmos_options,
            }
        ],
    }
    collection.database.command(index_command)


def generate_embedding(azure_openai_client, query_text: str,
                       model_name: str) -> List[float]:
    """Generate a single embedding for the query text."""
    response = azure_openai_client.embeddings.create(
        input=[query_text],
        model=model_name
    )
    return response.data[0].embedding


def vector_search_with_index(collection, query_embedding: List[float],
                             vector_field: str,
                             top_k: int) -> List[Dict[str, Any]]:
    """Run vector search using the single active index and return results."""
    pipeline = [
        {
            "$search": {
                "cosmosSearch": {
                    "vector": query_embedding,
                    "path": vector_field,
                    "k": top_k
                }
            }
        },
        {
            "$project": {
                "document": "$$ROOT",
                "score": {"$meta": "searchScore"}
            }
        }
    ]

    results = list(collection.aggregate(pipeline))

    return results


def vector_search_with_retry(collection, query_embedding: List[float],
                             vector_field: str, top_k: int,
                             index_name_value: str) -> List[Dict[str, Any]]:
    """Retry vector search until results are available or retries are exhausted."""
    max_retries = 5
    retry_delay_seconds = 2

    for attempt in range(max_retries + 1):
        results = vector_search_with_index(
            collection, query_embedding, vector_field, top_k
        )
        if results:
            return results

        if attempt < max_retries:
            print(
                f"  No results for '{index_name_value}' yet. "
                f"Retrying in {retry_delay_seconds} seconds "
                f"({attempt + 1}/{max_retries})..."
            )
            time.sleep(retry_delay_seconds)

    print(
        f"  Search for '{index_name_value}' did not return results "
        f"after {max_retries} retries. Recording as failed."
    )
    return []


def main():
    print("=" * 70)
    print("  Compare All Algorithms — 9 Combinations")
    print("  (3 Algorithms × 3 Similarity Metrics)")
    print("=" * 70)

    config = get_compare_config()
    query_text = config["query_text"]
    top_k = config["top_k"]

    print(f"\n  Query:  \"{query_text}\"")
    print(f"  Top K:  {top_k}\n")

    mongo_client, azure_openai_client = get_clients_passwordless()

    try:
        database = mongo_client[config["database_name"]]

        # Drop collection for a clean comparison
        database.drop_collection("hotels")
        print("Dropped existing 'hotels' collection (if any)")

        # Create fresh collection and load data
        collection = database["hotels"]
        data = read_file_return_json(config["data_file"])
        documents = [doc for doc in data if config["vector_field"] in doc]
        print(f"Loaded {len(documents)} documents with embeddings")
        insert_data(collection, documents, config["batch_size"])

        # Generate ONE embedding for the query
        print("\nGenerating embedding for query...")
        query_embedding = generate_embedding(
            azure_openai_client, query_text, config["model_name"]
        )

        # Run all 9 searches sequentially (create→search→drop for each)
        print("Running 9 vector searches...\n")
        table_rows = []

        for algo_label, kind, extra_params in ALGORITHMS:
            for metric in METRICS:
                name = index_name(algo_label, metric)
                # Drop all vector indexes first
                drop_vector_indexes(collection, config["vector_field"])
                # Create this specific index
                create_vector_index(
                    collection, name, kind, config["vector_field"],
                    config["dimensions"], metric, extra_params
                )
                print(f"  Created index '{name}'")
                results = vector_search_with_retry(
                    collection, query_embedding, config["vector_field"], top_k, name
                )

                if not results:
                    table_rows.append([
                        algo_label,
                        metric,
                        "(failed)",
                        f"{0:.4f}",
                        "(failed)",
                        f"{0:.4f}",
                        f"{0:.4f}",
                    ])
                    continue

                top1_name = results[0].get("document", results[0]).get("HotelName", "Unknown") if len(results) > 0 else "(no results)"
                top1_score = results[0].get("score", 0) if len(results) > 0 else 0
                top2_name = results[1].get("document", results[1]).get("HotelName", "Unknown") if len(results) > 1 else "(no results)"
                top2_score = results[1].get("score", 0) if len(results) > 1 else 0

                table_rows.append([
                    algo_label,
                    metric,
                    top1_name,
                    f"{top1_score:.4f}",
                    top2_name,
                    f"{top2_score:.4f}",
                    f"{abs(top1_score - top2_score):.4f}",
                ])

        # Print comparison table
        headers = ["Algorithm", "Metric", "Top 1 Result", "Score",
                   "Top 2 Result", "Score", "Diff"]
        print(tabulate(table_rows, headers=headers, tablefmt="grid"))

        success_count = sum(1 for row in table_rows if row[2] != "(failed)")
        if success_count == 0:
            print("\n❌ All 9 comparisons failed — no algorithm returned results.")
            sys.exit(1)
        else:
            print(f"\nSummary: {success_count} succeeded, {9 - success_count} failed")

    finally:
        # Cleanup: drop the comparison collection
        try:
            database = mongo_client[config["database_name"]]
            database.drop_collection("hotels")
            print("\nCleanup: dropped collection 'hotels'")
        except Exception as e:
            print(f"Cleanup warning: {e}")
        mongo_client.close()


if __name__ == "__main__":
    main()

This script orchestrates the algorithm comparison by:

  • Loading configuration from environment variables.
  • Initializing MongoDB and Azure OpenAI clients with passwordless authentication.
  • Loading hotel data with precalculated embeddings.
  • Testing each algorithm/similarity combination by creating a collection, inserting data, creating an index, and executing a search.
  • Measuring and comparing search performance across all configurations.
  • Displaying results in a comparison table.

Create utility functions

Create the src/utils.py file with the following code:

import json
import os
import time
import warnings
from typing import Dict, List, Any, Optional, Tuple

# Suppress the PyMongo DocumentDB cluster detection warning
warnings.filterwarnings(
    "ignore",
    message="You appear to be connected to a DocumentDB cluster.*",
)

from pymongo import MongoClient, InsertOne
from pymongo.collection import Collection
from pymongo.errors import BulkWriteError
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from pymongo.auth_oidc import OIDCCallback, OIDCCallbackContext, OIDCCallbackResult
from openai import AzureOpenAI


class AzureIdentityTokenCallback(OIDCCallback):
    def __init__(self, credential):
        self.credential = credential

    def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
        token = self.credential.get_token(
            "https://ossrdbms-aad.database.windows.net/.default").token
        return OIDCCallbackResult(access_token=token)


def get_clients_passwordless() -> Tuple[MongoClient, AzureOpenAI]:
    """Create MongoDB and Azure OpenAI clients using passwordless auth."""
    cluster_name = os.getenv("DOCUMENTDB_CLUSTER_NAME")
    if not cluster_name:
        raise ValueError("DOCUMENTDB_CLUSTER_NAME environment variable is required")

    credential = DefaultAzureCredential()

    mongo_client = MongoClient(
        f"mongodb+srv://{cluster_name}.global.mongocluster.cosmos.azure.com/",
        connectTimeoutMS=120000,
        tls=True,
        retryWrites=False,
        authMechanism="MONGODB-OIDC",
        authMechanismProperties={"OIDC_CALLBACK": AzureIdentityTokenCallback(credential)}
    )

    azure_openai_endpoint = os.getenv("AZURE_OPENAI_EMBEDDING_ENDPOINT")
    if not azure_openai_endpoint:
        raise ValueError("AZURE_OPENAI_EMBEDDING_ENDPOINT environment variable is required")

    token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")

    azure_openai_client = AzureOpenAI(
        azure_endpoint=azure_openai_endpoint,
        azure_ad_token_provider=token_provider,
        api_version=os.getenv("AZURE_OPENAI_EMBEDDING_API_VERSION", "2023-05-15")
    )

    return mongo_client, azure_openai_client


def get_config() -> Dict[str, Any]:
    """Load configuration from environment variables."""
    return {
        'database_name': os.getenv('AZURE_DOCUMENTDB_DATABASENAME', 'Hotels'),
        'data_file': os.getenv('DATA_FILE_WITH_VECTORS', 'data/Hotels_Vector.json'),
        'vector_field': os.getenv('EMBEDDED_FIELD', 'DescriptionVector'),
        'model_name': os.getenv('AZURE_OPENAI_EMBEDDING_MODEL', 'text-embedding-3-small'),
        'dimensions': int(os.getenv('EMBEDDING_DIMENSIONS', '1536')),
        'batch_size': int(os.getenv('LOAD_SIZE_BATCH', '100')),
        'similarity': os.getenv('SIMILARITY', ''),
    }


def read_file_return_json(file_path: str) -> List[Dict[str, Any]]:
    """Read a JSON file and return the parsed data."""
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return json.load(file)
    except FileNotFoundError:
        print(f"Error: File '{file_path}' not found")
        raise


def insert_data(collection: Collection, data: List[Dict[str, Any]],
                batch_size: int = 100) -> Dict[str, Any]:
    """Insert data into collection in batches, skipping if already populated."""
    total_documents = len(data)

    existing_count = collection.count_documents({})
    if existing_count >= total_documents:
        print(f"Collection already has {existing_count} documents, skipping insert")
        return {'total': total_documents, 'inserted': 0, 'skipped': True}

    if existing_count > 0:
        collection.delete_many({})

    inserted_count = 0
    for i in range(0, total_documents, batch_size):
        batch = data[i:i + batch_size]
        try:
            operations = [InsertOne(doc) for doc in batch]
            result = collection.bulk_write(operations, ordered=False)
            inserted_count += result.inserted_count
        except BulkWriteError as e:
            inserted_count += e.details.get('nInserted', 0)
        time.sleep(0.1)

    print(f"Inserted {inserted_count}/{total_documents} documents")
    return {'total': total_documents, 'inserted': inserted_count, 'skipped': False}


def drop_vector_indexes(collection: Collection, vector_field: str) -> None:
    """Drop any existing vector indexes on the specified field."""
    try:
        indexes = list(collection.list_indexes())
        for index in indexes:
            if 'key' in index and vector_field in index['key']:
                if index['key'][vector_field] == 'cosmosSearch':
                    collection.drop_index(index['name'])
                    print(f"Dropped existing vector index: {index['name']}")
    except Exception as e:
        print(f"Warning: Error dropping indexes: {e}")


def perform_vector_search(collection: Collection,
                          azure_openai_client: AzureOpenAI,
                          query_text: str,
                          vector_field: str,
                          model_name: str,
                          top_k: int = 5) -> List[Dict[str, Any]]:
    """Perform vector search using the $search aggregation stage."""
    embedding_response = azure_openai_client.embeddings.create(
        input=[query_text],
        model=model_name
    )
    query_embedding = embedding_response.data[0].embedding

    pipeline = [
        {
            "$search": {
                "cosmosSearch": {
                    "vector": query_embedding,
                    "path": vector_field,
                    "k": top_k
                }
            }
        },
        {
            "$project": {
                "document": "$$ROOT",
                "score": {"$meta": "searchScore"}
            }
        }
    ]

    return list(collection.aggregate(pipeline))


def print_search_results(results: List[Dict[str, Any]], algorithm: str) -> None:
    """Print formatted search results."""
    print(f"\n{'='*60}")
    print(f"  {algorithm} Search Results ({len(results)} found)")
    print(f"{'='*60}")
    for i, result in enumerate(results, 1):
        doc = result.get('document', result)
        name = doc.get('HotelName', doc.get('name', 'Unknown'))
        score = result.get('score', 0)
        print(f"  {i}. {name} (score: {score:.4f})")
    print()

The utilities provide essential functions for:

  • Passwordless authentication to DocumentDB and Azure OpenAI using DefaultAzureCredential.
  • Reading JSON data files with error handling.
  • Batch insertion of documents with DocumentDB's 16 MB payload limit in mind.
  • Formatted display of comparison results showing algorithm performance.

Note

The Python sample configures the DocumentDB connection with retryWrites=false, which is required for DocumentDB vector search operations.

Run the code

Sign in to Azure

az login

Create the output directory, then execute the comparison script to run all 9 combinations:

mkdir -p output
python src/compare_all.py

Expected output:

======================================================================
  Compare All Algorithms — 9 Combinations
  (3 Algorithms × 3 Similarity Metrics)
======================================================================

  Query:  "luxury hotel near the beach"
  Top K:  5

Dropped existing 'hotels' collection (if any)
Loaded 50 documents with embeddings
Inserted 50/50 documents

Generating embedding for query...
Running 9 vector searches...

  Created index 'vector_ivf_cos'
  Created index 'vector_ivf_l2'
  Created index 'vector_ivf_ip'
  Created index 'vector_hnsw_cos'
  Created index 'vector_hnsw_l2'
  Created index 'vector_hnsw_ip'
  Created index 'vector_diskann_cos'
  Created index 'vector_diskann_l2'
  Created index 'vector_diskann_ip'
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| Algorithm   | Metric   | Top 1 Result             |   Score | Top 2 Result      |   Score |   Diff |
+=============+==========+==========================+=========+===================+=========+========+
| IVF         | COS      | Ocean Water Resort & Spa |  0.6184 | Windy Ocean Motel |  0.5057 | 0.1128 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| IVF         | L2       | Ocean Water Resort & Spa |  0.8735 | Windy Ocean Motel |  0.9942 | 0.1207 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| IVF         | IP       | Ocean Water Resort & Spa |  0.6183 | Windy Ocean Motel |  0.5056 | 0.1127 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| HNSW        | COS      | Ocean Water Resort & Spa |  0.6184 | Windy Ocean Motel |  0.5057 | 0.1128 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| HNSW        | L2       | Ocean Water Resort & Spa |  0.8735 | Windy Ocean Motel |  0.9942 | 0.1207 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| HNSW        | IP       | Ocean Water Resort & Spa |  0.6183 | Windy Ocean Motel |  0.5056 | 0.1127 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| DiskANN     | COS      | Ocean Water Resort & Spa |  0.6184 | Windy Ocean Motel |  0.5057 | 0.1128 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| DiskANN     | L2       | Ocean Water Resort & Spa |  0.8735 | Windy Ocean Motel |  0.9942 | 0.1207 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+
| DiskANN     | IP       | Ocean Water Resort & Spa |  0.6183 | Windy Ocean Motel |  0.5056 | 0.1127 |
+-------------+----------+--------------------------+---------+-------------------+---------+--------+

Summary: 9 succeeded, 0 failed

Cleanup: dropped collection 'hotels'

The Diff column shows the score gap between the top-1 and top-2 results. A smaller diff indicates the algorithm found results with more similar relevance scores.

Understanding the results

The comparison table helps you choose the best configuration for your workload:

  • Algorithm: The vector index type used for the query (IVF, HNSW, or DiskANN).
  • Metric: The similarity metric used for scoring (COS, L2, or IP).
  • Top 1 Result and Top 2 Result: The two highest-ranked hotels returned for each algorithm and metric combination.
  • Score: The relevance score for each returned result.
  • Diff: The score gap between the top two results.

Choosing the right algorithm

Important

For production workloads, start with DiskANN on an M30+ cluster. DiskANN supports higher embedding dimensions, uses less cluster memory, and is less likely to require an index redesign as your models evolve.

Use this quick-reference table to select the right algorithm for your workload:

Scenario Algorithm Cluster tier Max dimensions
Dev/test, demos, small datasets IVF M10+ 2,000
Production (default) DiskANN M30+ 16,000
Production (max recall priority) HNSW M30+ 8,000

IVF (inverted file index):

  • Best for: Test environments, demos, and small clusters
  • Pros: Fast to build, low resource requirements
  • Cons: Lower recall compared to graph-based algorithms at scale
  • Tune: Increase numLists for larger datasets, increase nProbes for better recall

DiskANN (disk-based approximate nearest neighbor) — recommended for production:

  • Best for: Production workloads on M30+ clusters
  • Pros: Supports embeddings up to 16,000 dimensions, keeps most index data on disk freeing cluster memory for reads and writes, lighter index updates, easier backups, faster recovery
  • Cons: Requires M30+ cluster tier
  • Tune: Increase maxDegree and lBuild for better accuracy, increase lSearch for better recall
  • Why default: As embedding models evolve (some already exceed 8,000 dimensions), DiskANN avoids costly index redesigns. Its disk-based architecture also means your cluster memory stays available for operational workloads rather than index storage.

HNSW (hierarchical navigable small world):

  • Best for: Production workloads on M30+ clusters where maximum recall is the top priority
  • Pros: Excellent recall, fast queries
  • Cons: Requires M30+ cluster tier, supports embeddings up to 8,000 dimensions (vs 16,000 for DiskANN), higher memory usage since the full graph lives in RAM
  • Tune: Increase m and efConstruction for better index quality, increase efSearch for better recall

Choosing the right similarity function

Function Score meaning Best for
COS (Cosine) Higher = more similar (0–1) Text embeddings (normalized vectors)
L2 (Euclidean) Lower = more similar (distance) When magnitude matters
IP (Inner Product) Higher = more similar Equivalent to COS for normalized vectors

For the text-embedding-3-small model used in this quickstart, COS (cosine similarity) is recommended because OpenAI embeddings are normalized and optimized for cosine similarity.

Troubleshooting

Issue Solution
ServerSelectionTimeoutError Verify that your environment variables are set in the current shell. Ensure your IP is in the DocumentDB firewall rules.
AuthenticationFailed Verify your Microsoft Entra token is valid. Run az login to refresh your credentials.
pymongo.errors.OperationFailure Ensure the database and collection exist. Check that the vector index was created successfully.
ModuleNotFoundError: No module named 'pymongo' Activate your virtual environment and run pip install "pymongo>=4.7".
Empty search results The vector index may take a few minutes to build. Wait 2-3 minutes after index creation, then rerun the script.

Clean up resources

Remove the database using the DocumentDB for VS Code extension:

  1. Install the DocumentDB for VS Code extension.
  2. Connect to your Azure DocumentDB cluster.
  3. Expand the cluster, right-click the Hotels database, and select Drop Database.

Related content

  • Last updated on