Edit

Quickstart: Vector index with Go in Azure DocumentDB

This quickstart walks you through building a Go application that compares all three vector index algorithms (DiskANN, HNSW, and IVF) side by side with different similarity functions to help you choose the best configuration for your workload. The sample uses a hotels 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.

  • Azure Developer CLI (optional). Use azd up to deploy all required Azure resources in one command.

  • Go 1.24 or greater

Create a Go project

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

    mkdir select-algorithm-go
cd select-algorithm-go
code .

  1. Initialize a new Go module:

    go mod init documentdb-vector-samples

    Verify the module was initialized:

    cat go.mod

  1. Install the required packages:

    go get github.com/Azure/azure-sdk-for-go/sdk/azcore@v1.20.0
go get github.com/Azure/azure-sdk-for-go/sdk/azidentity@v1.13.1
go get github.com/openai/openai-go/v3@v3.12.0
go get go.mongodb.org/mongo-driver@v1.17.6
go mod tidy
    • azcore: Core Azure SDK functionality for Go.
    • azidentity: Azure Identity library for passwordless authentication with DefaultAzureCredential.
    • openai-go/v3: OpenAI client library with Azure support to generate embeddings.
    • mongo-driver: Official MongoDB driver for Go to work with DocumentDB.

    Verify the packages are installed:

    go list -m all | grep mongo

  1. Create a src directory:

    mkdir src

Create data file with vectors

  1. Create a new data directory for the hotels data file:

    mkdir data

  1. Download the Hotels_Vector.json raw data file with vectors to your data directory:

    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/Hotels_Vector.json

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 in your current shell session 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. For more information on setting up managed identity and the full range of your authentication options, see Authenticate Go apps to Azure services by using the Azure SDK for Go.

Create code files

Create the main application file:

touch src/main.go

When you're done, the project structure should look like this:

select-algorithm-go/
├── data/
│   └── Hotels_Vector.json
├── src/
│   ├── compare_all.go
│   ├── main.go
│   └── utils.go
└── go.mod

Create the algorithm comparison code

Create the following source files in the src directory.

src/main.go

package main

import (
    "context"
    "fmt"
    "log"
)

func main() {
    fmt.Println("Starting vector algorithm comparison...")

    ctx := context.Background()
    config, err := LoadConfig()
    if err != nil {
        log.Fatalf("Invalid configuration: %v", err)
    }

    fmt.Println("\nInitializing clients with passwordless authentication...")
    mongoClient, azureOpenAIClient, err := GetClientsPasswordless(ctx, config)
    if err != nil {
        log.Fatalf("Failed to initialize clients: %v", err)
    }
    defer mongoClient.Disconnect(ctx)

    err = RunCompareAll(ctx, config, mongoClient, azureOpenAIClient)
    if err != nil {
        log.Fatalf("Compare all failed: %v", err)
    }

    fmt.Println("\nComparison completed successfully!")
}

src/compare_all.go

package main

import (
    "context"
    "fmt"
    "math"
    "strconv"
    "strings"
    "time"

    "github.com/openai/openai-go/v3"
    "go.mongodb.org/mongo-driver/bson"
    "go.mongodb.org/mongo-driver/mongo"
)

// CompareResult holds the result of a single algorithm+metric search
type CompareResult struct {
    Algorithm string
    Metric    string
    Results   []SearchResult
    Top1Name  string
    Top1Score float64
    Top2Name  string
    Top2Score float64
    Error     error
}

// indexSpec defines one of the 9 combinations
type indexSpec struct {
    Algorithm string
    Kind      string
    Metric    string
    IndexName string
    Options   bson.D
}

// RunCompareAll executes all 9 algorithm×metric combinations on a single collection
func RunCompareAll(ctx context.Context, config *Config, dbClient *mongo.Client, aiClient openai.Client) error {
    queryText := getEnvOrDefault("QUERY_TEXT", "luxury hotel near the beach")
    topK, _ := strconv.Atoi(getEnvOrDefault("TOP_K", "5"))

    fmt.Println("\n" + strings.Repeat("=", 70))
    fmt.Println("  COMPARE ALL: 3 Algorithms × 3 Similarity Metrics (9 combinations)")
    fmt.Println(strings.Repeat("=", 70))
    fmt.Printf("Query:  %q\n", queryText)
    fmt.Printf("Top-K:  %d\n", topK)

    // 1. Drop collection for clean comparison, then load data
    database := dbClient.Database(config.DatabaseName)
    collection := database.Collection("hotels")

    // Drop existing collection for a clean comparison
    if err := collection.Drop(ctx); err != nil {
        fmt.Printf("Note: could not drop collection (may not exist): %v\n", err)
    } else {
        fmt.Println("Dropped existing 'hotels' collection")
    }

    // Ensure cleanup on exit
    defer func() {
        fmt.Println("\nCleanup: dropping comparison collection...")
        if dropErr := collection.Drop(ctx); dropErr != nil {
            fmt.Printf("Cleanup warning: %v\n", dropErr)
        } else {
            fmt.Println("Cleanup: dropped collection 'hotels'")
        }
    }()

    fmt.Printf("\nLoading data from %s...\n", config.DataFile)
    data, err := ReadFileReturnJSON(config.DataFile)
    if err != nil {
        return fmt.Errorf("failed to load data: %v", err)
    }

    documentsWithEmbeddings := FilterDocumentsWithEmbeddings(data, config.VectorField)
    if len(documentsWithEmbeddings) == 0 {
        return fmt.Errorf("no documents found with embeddings in field '%s'", config.VectorField)
    }
    fmt.Printf("Loaded %d documents with embeddings\n", len(documentsWithEmbeddings))

    stats, err := PrepareCollection(ctx, collection, documentsWithEmbeddings, config.BatchSize)
    if err != nil {
        return err
    }
    fmt.Printf("Insertion completed: %d inserted, %d failed\n", stats.Inserted, stats.Failed)

    // 2. Generate ONE embedding for the query (reused for all 9 searches)
    fmt.Printf("\nGenerating embedding for query: %q\n", queryText)
    queryEmbedding, err := GenerateEmbedding(ctx, aiClient, queryText, config.ModelName)
    if err != nil {
        return fmt.Errorf("failed to generate query embedding: %v", err)
    }
    fmt.Printf("Embedding generated (%d dimensions)\n", len(queryEmbedding))

    // 3. Define all 9 index specs
    metrics := []string{"COS", "L2", "IP"}
    specs := buildIndexSpecs(config.VectorField, config.Dimensions, metrics)

    // 4. Create→search→drop each index sequentially (DocumentDB only allows one vector index per field)
    fmt.Printf("\nRunning %d vector index comparisons (create→search→drop)...\n", len(specs))
    var results []CompareResult
    successfulComparisons := 0
    failedComparisons := 0

    for _, spec := range specs {
        // Drop all existing vector indexes on this field
        DropVectorIndexes(ctx, collection, config.VectorField)

        // Create this specific index with retry (drop may still be in progress)
        var createErr error
        for attempt := 0; attempt < 3; attempt++ {
            if attempt > 0 {
                time.Sleep(3 * time.Second)
            }
            createErr = createNamedVectorIndex(ctx, collection, config.VectorField, spec)
            if createErr == nil {
                break
            }
        }
        if createErr != nil {
            results = append(results, CompareResult{
                Algorithm: spec.Algorithm,
                Metric:    spec.Metric,
                Error:     createErr,
            })
            failedComparisons++
            fmt.Printf("  ⚠ %s: %v\n", spec.IndexName, createErr)
            continue
        }
        fmt.Printf("  ✓ %s created\n", spec.IndexName)

        // Search using simple cosmosSearch with bounded retry for index readiness.
        searchResults, searchErr := runVectorSearchWithRetry(ctx, collection, queryEmbedding, config.VectorField, topK)

        top1Name, top1Score := extractResult(searchResults, 0)
        top2Name, top2Score := extractResult(searchResults, 1)

        cr := CompareResult{
            Algorithm: spec.Algorithm,
            Metric:    spec.Metric,
            Results:   searchResults,
            Top1Name:  top1Name,
            Top1Score: top1Score,
            Top2Name:  top2Name,
            Top2Score: top2Score,
            Error:     searchErr,
        }
        results = append(results, cr)
        if searchErr != nil {
            failedComparisons++
        } else {
            successfulComparisons++
        }
    }

    // 6. Print comparison table
    fmt.Println()
    printComparisonTable(results)
    fmt.Printf("\nSummary: %d succeeded, %d failed\n", successfulComparisons, failedComparisons)
    if successfulComparisons == 0 {
        return fmt.Errorf("all %d comparisons failed", failedComparisons)
    }

    return nil
}

func runVectorSearchWithRetry(ctx context.Context, collection *mongo.Collection, queryEmbedding []float64, vectorField string, topK int) ([]SearchResult, error) {
    const maxAttempts = 6
    const retryDelay = 2 * time.Second

    var searchResults []SearchResult
    var searchErr error

    for attempt := 1; attempt <= maxAttempts; attempt++ {
        searchResults, searchErr = vectorSearchSimple(ctx, collection, queryEmbedding, vectorField, topK)
        if searchErr == nil {
            if len(searchResults) > 0 {
                return searchResults, nil
            }
            searchErr = fmt.Errorf("search returned no results")
        }

        if attempt < maxAttempts {
            time.Sleep(retryDelay)
        }
    }

    return searchResults, searchErr
}

// buildIndexSpecs creates the 9 index specifications
func buildIndexSpecs(vectorField string, dimensions int, metrics []string) []indexSpec {
    var specs []indexSpec

    type algoConfig struct {
        name    string
        kind    string
        options bson.D
    }

    algos := []algoConfig{
        {"IVF", "vector-ivf", bson.D{{"numLists", 1}}},
        {"HNSW", "vector-hnsw", bson.D{{"m", 16}, {"efConstruction", 64}}},
        {"DiskANN", "vector-diskann", bson.D{{"maxDegree", 32}, {"lBuild", 50}}},
    }

    for _, algo := range algos {
        for _, metric := range metrics {
            metricLower := strings.ToLower(metric)
            opts := bson.D{
                {"kind", algo.kind},
                {"dimensions", dimensions},
                {"similarity", metric},
            }
            for _, o := range algo.options {
                opts = append(opts, o)
            }

            specs = append(specs, indexSpec{
                Algorithm: algo.name,
                Kind:      algo.kind,
                Metric:    metric,
                IndexName: fmt.Sprintf("vector_%s_%s", strings.ToLower(algo.name), metricLower),
                Options:   opts,
            })
        }
    }

    return specs
}

// createNamedVectorIndex creates a single named vector index
func createNamedVectorIndex(ctx context.Context, collection *mongo.Collection, vectorField string, spec indexSpec) error {
    indexCommand := bson.D{
        {"createIndexes", collection.Name()},
        {"indexes", []bson.D{
            {
                {"name", spec.IndexName},
                {"key", bson.D{
                    {vectorField, "cosmosSearch"},
                }},
                {"cosmosSearchOptions", spec.Options},
            },
        }},
    }

    var result bson.M
    err := collection.Database().RunCommand(ctx, indexCommand).Decode(&result)
    if err != nil {
        if strings.Contains(err.Error(), "already exists") || strings.Contains(err.Error(), "IndexAlreadyExists") {
            return nil
        }
        return err
    }
    return nil
}

// vectorSearchSimple performs a vector search using the active vector index
func vectorSearchSimple(ctx context.Context, collection *mongo.Collection, embedding []float64, vectorField string, topK int) ([]SearchResult, error) {
    pipeline := []bson.M{
        {
            "$search": bson.M{
                "cosmosSearch": bson.M{
                    "vector": embedding,
                    "path":   vectorField,
                    "k":      topK,
                },
            },
        },
        {
            "$project": bson.M{
                "document": "$$ROOT",
                "score":    bson.M{"$meta": "searchScore"},
            },
        },
    }

    cursor, err := collection.Aggregate(ctx, pipeline)
    if err != nil {
        return nil, err
    }
    defer cursor.Close(ctx)

    var results []SearchResult
    for cursor.Next(ctx) {
        var result SearchResult
        if err := cursor.Decode(&result); err != nil {
            continue
        }
        results = append(results, result)
    }

    if err := cursor.Err(); err != nil {
        return nil, err
    }

    return results, nil
}

// extractResult returns the name and score of the result at the given index
func extractResult(results []SearchResult, idx int) (string, float64) {
    if idx >= len(results) {
        return "(no results)", 0
    }
    doc := results[idx].Document.(bson.D)
    var name string
    for _, elem := range doc {
        if elem.Key == "HotelName" {
            name = fmt.Sprintf("%v", elem.Value)
            break
        }
    }
    if name == "" {
        name = "Unknown"
    }
    return name, results[idx].Score
}

// printComparisonTable outputs a formatted table of results
func printComparisonTable(results []CompareResult) {
    fmt.Println("┌──────────┬────────┬────────────────────────────┬────────┬────────────────────────────┬────────┬───────┐")
    fmt.Printf("│ %-8s │ %-6s │ %-26s │ %-6s │ %-26s │ %-6s │ %-5s │\n",
        "Algorithm", "Metric", "Top 1 Result", "Score", "Top 2 Result", "Score", "Diff")
    fmt.Println("├──────────┼────────┼────────────────────────────┼────────┼────────────────────────────┼────────┼───────┤")

    for _, r := range results {
        if r.Error != nil {
            fmt.Printf("│ %-8s │ %-6s │ %-26s │ %-6s │ %-26s │ %-6s │ %-5s │\n",
                r.Algorithm, r.Metric, "ERROR", "-", "-", "-", "-")
            continue
        }

        top1 := r.Top1Name
        if len(top1) > 26 {
            top1 = top1[:26]
        }
        top2 := r.Top2Name
        if len(top2) > 26 {
            top2 = top2[:26]
        }
        diff := math.Abs(r.Top1Score - r.Top2Score)

        fmt.Printf("│ %-8s │ %-6s │ %-26s │ %6.4f │ %-26s │ %6.4f │%6.4f │\n",
            r.Algorithm, r.Metric, top1, r.Top1Score, top2, r.Top2Score, diff)
    }

    fmt.Println("└──────────┴────────┴────────────────────────────┴────────┴────────────────────────────┴────────┴───────┘")
}

src/utils.go

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "os"
    "strconv"
    "strings"
    "time"

    "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
    "github.com/Azure/azure-sdk-for-go/sdk/azidentity"
    "github.com/openai/openai-go/v3"
    "github.com/openai/openai-go/v3/azure"
    "github.com/openai/openai-go/v3/option"
    "go.mongodb.org/mongo-driver/bson"
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
)

// Config holds the application configuration
type Config struct {
    ClusterName  string
    DatabaseName string
    DataFile     string
    VectorField  string
    ModelName    string
    Dimensions   int
    BatchSize    int
    Similarity   string
    Algorithm    string
}

// SearchResult represents a search result document
type SearchResult struct {
    Document interface{} `bson:"document"`
    Score    float64     `bson:"score"`
}

// InsertStats holds statistics about data insertion
type InsertStats struct {
    Total    int `json:"total"`
    Inserted int `json:"inserted"`
    Failed   int `json:"failed"`
}

// LoadConfig loads configuration from environment variables
func LoadConfig() (*Config, error) {
    dimensions, err := parsePositiveIntEnv("EMBEDDING_DIMENSIONS", "1536")
    if err != nil {
        return nil, err
    }

    batchSize, err := parsePositiveIntEnv("LOAD_SIZE_BATCH", "100")
    if err != nil {
        return nil, err
    }

    return &Config{
        ClusterName:  getEnvOrDefault("DOCUMENTDB_CLUSTER_NAME", ""),
        DatabaseName: getEnvOrDefault("AZURE_DOCUMENTDB_DATABASENAME", "Hotels"),
        DataFile:     getEnvOrDefault("DATA_FILE_WITH_VECTORS", "data/Hotels_Vector.json"),
        VectorField:  getEnvOrDefault("EMBEDDED_FIELD", "DescriptionVector"),
        ModelName:    getEnvOrDefault("AZURE_OPENAI_EMBEDDING_MODEL", "text-embedding-3-small"),
        Dimensions:   dimensions,
        BatchSize:    batchSize,
        Similarity:   getEnvOrDefault("SIMILARITY", ""),
        Algorithm:    strings.ToLower(getEnvOrDefault("ALGORITHM", "")),
    }, nil
}

func parsePositiveIntEnv(key, defaultValue string) (int, error) {
    value := getEnvOrDefault(key, defaultValue)
    parsedValue, err := strconv.Atoi(value)
    if err != nil {
        return 0, fmt.Errorf("%s must be a positive integer, got %q", key, value)
    }
    if parsedValue <= 0 {
        return 0, fmt.Errorf("%s must be greater than 0, got %q", key, value)
    }
    return parsedValue, nil
}

// getEnvOrDefault returns environment variable value or default if not set
func getEnvOrDefault(key, defaultValue string) string {
    if value := os.Getenv(key); value != "" {
        return value
    }
    return defaultValue
}

// GetClientsPasswordless creates MongoDB and Azure OpenAI clients with passwordless authentication
func GetClientsPasswordless(ctx context.Context, config *Config) (*mongo.Client, openai.Client, error) {
    if config.ClusterName == "" {
        return nil, openai.Client{}, fmt.Errorf("DOCUMENTDB_CLUSTER_NAME environment variable is required")
    }

    // Create Azure credential
    credential, err := azidentity.NewDefaultAzureCredential(nil)
    if err != nil {
        return nil, openai.Client{}, fmt.Errorf("failed to create Azure credential: %v", err)
    }

    // Connect to DocumentDB with OIDC authentication
    mongoURI := fmt.Sprintf("mongodb+srv://%s.global.mongocluster.cosmos.azure.com/", config.ClusterName)

    fmt.Println("Attempting OIDC authentication...")
    mongoClient, err := connectWithOIDC(ctx, mongoURI, credential)
    if err != nil {
        return nil, openai.Client{}, fmt.Errorf("OIDC authentication failed: %v", err)
    }
    fmt.Println("OIDC authentication successful!")

    // Get Azure OpenAI endpoint
    azureOpenAIEndpoint := os.Getenv("AZURE_OPENAI_EMBEDDING_ENDPOINT")
    if azureOpenAIEndpoint == "" {
        return nil, openai.Client{}, fmt.Errorf("AZURE_OPENAI_EMBEDDING_ENDPOINT environment variable is required")
    }

    // Create Azure OpenAI client with credential-based authentication
    openAIClient := openai.NewClient(
        option.WithBaseURL(fmt.Sprintf("%s/openai/v1", azureOpenAIEndpoint)),
        azure.WithTokenCredential(credential))

    return mongoClient, openAIClient, nil
}

// connectWithOIDC attempts to connect using OIDC authentication
func connectWithOIDC(ctx context.Context, mongoURI string, credential *azidentity.DefaultAzureCredential) (*mongo.Client, error) {
    oidcCallback := func(ctx context.Context, args *options.OIDCArgs) (*options.OIDCCredential, error) {
        scope := "https://ossrdbms-aad.database.windows.net/.default"
        fmt.Printf("Getting token with scope: %s\n", scope)
        token, err := credential.GetToken(ctx, policy.TokenRequestOptions{
            Scopes: []string{scope},
        })
        if err != nil {
            return nil, fmt.Errorf("failed to get token with scope %s: %v", scope, err)
        }

        fmt.Printf("Successfully obtained token\n")

        return &options.OIDCCredential{
            AccessToken: token.Token,
        }, nil
    }

    clientOptions := options.Client().
        ApplyURI(mongoURI).
        SetConnectTimeout(30 * time.Second).
        SetServerSelectionTimeout(30 * time.Second).
        SetRetryWrites(false).
        SetAuth(options.Credential{
            AuthMechanism: "MONGODB-OIDC",
            AuthMechanismProperties: map[string]string{
                "TOKEN_RESOURCE": "https://ossrdbms-aad.database.windows.net",
            },
            OIDCMachineCallback: oidcCallback,
        })

    mongoClient, err := mongo.Connect(ctx, clientOptions)
    if err != nil {
        return nil, err
    }

    return mongoClient, nil
}

// ReadFileReturnJSON reads a JSON file and returns the data as a slice of maps
func ReadFileReturnJSON(filePath string) ([]map[string]interface{}, error) {
    file, err := os.ReadFile(filePath)
    if err != nil {
        return nil, fmt.Errorf("error reading file '%s': %v", filePath, err)
    }

    var data []map[string]interface{}
    err = json.Unmarshal(file, &data)
    if err != nil {
        return nil, fmt.Errorf("error parsing JSON in file '%s': %v", filePath, err)
    }

    return data, nil
}

// InsertData inserts data into a MongoDB collection in batches
func InsertData(ctx context.Context, collection *mongo.Collection, data []map[string]interface{}, batchSize int) (*InsertStats, error) {
    totalDocuments := len(data)
    insertedCount := 0
    failedCount := 0

    fmt.Printf("Starting batch insertion of %d documents...\n", totalDocuments)

    for i := 0; i < totalDocuments; i += batchSize {
        end := i + batchSize
        if end > totalDocuments {
            end = totalDocuments
        }

        batch := data[i:end]
        batchNum := (i / batchSize) + 1

        documents := make([]interface{}, len(batch))
        for j, doc := range batch {
            documents[j] = doc
        }

        result, err := collection.InsertMany(ctx, documents, options.InsertMany().SetOrdered(false))
        if err != nil {
            if bulkErr, ok := err.(mongo.BulkWriteException); ok {
                errorCount := len(bulkErr.WriteErrors)
                insertedCount += len(batch) - errorCount
                failedCount += errorCount
                fmt.Printf("Batch %d had errors: %d inserted, %d failed\n", batchNum, len(batch)-errorCount, errorCount)
                for _, writeErr := range bulkErr.WriteErrors {
                    fmt.Printf("  Error: %s\n", writeErr.Message)
                }
            } else {
                failedCount += len(batch)
                fmt.Printf("Batch %d failed completely: %v\n", batchNum, err)
            }
        } else {
            insertedCount += len(result.InsertedIDs)
            fmt.Printf("Batch %d completed: %d documents inserted\n", batchNum, len(result.InsertedIDs))
        }

        time.Sleep(100 * time.Millisecond)
    }

    return &InsertStats{
        Total:    totalDocuments,
        Inserted: insertedCount,
        Failed:   failedCount,
    }, nil
}

// DropVectorIndexes drops existing vector indexes on the specified field
func DropVectorIndexes(ctx context.Context, collection *mongo.Collection, vectorField string) error {
    cursor, err := collection.Indexes().List(ctx)
    if err != nil {
        return fmt.Errorf("could not list indexes: %v", err)
    }
    defer cursor.Close(ctx)

    var vectorIndexes []string
    for cursor.Next(ctx) {
        var index bson.M
        if err := cursor.Decode(&index); err != nil {
            continue
        }

        if key, ok := index["key"].(bson.M); ok {
            if indexType, exists := key[vectorField]; exists && indexType == "cosmosSearch" {
                if name, ok := index["name"].(string); ok {
                    vectorIndexes = append(vectorIndexes, name)
                }
            }
        }
    }

    for _, indexName := range vectorIndexes {
        fmt.Printf("Dropping existing vector index: %s\n", indexName)
        _, err := collection.Indexes().DropOne(ctx, indexName)
        if err != nil {
            fmt.Printf("Warning: Could not drop index %s: %v\n", indexName, err)
        }
    }

    if len(vectorIndexes) > 0 {
        fmt.Printf("Dropped %d existing vector index(es)\n", len(vectorIndexes))
    } else {
        fmt.Println("No existing vector indexes found to drop")
    }

    return nil
}

// PerformVectorSearch performs a vector search using the cosmosSearch aggregation pipeline
func PerformVectorSearch(ctx context.Context, collection *mongo.Collection, client openai.Client, query, vectorField, model string, topK int) ([]SearchResult, error) {
    fmt.Printf("Performing vector search for: '%s'\n", query)

    queryEmbedding, err := GenerateEmbedding(ctx, client, query, model)
    if err != nil {
        return nil, fmt.Errorf("error generating embedding: %v", err)
    }

    pipeline := []bson.M{
        {
            "$search": bson.M{
                "cosmosSearch": bson.M{
                    "vector": queryEmbedding,
                    "path":   vectorField,
                    "k":      topK,
                },
            },
        },
        {
            "$project": bson.M{
                "document": "$$ROOT",
                "score":    bson.M{"$meta": "searchScore"},
            },
        },
    }

    cursor, err := collection.Aggregate(ctx, pipeline)
    if err != nil {
        return nil, fmt.Errorf("error performing vector search: %v", err)
    }
    defer cursor.Close(ctx)

    var results []SearchResult
    for cursor.Next(ctx) {
        var result SearchResult
        if err := cursor.Decode(&result); err != nil {
            fmt.Printf("Warning: Could not decode result: %v\n", err)
            continue
        }
        results = append(results, result)
    }

    if err := cursor.Err(); err != nil {
        return nil, fmt.Errorf("cursor error: %v", err)
    }

    return results, nil
}

// GenerateEmbedding generates an embedding for the given text using Azure OpenAI
func GenerateEmbedding(ctx context.Context, client openai.Client, text, modelName string) ([]float64, error) {
    resp, err := client.Embeddings.New(ctx, openai.EmbeddingNewParams{
        Input: openai.EmbeddingNewParamsInputUnion{
            OfString: openai.String(text),
        },
        Model: modelName,
    })
    if err != nil {
        return nil, fmt.Errorf("failed to generate embedding: %v", err)
    }

    if len(resp.Data) == 0 {
        return nil, fmt.Errorf("no embedding data received")
    }

    embedding := make([]float64, len(resp.Data[0].Embedding))
    for i, v := range resp.Data[0].Embedding {
        embedding[i] = float64(v)
    }

    return embedding, nil
}

// PrintSearchResults prints search results in a formatted way
func PrintSearchResults(results []SearchResult, algorithm string) {
    if len(results) == 0 {
        fmt.Println("No search results found.")
        return
    }

    fmt.Printf("\n%s Search Results (top %d):\n", strings.ToUpper(algorithm), len(results))
    fmt.Println(strings.Repeat("=", 80))

    for i, result := range results {
        doc := result.Document.(bson.D)
        var hotelName string
        for _, elem := range doc {
            if elem.Key == "HotelName" {
                hotelName = fmt.Sprintf("%v", elem.Value)
                break
            }
        }

        fmt.Printf("%d. HotelName: %s, Score: %.4f\n", i+1, hotelName, result.Score)
    }
}

// FilterDocumentsWithEmbeddings returns only documents that contain the vector field
func FilterDocumentsWithEmbeddings(data []map[string]interface{}, vectorField string) []map[string]interface{} {
    var filtered []map[string]interface{}
    for _, doc := range data {
        if _, exists := doc[vectorField]; exists {
            filtered = append(filtered, doc)
        }
    }
    return filtered
}

// PrepareCollection clears existing data and inserts new documents
func PrepareCollection(ctx context.Context, collection *mongo.Collection, data []map[string]interface{}, batchSize int) (*InsertStats, error) {
    fmt.Printf("Preparing collection '%s'...\n", collection.Name())

    deleteResult, err := collection.DeleteMany(ctx, bson.M{})
    if err != nil {
        return nil, fmt.Errorf("failed to clear existing data: %v", err)
    }
    if deleteResult.DeletedCount > 0 {
        fmt.Printf("Cleared %d existing documents from collection\n", deleteResult.DeletedCount)
    }

    stats, err := InsertData(ctx, collection, data, batchSize)
    if err != nil {
        return nil, fmt.Errorf("failed to insert data: %v", err)
    }

    return stats, nil
}

This code provides a vector algorithm comparison application with these key features:

  • Passwordless authentication: Uses DefaultAzureCredential for both Azure OpenAI and DocumentDB via OIDC.
  • Three vector algorithms: Implements DiskANN, HNSW, and IVF with algorithm-specific tuning parameters.
  • Three similarity functions: Supports COS (cosine), L2 (Euclidean), and IP (inner product).
  • Single compare-all entry point: Always runs all 9 algorithm × similarity combinations in one pass.
  • Index lifecycle automation: Creates, queries, and drops each vector index in sequence.
  • Comparison output: Generates a formatted table showing the top two results and score gap for each combination.
  • Production-ready patterns: Includes batched insertion, error handling, and connection pooling.

Note

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

Run the code

After setting the environment variables in your shell session, run the application:

go run ./src/

The application does the following:

  1. Connect to Azure DocumentDB and Azure OpenAI using passwordless authentication
  2. Load the hotel data and insert it into the hotels collection
  3. Generate an embedding for the search query
  4. Run all 9 vector index comparisons by creating, querying, and dropping each index in sequence
  5. Display a comparison table with the top two results and score gap for each combination
  6. Drop the hotels collection during cleanup

Expected output:

======================================================================
  COMPARE ALL: 3 Algorithms × 3 Similarity Metrics (9 combinations)
======================================================================
Query:  "luxury hotel near the beach"
Top-K:  5

Loading data from data/Hotels_Vector.json...
Loaded 50 documents with embeddings
Insertion completed: 50 inserted, 0 failed

Generating embedding for query: "luxury hotel near the beach"
Embedding generated (1536 dimensions)

Running 9 vector index comparisons (create→search→drop)...
  ✓ vector_ivf_cos created
  ✓ vector_ivf_l2 created
  ✓ vector_ivf_ip created
  ✓ vector_hnsw_cos created
  ✓ vector_hnsw_l2 created
  ✓ vector_hnsw_ip created
  ✓ vector_diskann_cos created
  ✓ vector_diskann_l2 created
  ✓ vector_diskann_ip created

┌──────────┬────────┬────────────────────────────┬────────┬────────────────────────────┬────────┬───────┐
│ Algorithm│ Metric │ Top 1 Result               │ Score  │ Top 2 Result               │ Score  │ Diff  │
├──────────┼────────┼────────────────────────────┼────────┼────────────────────────────┼────────┼───────┤
│ IVF      │ COS    │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
│ IVF      │ L2     │ Ocean Water Resort & Spa   │ 0.8736 │ Windy Ocean Motel          │ 0.9943 │ 0.1208│
│ IVF      │ IP     │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
│ HNSW     │ COS    │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
│ HNSW     │ L2     │ Ocean Water Resort & Spa   │ 0.8736 │ Windy Ocean Motel          │ 0.9943 │ 0.1208│
│ HNSW     │ IP     │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
│ DiskANN  │ COS    │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
│ DiskANN  │ L2     │ Ocean Water Resort & Spa   │ 0.8736 │ Windy Ocean Motel          │ 0.9943 │ 0.1208│
│ DiskANN  │ IP     │ Ocean Water Resort & Spa   │ 0.6184 │ Windy Ocean Motel          │ 0.5056 │ 0.1128│
└──────────┴────────┴────────────────────────────┴────────┴────────────────────────────┴────────┴───────┘

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 shows how different algorithms perform on the same dataset with the same query:

  • Algorithm: DiskANN, HNSW, or IVF
  • Metric: The similarity metric (COS, L2, or IP)
  • Top 1 Result: The highest-ranked hotel for that algorithm and metric
  • Score: The relevance score for the corresponding result
  • Top 2 Result: The second-highest-ranked hotel for that algorithm and metric
  • 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
server selection error Verify your environment variables are set correctly. Ensure your IP is in the DocumentDB firewall rules.
authentication failed Verify your Microsoft Entra token is valid. Run az login to refresh your credentials.
go: module not found Run go mod tidy to resolve dependencies.
Build errors Ensure Go 1.24+ is installed. Run go version to check.
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