向量数据库选型与检索优化：从 ANN 索引到生产运维

import numpy as np

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """余弦相似度：文本语义检索的默认度量。"""
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

# 内积在归一化后等价于余弦相似度
a = np.random.randn(1024); a = a / np.linalg.norm(a)
b = np.random.randn(1024); b = b / np.linalg.norm(b)
print(np.dot(a, b))  # 归一化后的内积 = 余弦相似度

# Qdrant HNSW 建索引配置
from qdrant_client import QdrantClient, models

client = QdrantClient(url="http://localhost:6333")

client.create_collection(
    collection_name="kb_embeddings",
    vectors_config=models.VectorParams(
        size=1024,                      # 向量维度
        distance=models.Distance.COSINE,
    ),
    hnsw_config=models.HnswConfigDiff(
        m=16,             # 出度：16-64 常见，默认 16
        ef_construct=100, # 建图宽度：默认 100
    ),
)

# 查询时通过 search_params 控制 efSearch
client.search(
    collection_name="kb_embeddings",
    query_vector=query_vec,
    limit=20,
    search_params=models.SearchParams(ef=256, hnsw_ef=256),
)

# Milvus 中使用 IVF_SQ8 或 IVF_PQ 索引
from pymilvus import connections, CollectionSchema, FieldSchema, DataType, Collection

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
]
schema = CollectionSchema(fields, description="knowledge base")
collection = Collection("kb", schema)

index_params = {
    "index_type": "IVF_SQ8",      # 分区 + 标量量化（SQ8）
    "metric_type": "COSINE",
    "params": {"nlist": 4096},    # 聚类中心数
}
collection.create_index("vector", index_params)

# 查询：nprobe 越大召回越高、延迟越高
search_params = {"metric_type": "COSINE", "params": {"nprobe": 32}}
results = collection.search(
    data=[query_vec], anns_field="vector",
    param=search_params, limit=20, output_fields=["text"],
)

# pgvector 最小示例
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1024)          -- 向量列
);

-- 建 HNSW 索引（pgvector 0.5+）
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);

-- 检索 Top-5
SELECT id, content, 1 - (embedding <=> $1) AS similarity
FROM documents
ORDER BY embedding <=> $1           -- <=> 为余弦距离
LIMIT 5;

import time
import numpy as np
from qdrant_client import QdrantClient, models

def recall_at_k(ground_truth, preds, k=10):
    gt = set(ground_truth[:k])
    hit = len(gt.intersection(preds[:k]))
    return hit / k

def tune_efsearch(client, queries, ef_range, k=10):
    for ef in ef_range:
        latencies, recalls = [], []
        for q, gt in queries:
            t0 = time.perf_counter()
            res = client.search(collection_name="kb", query_vector=q,
                                limit=k, search_params=models.SearchParams(ef=ef))
            latencies.append((time.perf_counter() - t0) * 1000)
            recalls.append(recall_at_k(gt, [r.id for r in res], k))
        print(f"efSearch={ef}: recall={np.mean(recalls):.3f}, p95={np.percentile(latencies, 95):.1f}ms")

# Qdrant 预过滤 + 向量检索
from qdrant_client import QdrantClient, models

client.search(
    collection_name="kb",
    query_vector=query_vec,
    query_filter=models.Filter(
        must=[
            models.FieldCondition(
                key="tenant_id", match=models.MatchValue(value="tenant-123")),
            models.FieldCondition(
                key="updated_at", range=models.Range(
                    gte=int(start_ts), lte=int(end_ts))),
        ]
    ),
    limit=20,
    search_params=models.SearchParams(hnsw_ef=256),
)

// Milvus 中布尔表达式过滤
{"bool": {
  "must": [
    {"term": {"tenant_id": "tenant-123"}},
    {"range": {"updated_at": {"gte": 1727280000}}}
  ]
}}

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("BAAI/bge-m3", device="cuda")

def batch_encode(texts: list[str], batch_size: int = 64):
    embeddings = model.encode(
        texts, batch_size=batch_size, normalize_embeddings=True,
        convert_to_numpy=True,
    )
    return embeddings.astype(np.float16)  # 存储用 FP16 减半显存

docs = ["向量数据库选型要点", "RAG 检索优化实践"]
vecs = batch_encode(docs)
print(vecs.shape, vecs.dtype)  # (2, 1024) float16

# Milvus 分布式容量规划参考
clusters:
  query_nodes:
    replica: 2          # 查询副本数
    resources:
      cpu: "16"
      memory: "64Gi"
    limits:
      maxMemory: "48Gi" # 留给索引 + 检索缓冲
  index_nodes:
    cpu: "8"
    memory: "32Gi"

# Qdrant 定期清理过期 collection 片段的简单脚本（cron）
#!/bin/bash
COLLECTION="news_kb"
BEFORE_TS=$(date -d "-90 days" +%s)
curl -s -X POST http://localhost:6333/collections/$COLLECTION/points/delete \
  -H "Content-Type: application/json" \
  -d "{\"filter\": {\"must\": [{\"key\": \"published_at\", \"range\": {\"lt\": $BEFORE_TS}}]}, \"wait\": true}"