向量数据库对比与集成：PGVector / Qdrant / Chroma / Milvus / Weaviate 选型与实战

-- PostgreSQL 中启用扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 创建向量表
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(768),  -- BGE-M3 维度
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMP DEFAULT NOW()
);

-- HNSW 索引（近似最近邻，推荐用于生产）
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- IVFFlat 索引（内存更小，适合低写入场景）
-- CREATE INDEX ON documents
-- USING ivfflat (embedding vector_cosine_ops)
-- WITH (lists = 100);

import psycopg
from pgvector.psycopg import register_vector
import numpy as np

class PGVectorStore:
    def __init__(self, dsn: str = "postgresql://user:pass@localhost/db"):
        self.dsn = dsn
        self._ensure_setup()

    def _ensure_setup(self):
        with psycopg.connect(self.dsn) as conn:
            conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
            conn.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id SERIAL PRIMARY KEY,
                    content TEXT NOT NULL,
                    embedding vector(768),
                    metadata JSONB DEFAULT '{}',
                    created_at TIMESTAMP DEFAULT NOW()
                )
            """)
            conn.commit()

    def upsert(self, contents: list[str], embeddings: np.ndarray, metadatas: list[dict] = None):
        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            for i, (content, emb) in enumerate(zip(contents, embeddings)):
                meta = metadatas[i] if metadatas else {}
                conn.execute("""
                    INSERT INTO documents (content, embedding, metadata)
                    VALUES (%s, %s, %s)
                    ON CONFLICT (id) DO UPDATE SET
                        content = EXCLUDED.content,
                        embedding = EXCLUDED.embedding,
                        metadata = EXCLUDED.metadata
                """, (content, emb.tolist(), json.dumps(meta)))
            conn.commit()

    def search(self, query_embedding: np.ndarray, top_k: int = 5, filters: dict = None) -> list[dict]:
        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            where_clauses = []
            params = [query_embedding.tolist(), top_k]

            if filters:
                for key, value in filters.items():
                    where_clauses.append(f"metadata->>'{key}' = %s")
                    params.append(value)

            where_sql = " AND ".join(where_clauses) if where_clauses else "TRUE"

            sql = f"""
                SELECT id, content, metadata,
                       1 - (embedding <=> %s) AS cosine_similarity
                FROM documents
                WHERE {where_sql}
                ORDER BY embedding <=> %s
                LIMIT %s
            """
            # 参数重排：第一个 %s (where) + 第二个 %s (order) + %s (limit)
            params = [query_embedding.tolist()] + params[2:] + [query_embedding.tolist(), top_k]

            with conn.cursor() as cur:
                cur.execute(sql, params)
                rows = cur.fetchall()

            return [
                {
                    "id": r[0],
                    "content": r[1],
                    "metadata": r[2],
                    "score": float(r[3]),
                }
                for r in rows
            ]

    # 混合搜索：全文检索 + 向量
    def hybrid_search(self, query_text: str, query_embedding: np.ndarray, top_k: int = 5) -> list[dict]:
        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            sql = """
                SELECT id, content, metadata,
                       0.5 * ts_rank(to_tsvector('english', content), plainto_tsquery('english', %s))
                       + 0.5 * (1 - (embedding <=> %s)) AS hybrid_score
                FROM documents
                ORDER BY hybrid_score DESC
                LIMIT %s
            """
            with conn.cursor() as cur:
                cur.execute(sql, (query_text, query_embedding.tolist(), top_k))
                rows = cur.fetchall()
            return [{"id": r[0], "content": r[1], "metadata": r[2], "score": float(r[3])} for r in rows]

-- 距离运算符
SELECT embedding <-> '[1,2,3]' AS l2_distance;        -- L2 欧氏距离
SELECT embedding <=> '[1,2,3]' AS cosine_distance;     -- 余弦距离 (1 - cosine_similarity)
SELECT embedding <#> '[1,2,3]' AS inner_product;      -- 内积

-- 组合过滤 + 向量排序
SELECT * FROM documents
WHERE metadata->>'category' = 'tech'
ORDER BY embedding <=> $1
LIMIT 10;

-- 批量相似度（k-NN join）
SELECT d1.id, d2.id,
       d1.embedding <=> d2.embedding AS distance
FROM documents d1
JOIN documents d2 ON d1.id != d2.id
WHERE d1.id = 1
ORDER BY distance
LIMIT 5;

docker run -p 6333:6333 -p 6334:6334 \
  -v $(pwd)/qdrant_storage:/qdrant/storage \
  qdrant/qdrant:latest

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue,
    HnswConfigDiff, OptimizersConfigDiff,
)
import numpy as np

class QdrantManager:
    def __init__(self, host: str = "localhost", port: int = 6333, grpc_port: int = 6334):
        # REST + gRPC 双协议客户端
        self.client = QdrantClient(host=host, port=port, grpc_port=grpc_port, prefer_grpc=True)

    def create_collection(
        self,
        name: str,
        dim: int = 768,
        distance: Distance = Distance.COSINE,
        on_disk: bool = False,
    ):
        """创建集合，支持 HNSW 参数自定义。"""
        self.client.recreate_collection(
            collection_name=name,
            vectors_config=VectorParams(
                size=dim,
                distance=distance,
                on_disk=on_disk,  # 内存映射到磁盘，降低内存占用
            ),
            hnsw_config=HnswConfigDiff(
                m=16,
                ef_construct=100,
                full_scan_threshold=10000,
            ),
            optimizers_config=OptimizersConfigDiff(
                indexing_threshold=20000,  # 自动触发索引的阈值
            ),
        )

    def upsert(
        self,
        collection: str,
        ids: list[int],
        vectors: np.ndarray,
        payloads: list[dict],
        batch_size: int = 100,
    ):
        """批量插入数据。"""
        points = [
            PointStruct(id=i, vector=v.tolist(), payload=p)
            for i, v, p in zip(ids, vectors, payloads)
        ]
        self.client.upsert(collection_name=collection, points=points, batch_size=batch_size)

    def search(
        self,
        collection: str,
        query_vector: np.ndarray,
        top_k: int = 10,
        filters: dict = None,
        with_payload: bool = True,
        score_threshold: float = None,
    ) -> list[dict]:
        """语义搜索，支持元数据过滤和分数阈值。"""
        query_filter = None
        if filters:
            conditions = [
                FieldCondition(key=k, match=MatchValue(value=v))
                for k, v in filters.items()
            ]
            query_filter = Filter(must=conditions)

        results = self.client.search(
            collection_name=collection,
            query_vector=query_vector.tolist(),
            limit=top_k,
            query_filter=query_filter,
            with_payload=with_payload,
            score_threshold=score_threshold,
        )
        return [
            {
                "id": r.id,
                "score": r.score,
                "payload": r.payload,
            }
            for r in results
        ]

    # 多向量搜索（同一文档多个 Embedding）
    def create_multivec_collection(self, name: str):
        self.client.recreate_collection(
            collection_name=name,
            vectors_config={
                "title": VectorParams(size=768, distance=Distance.COSINE),
                "content": VectorParams(size=768, distance=Distance.COSINE),
                "summary": VectorParams(size=384, distance=Distance.COSINE),
            },
        )

    def search_multivec(self, collection: str, vector_name: str, query: np.ndarray, top_k: int = 5):
        return self.client.search(
            collection_name=collection,
            query_vector=(vector_name, query.tolist()),
            limit=top_k,
        )

    # 推荐 API（基于正/负样本）
    def recommend(
        self,
        collection: str,
        positive_ids: list[int],
        negative_ids: list[int] = None,
        top_k: int = 5,
    ) -> list[dict]:
        results = self.client.recommend(
            collection_name=collection,
            positive=positive_ids,
            negative=negative_ids or [],
            limit=top_k,
        )
        return [{"id": r.id, "score": r.score, "payload": r.payload} for r in results]

from qdrant_client.models import SparseVector

# Qdrant 1.10+ 支持稀疏向量
class SparseVectorSearch:
    def create_sparse_collection(self, name: str):
        self.client.recreate_collection(
            collection_name=name,
            sparse_vectors_config={
                "splade": models.SparseVectorParams(
                    index=models.SparseIndexParams(
                        full_scan_threshold=1000,
                    )
                )
            },
        )

    def upsert_sparse(self, collection: str, id: int, sparse_vec: dict):
        """sparse_vec: {indices: [0, 5, 10], values: [0.8, 0.3, 0.9]}"""
        self.client.upsert(
            collection_name=collection,
            points=[PointStruct(
                id=id,
                vector={"splade": SparseVector(**sparse_vec)},
            )],
        )

import chromadb
from chromadb.config import Settings

class ChromaStore:
    def __init__(self, persist_dir: str = "./chroma_db"):
        # 本地持久化模式
        self.client = chromadb.PersistentClient(
            path=persist_dir,
            settings=Settings(anonymized_telemetry=False),
        )

    def get_or_create_collection(self, name: str):
        return self.client.get_or_create_collection(
            name=name,
            metadata={"hnsw:space": "cosine"},
        )

    def add(self, collection_name: str, texts: list[str], embeddings: list[list[float]], metadatas: list[dict] = None, ids: list[str] = None):
        coll = self.get_or_create_collection(collection_name)
        if ids is None:
            ids = [f"doc_{i}" for i in range(len(texts))]
        coll.add(
            documents=texts,
            embeddings=embeddings,
            metadatas=metadatas,
            ids=ids,
        )

    def query(self, collection_name: str, query_embedding: list[float], n_results: int = 5, where: dict = None) -> dict:
        coll = self.get_or_create_collection(collection_name)
        return coll.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
            include=["documents", "metadatas", "distances"],
        )

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

class MilvusStore:
    def __init__(self, host: str = "localhost", port: str = "19530"):
        connections.connect(alias="default", host=host, port=port)

    def create_collection(self, name: str, dim: int = 768):
        if utility.has_collection(name):
            return Collection(name)

        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="metadata", dtype=DataType.JSON),
        ]
        schema = CollectionSchema(fields, description="Document embeddings")
        collection = Collection(name, schema)

        # HNSW 索引
        index_params = {
            "index_type": "HNSW",
            "metric_type": "COSINE",
            "params": {"M": 16, "efConstruction": 64},
        }
        collection.create_index("embedding", index_params)
        return collection

    def insert(self, collection_name: str, embeddings: list[list[float]], contents: list[str], metadatas: list[dict]):
        coll = Collection(collection_name)
        coll.insert([embeddings, contents, metadatas])
        coll.flush()

    def search(self, collection_name: str, query_embedding: list[float], top_k: int = 5) -> list:
        coll = Collection(collection_name)
        coll.load()
        results = coll.search(
            data=[query_embedding],
            anns_field="embedding",
            param={"metric_type": "COSINE", "params": {"ef": 64}},
            limit=top_k,
            output_fields=["content", "metadata"],
        )
        return results[0]

import weaviate

class WeaviateStore:
    def __init__(self, url: str = "http://localhost:8080"):
        self.client = weaviate.Client(url=url)

    def create_schema(self, class_name: str = "Document"):
        schema = {
            "class": class_name,
            "vectorizer": "none",  # 手动提供向量
            "properties": [
                {"name": "content", "dataType": ["text"]},
                {"name": "category", "dataType": ["text"]},
            ],
        }
        if not self.client.schema.contains({"class": class_name}):
            self.client.schema.create_class(schema)

    def insert(self, class_name: str, objects: list[dict], vectors: list[list[float]]):
        with self.client.batch as batch:
            batch.batch_size = 100
            for obj, vec in zip(objects, vectors):
                batch.add_data_object(obj, class_name, vector=vec)

    def search(self, class_name: str, query_vec: list[float], top_k: int = 5) -> list[dict]:
        result = (
            self.client.query
            .get(class_name, ["content", "category"])
            .with_near_vector({"vector": query_vec})
            .with_limit(top_k)
            .with_additional(["distance"])
            .do()
        )
        return result["data"]["Get"][class_name]