Posted inAI

Building Large-Scale Semantic Search with Milvus and Python: Data Partitioning and Query Optimization for Millions of Vectors

AI tutorial - IT technology blog
AI tutorial - IT technology blog

The Real-World Bottleneck: Scaling Semantic Search Past 100k Vectors

Everything worked fine in our local sandbox. But the moment we pushed 5 million enterprise documents into our staging cluster, our search pipeline collapsed. Simple similarity lookups that used to return in 15ms suddenly stalled with p99 latencies spiking past 1,200ms, quickly running our worker nodes out of memory.

Most tutorials demonstrate semantic search on toy datasets of a few thousand vectors. Flat brute-force scans and unoptimized in-memory indices work well at that size. Production systems tell a different story. Scaling to millions of records introduces three stubborn friction points:

  • Memory bloat: A 1536-dimensional float32 embedding (like OpenAI’s text-embedding-3-small) takes roughly 6.14 KB per document. For 10 million vectors, raw storage alone demands over 60 GB of RAM before factoring in graph links and indexing overhead.
  • Wasted collection scans: When an application serves multi-tenant data, searching an entire monolithic vector space burns CPU cycles on millions of irrelevant embeddings that the user never had permission to view anyway.
  • Index build starvation: Constructing dense proximity graphs while simultaneously ingesting high-throughput real-time updates spikes CPU utilization to 100%, causing query timeouts for active users.

Surviving these scale hurdles requires a purpose-built vector database, automated physical partitioning, and an index configuration tuned for real-world memory and query constraints.

Core Concepts: Why Milvus, Partition Keys, and HNSW Matter

Milvus is an open-source, distributed vector database built around storage-compute disaggregation. Rather than bolting vector arithmetic onto relational engines, Milvus splits collections into discrete physical segments and computes similarity using hardware-level SIMD and GPU acceleration.

1. Physical Partitions vs. Partition Keys

Creating manual partitions in Milvus works for simple setups, but managing thousands of individual partition tags inside application code quickly turns into an operational nightmare. Milvus 2.2+ resolved this with Partition Keys.

You flag a scalar field—such as tenant_id or workspace_id—with is_partition_key=True. Milvus then automatically hashes that scalar value and assigns records into a fixed pool of underlying physical partitions (typically 64 to 128 partitions). At query time, passing a tenant filter routes the vector search directly to relevant segments, eliminating 90%+ of unnecessary vector comparisons.

2. Index Trade-Offs: IVF_FLAT vs. HNSW

Choosing an Approximate Nearest Neighbor (ANN) algorithm comes down to balancing latency budgets against hardware costs:

  • IVF_FLAT (Inverted File Flat): Clusters the vector space into Voronoi cells. It adds almost zero extra RAM overhead on top of the raw embeddings. However, query latency degrades sharply as throughput grows, because achieving high recall requires scanning multiple clusters (higher nprobe).
  • HNSW (Hierarchical Navigable Small World): Constructs a multi-layer graph where top layers provide fast routing and bottom layers capture local neighborhood detail. Query times consistently stay under 10ms with over 98% recall. The trade-off? It requires 1.2x to 1.5x additional RAM over raw vector size to store node connection lists.

When high query throughput and sub-20ms p99 SLAs matter, HNSW combined with Partition Keys is the most reliable production setup.

Hands-On Practice: Building and Tuning the Pipeline

Here is a complete setup: deploying Milvus Standalone, defining a partition-key schema in Python, building an HNSW index, and executing filtered vector searches.

Step 1: Running Milvus Standalone with Docker Compose

Save this configuration as docker-compose.yml:

version: '3.5'

services:
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
    volumes:
      - ./volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls=http://0.0.0.0:2379 --data-dir=/etcd

  minio:
    container_name: milvus-minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    volumes:
      - ./volumes/minio:/minio_data
    command: minio server /minio_data

  standalone:
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.3.4
    command: ["milvus", "run", "standalone"]
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ./volumes/milvus:/var/lib/milvus
    ports:
      - "19530:19530"
      - "9091:9091"
    depends_on:
      - "etcd"
      - "minio"

Spin up the stack:

docker compose up -d
docker compose ps

Step 2: Defining the Schema with Partition Keys

Install the required Python dependencies:

pip install pymilvus sentence-transformers numpy

Create milvus_semantic_search.py to initialize your connection and define the schema:

import numpy as np
from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    utility
)

# Connect to Milvus server
connections.connect(alias="default", host="localhost", port="19530")

COLLECTION_NAME = "enterprise_docs"
DIMENSION = 384  # Standard output dimension for all-MiniLM-L6-v2

if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)

# Configure fields with tenant_id designated as the partition key
fields = [
    FieldSchema(name="doc_id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="tenant_id", dtype=DataType.VARCHAR, max_length=64, is_partition_key=True),
    FieldSchema(name="doc_category", dtype=DataType.VARCHAR, max_length=64),
    FieldSchema(name="raw_text", dtype=DataType.VARCHAR, max_length=2048),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]

schema = CollectionSchema(fields=fields, description="Tenant-partitioned semantic knowledge base")
collection = Collection(name=COLLECTION_NAME, schema=schema, num_partitions=64)
print(f"Collection '{COLLECTION_NAME}' created with 64 hash partitions.")

Step 3: Configuring the HNSW Index

Two primary parameters govern HNSW build performance and recall accuracy:

  • M (Connectivity limit, typically 16 to 64): Maximum outgoing links per node. Higher numbers yield better accuracy on noisy embeddings but increase memory usage. Setting M=16 is the standard sweet spot.
  • efConstruction (Search depth during construction, 64 to 512): Controls the exploration depth when inserting vectors. Setting efConstruction=200 strikes a clean balance between build time and index quality.
index_params = {
    "metric_type": "COSINE",
    "index_type": "HNSW",
    "params": {
        "M": 16,
        "efConstruction": 200
    }
}

collection.create_index(field_name="vector", index_params=index_params)
print("HNSW Index built successfully.")

Step 4: Ingestion and Partition-Isolated Search

Let’s embed sample documents, insert them across multiple tenant IDs, and run a targeted query:

from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("all-MiniLM-L6-v2")

sample_docs = [
    {"tenant": "tenant_a", "cat": "network", "text": "Configuring BGP routing on enterprise edge gateways."},
    {"tenant": "tenant_a", "cat": "security", "text": "Setting up MFA authentication policies for remote engineers."},
    {"tenant": "tenant_b", "cat": "database", "text": "PostgreSQL connection pool optimization and query caching strategies."},
    {"tenant": "tenant_b", "cat": "network", "text": "VPC peering latency debugging in multi-region clusters."}
]

# Prepare data batches
tenant_ids = [d["tenant"] for d in sample_docs]
categories = [d["cat"] for d in sample_docs]
texts = [d["text"] for d in sample_docs]
vectors = encoder.encode(texts).tolist()

# Ingest data and flush to disk
collection.insert([
    tenant_ids,
    categories,
    texts,
    vectors
])
collection.flush()

# Load collection into RAM for searching
collection.load()

# Execute search scoped to tenant_b
query_text = "How do I resolve multi-region VPC networking lag?"
query_vector = encoder.encode([query_text]).tolist()

search_params = {
    "metric_type": "COSINE",
    "params": {"ef": 64}  # ef controls runtime search depth vs latency
}

results = collection.search(
    data=query_vector,
    anns_field="vector",
    param=search_params,
    limit=2,
    expr="tenant_id == 'tenant_b'",
    output_fields=["tenant_id", "doc_category", "raw_text"]
)

for hits in results:
    for hit in hits:
        print(f"Score: {hit.distance:.4f} | Tenant: {hit.entity.get('tenant_id')} | Category: {hit.entity.get('doc_category')} | Text: {hit.entity.get('raw_text')}")

Operational Best Practices

Running millions of vectors in production requires disciplined resource planning and query hygiene. Keep these rules in mind:

  • Always include the partition key in search filters: Omitting the key triggers a global scan across every segment. In a 10-million-vector benchmark, scoping queries by tenant partition key dropped p99 latency from 185ms down to 14ms.
  • Tune runtime ef dynamically: Match ef to your SLA. Setting ef=32 or ef=64 typically yields 95%+ recall while keeping search latency under 15ms. Reserve ef=128 for tasks where recall precision outranks speed.
  • Calculate memory capacity ahead of time: Size your RAM requirements with the formula: Total RAM ≈ Vector Count × (Dimension × 4 bytes + M × 8 bytes) × 1.35 (incorporating a 35% safety margin for OS buffers and query execution).
  • Avoid calling collection.flush() on single inserts: Flushing after every small payload generates hundreds of tiny segment files, causing severe segment fragmentation and degraded query performance. Instead, buffer inserts into batches of 2,000 to 10,000 rows.

Combining Milvus partition keys with well-parameterized HNSW indexes keeps search latencies predictable as data scales from thousands to tens of millions of records.

Share:
Tags: