AI 原生数据库设计：向量检索与结构化查询的融合

┌──────────────────────────────────────────────────┐
│                  SQL / GraphQL / REST API         │
├──────────────────────────────────────────────────┤
│              Query Optimizer & Planner             │
│   ┌─────────────┐   ┌─────────────┐   ┌────────┐ │
│   │Scalar Filter│   │ Vector Index│   │Sort/Max│ │
│   │  Planner    │   │  Planner    │   │ Planner│ │
│   └─────────────┘   └─────────────┘   └────────┘ │
├──────────────────────────────────────────────────┤
│           Execution Engine (Vectorized)           │
├──────────────────────────────────────────────────┤
│  ┌──────────┐  ┌──────────┐  ┌──────────────┐   │
│  │ Vector   │  │  Row     │  │  Column      │   │
│  │  Index   │  │  Store   │  │   Store      │   │
│  │(HNSW/IVF)│  │ (Pilosa) │  │(Arrow/Parquet)│  │
│  └──────────┘  └──────────┘  └──────────────┘   │
├──────────────────────────────────────────────────┤
│              Distributed Storage Layer            │
│   ┌────────────────────────────────────────────┐ │
│   │ S3 / HDFS / Local SSD / Remote Memory      │ │
│   └────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────┘

from dataclasses import dataclass, field
from typing import Optional
import numpy as np

@dataclass
class VectorRecord:
    id: str
    vector: np.ndarray                    # 原始向量（float32/float16)
    metadata: dict = field(default_factory=dict)
    version: int = 0                      # MVCC 版本号
    deleted: bool = False

    def to_bytes(self) -> bytes:
        """序列化：4字节维度 + 原始向量 + 元数据JSON + 版本"""
        dim = len(self.vector)
        vec_bytes = self.vector.astype(np.float32).tobytes()
        meta_bytes = json.dumps(self.metadata).encode("utf-8")
        header = struct.pack("<i", dim)
        version_bytes = struct.pack("<i", self.version)
        deleted_byte = b"\x01" if self.deleted else b"\x00"
        return header + version_bytes + deleted_byte + vec_bytes + meta_bytes

    @classmethod
    def from_bytes(cls, data: bytes) -> "VectorRecord":
        dim, version = struct.unpack("<ii", data[:8])
        deleted = data[8] == 1
        vec_bytes = data[9:9 + dim * 4]
        meta_bytes = data[9 + dim * 4:].decode("utf-8")
        return cls(
            id="",  # id 需要从索引层获取
            vector=np.frombuffer(vec_bytes, dtype=np.float32),
            metadata=json.loads(meta_bytes),
            version=version,
            deleted=deleted
        )

class VectorPage:
    """向量分页存储单元"""
    def __init__(self, page_size: int = 4096):
        self.page_size = page_size
        self.records: dict[str, bytes] = {}  # id -> serialized bytes
        self.current_size = 0

    def add(self, record: VectorRecord) -> bool:
        record_bytes = record.to_bytes()
        if self.current_size + len(record_bytes) > self.page_size:
            return False
        self.records[record.id] = record_bytes
        self.current_size += len(record_bytes)
        return True

import heapq
import random
from dataclasses import dataclass, field
from typing import Callable

@dataclass
class HNSWNode:
    id: str
    vector: np.ndarray
    level: int
    neighbors: dict[int, list[str]] = field(default_factory=dict)
    # level -> [neighbor_ids]

class HNSWIndex:
    """
    HNSW 实现：多层跳表 + NSW 图
    - 最底层（L0）包含所有向量
    - 上层节点稀疏，构成高速公路
    - 搜索时从顶层贪心下降
    - 构建时通过概率分布决定节点层数
    """
    def __init__(self, dim: int, m: int = 16, ef_construction: int = 200,
                 m0: int = None, max_level: int = None):
        self.dim = dim
        self.m = m                      # 每个节点每层最多邻居数
        self.ef_construction = ef_construction  # 构建时搜索宽度
        self.m0 = m0 or m
        self.max_level = max_level or int(np.log2(1_000_000))
        self.nodes: dict[str, HNSWNode] = {}
        self.entry_point: Optional[str] = None
        self.top_level = 0
        # 各层节点列表（用于计算层）
        self.level_nodes: dict[int, set[str]] = {i: set() for i in range(self.max_level + 1)}

    def _calc_level(self) -> int:
        """指数分布决定节点层数，越高层节点越少"""
        return min(int(-np.log(random.random()) * self.L), self.max_level)

    def _distance(self, v1: np.ndarray, v2: np.ndarray) -> float:
        return float(np.linalg.norm(v1 - v2))

    def _search_layer(self, query: np.ndarray, ep_id: str, ef: int, level: int) -> list[tuple[str, float]]:
        """贪婪搜索单层 NSW 图"""
        visited = {ep_id}
        candidates: list[tuple[float, str]] = [(0.0, ep_id)]
        results: list[tuple[float, str]] = [(float("inf"), "")]

        while candidates:
            dist, current_id = heapq.heappop(candidates)
            # 最差结果不够好则退出
            if dist > results[0][0]:
                break
            for neighbor_id in self.nodes[current_id].neighbors.get(level, []):
                if neighbor_id not in visited:
                    visited.add(neighbor_id)
                    d = self._distance(query, self.nodes[neighbor_id].vector)
                    if d < results[0][0] or len(results) < ef:
                        heapq.heappush(results, (d, neighbor_id))
                        heapq.heappush(candidates, (d, neighbor_id))
                        results = heapq.nlargest(ef, results)

        results = [(d, i) for d, i in results if i]
        return sorted(results)[:ef]

    def insert(self, id: str, vector: np.ndarray):
        """插入向量到 HNSW"""
        node = HNSWNode(id=id, vector=vector, level=self._calc_level())
        self.nodes[id] = node

        if self.entry_point is None:
            self.entry_point = id
            self.top_level = node.level
            return

        # 从顶层向下搜索找到各层最近邻
        ep_id = self.entry_point
        for level in range(self.top_level, node.level, -1):
            results = self._search_layer(vector, ep_id, 1, level)
            ep_id = results[0][1] if results else ep_id

        # 自上而下插入各层
        for level in range(min(node.level, self.top_level), -1, -1):
            neighbors = self._search_layer(vector, ep_id, self.ef_construction, level)
            # 选择 top-m 个邻居
            neighbor_ids = [nid for _, nid in neighbors[:self.m if level > 0 else self.m0]]
            node.neighbors[level] = neighbor_ids
            self.level_nodes[level].add(id)
            # 双向链接
            for nid in neighbor_ids:
                if len(self.nodes[nid].neighbors.setdefault(level, [])) < (self.m if level > 0 else self.m0):
                    self.nodes[nid].neighbors[level].append(id)
            ep_id = neighbor_ids[0] if neighbor_ids else ep_id

        self.top_level = max(self.top_level, node.level)

    def search(self, query: np.ndarray, k: int = 10, ef: int = 100) -> list[tuple[str, float]]:
        """近似最近邻搜索"""
        if not self.nodes:
            return []

        ep_id = self.entry_point
        for level in range(self.top_level, 1, -1):
            results = self._search_layer(query, ep_id, 1, level)
            ep_id = results[0][1] if results else ep_id

        results = self._search_layer(query, ep_id, ef, 0)
        return sorted(results)[:k]

class IVFPQIndex:
    """IVF-PQ：先聚类后搜索，适合超大规模数据"""
    def __init__(self, n_clusters: int = 1024, n_subvectors: int = 8, n_probe: int = 64):
        self.n_clusters = n_clusters
        self.n_subvectors = n_subvectors
        self.n_probe = n_probe
        self.centroids: np.ndarray = None  # (n_clusters, dim)
        self.subCentroids: np.ndarray = None  # (n_subvectors, n_sub_clusters, sub_dim)
        self.posting_lists: dict[int, list[str]] = {i: [] for i in range(n_clusters)}
        self.id_to_cluster: dict[str, int] = {}

    def build(self, vectors: np.ndarray, ids: list[str]):
        """训练阶段：计算 PQ 码本"""
        dim = vectors.shape[1]
        sub_dim = dim // self.n_subvectors
        # K-Means 聚类
        self.centroids, _ = self._kmeans(vectors, self.n_clusters)
        # 分配向量到聚类
        for i, (vec, vid) in enumerate(zip(vectors, ids)):
            cluster = self._assign_cluster(vec)
            self.posting_lists[cluster].append(vid)
            self.id_to_cluster[vid] = cluster
        # PQ 训练
        residuals = vectors - self.centroids[self.id_to_cluster[id_]]
        self.subCentroids = self._train_pq(residuals, self.n_subvectors, sub_dim)

    def _kmeans(self, vectors, k):
        """简化 K-Means"""
        from sklearn.cluster import MiniBatchKMeans
        km = MiniBatchKMeans(n_clusters=k, random_state=42)
        km.fit(vectors)
        return km.cluster_centers_, km.labels_

    def _train_pq(self, residuals, n_sub, sub_dim):
        """训练产品量化码本"""
        subCentroids = np.zeros((n_sub, 256, sub_dim), dtype=np.float32)
        for i in range(n_sub):
            chunk = residuals[:, i*sub_dim:(i+1)*sub_dim]
            km = MiniBatchKMeans(n_clusters=256, random_state=42)
            km.fit(chunk)
            subCentroids[i] = km.cluster_centers_
        return subCentroids

    def search(self, query: np.ndarray, k: int = 10) -> list[tuple[str, float]]:
        """搜索"""
        # 1. 找到最近的 n_probe 个聚类中心
        distances = np.linalg.norm(self.centroids - query, axis=1)
        nearest_clusters = np.argsort(distances)[:self.n_probe]
        # 2. 在这些聚类中搜索
        candidates = []
        for c in nearest_clusters:
            for vid in self.posting_lists[c]:
                d = self._distance_pq(query, vid)
                candidates.append((vid, d))
        return sorted(candidates, key=lambda x: x[1])[:k]

from enum import Enum

class FilterType(str, Enum):
    EQ = "eq"
    GT = "gt"
    LT = "lt"
    IN = "in"
    CONTAINS = "contains"
    RANGE = "range"

@dataclass
class ScalarCondition:
    field: str
    op: FilterType
    value: any

@dataclass
class HybridQuery:
    vector: np.ndarray          # 查询向量
    vector_top_k: int           # 向量检索数量
    scalar_filters: list[ScalarCondition]  # 标量过滤条件
    post_filter: bool = True    # True=先向量后过滤, False=先过滤后向量

class HybridQueryPlanner:
    """混合查询计划器"""

    def plan(self, query: HybridQuery) -> dict:
        has_scalar = len(query.scalar_filters) > 0
        selectivity = self._estimate_selectivity(query.scalar_filters)

        if has_scalar and selectivity < 0.01:
            # 高选择性：先过滤再向量检索
            return {
                "strategy": "filter_then_search",
                "steps": [
                    {"type": "scalar_scan", "filters": query.scalar_filters},
                    {"type": "vector_search", "within_ids": "<from_previous>", "top_k": query.vector_top_k}
                ]
            }
        elif has_scalar and selectivity < 0.3:
            # 中等选择性：HNSW + 运行时过滤
            return {
                "strategy": "indexed_with_filter",
                "ef": max(query.vector_top_k * 10, 100),
                "steps": [{"type": "hnsw_search", "post_filters": query.scalar_filters, "top_k": query.vector_top_k}]
            }
        else:
            # 低选择性或无过滤：直接向量检索
            return {
                "strategy": "direct_search",
                "steps": [{"type": "hnsw_search", "top_k": query.vector_top_k}]
            }

    def _estimate_selectivity(self, filters: list[ScalarCondition]) -> float:
        """估算过滤器的选择性（返回行数比例）"""
        if not filters:
            return 1.0
        # 简化估算，实际需要统计信息
        return min(0.5 ** len(filters), 0.5)

class ShardingStrategy(Enum):
    STRATIFIED = "stratified"    # 按向量维度范围分片（不适合向量）
    ROUND_ROBIN = "round_robin"  # 轮询分片
    SEMANTIC = "semantic"        # 按向量语义分片（用聚类中心分桶）
    CONSISTENT_HASH = "hash"     # 按 ID 哈希

class SemanticSharding:
    """语义分片：向量相似的数据在同一分片"""
    def __init__(self, n_shards: int = 8, samples: int = 10000):
        self.n_shards = n_shards
        self.centroids: np.ndarray = None  # (n_shards, dim)
        self.sample_vectors: np.ndarray = None
        self.sample_ids: list[str] = None

    def initialize(self, sample_vectors: np.ndarray, sample_ids: list[str]):
        """用采样数据确定分片边界"""
        self.sample_vectors = sample_vectors
        self.sample_ids = sample_ids
        # K-Means 确定分片中心
        self.centroids, _ = self._kmeans(sample_vectors, self.n_shards)

    def get_shard(self, vector: np.ndarray) -> int:
        distances = np.linalg.norm(self.centroids - vector, axis=1)
        return int(np.argmin(distances))

    def get_shard_by_id(self, id: str) -> int:
        """通过 ID 找到对应分片"""
        idx = self.sample_ids.index(id)
        shard = self.get_shard(self.sample_vectors[idx])
        return shard

    async def distributed_search(self, query: np.ndarray, k: int, shards: dict) -> list[tuple[str, float]]:
        """并行在所有分片上搜索"""
        tasks = [
            shard.search(query, k * 2)  # 每分片多取一些
            for shard in shards.values()
        ]
        all_results = await asyncio.gather(*tasks)
        # 合并并重新排序
        merged = []
        for results in all_results:
            merged.extend(results)
        return sorted(merged, key=lambda x: x[1])[:k]

class PerformanceTuning:
    @staticmethod
    def recommend_index_params(vectors: np.ndarray, n: int) -> dict:
        """根据数据规模推荐索引参数"""
        if n < 10000:
            return {"type": "flat", "metric": "cosine"}
        elif n < 1_000_000:
            return {"type": "hnsw", "m": 16, "efConstruction": 200, "ef": 100}
        elif n < 100_000_000:
            return {"type": "hnsw", "m": 32, "efConstruction": 512, "ef": 256}
        else:
            return {"type": "ivf_pq", "n_clusters": max(1024, int(np.sqrt(n))),
                    "n_subvectors": 8, "n_probe": 64}

    @staticmethod
    def tune_ef(query: np.ndarray, index, target_recall: float = 0.95) -> int:
        """二分搜索找到满足召回率的最小 ef"""
        low, high = 10, 1000
        while low < high:
            mid = (low + high) // 2
            # ground_truth 需要预先计算
            results = index.search(query, k=10, ef=mid)
            recall = index.compute_recall(results, ground_truth)
            if recall >= target_recall:
                high = mid
            else:
                low = mid + 1
        return low

class MetadataIndex:
    """元数据倒排索引（独立于向量索引）"""
    def __init__(self):
        self.inverted_index: dict[str, dict[any, set[str]]] = {}  # field -> value -> {doc_ids}

    def add(self, doc_id: str, metadata: dict):
        for field, value in metadata.items():
            if field not in self.inverted_index:
                self.inverted_index[field] = {}
            if isinstance(value, list):
                for v in value:
                    self.inverted_index[field].setdefault(v, set()).add(doc_id)
            else:
                self.inverted_index[field].setdefault(value, set()).add(doc_id)

    def filter_ids(self, conditions: list[ScalarCondition]) -> set[str]:
        result_sets = []
        for cond in conditions:
            if cond.op == FilterType.EQ:
                result_sets.append(self.inverted_index.get(cond.field, {}).get(cond.value, set()))
            elif cond.op == FilterType.IN:
                s = set()
                for v in cond.value:
                    s |= self.inverted_index.get(cond.field, {}).get(v, set())
                result_sets.append(s)
            elif cond.op == FilterType.RANGE:
                # 范围查询需遍历
                all_ids = set()
                for v, ids in self.inverted_index.get(cond.field, {}).items():
                    if cond.value[0] <= v <= cond.value[1]:
                        all_ids |= ids
                result_sets.append(all_ids)
        # AND 条件取交集
        if not result_sets:
            return set()
        return set.intersection(*result_sets)