soupdb-vector 0.1.0__py3-none-any.whl

soupdb_vector/__init__.py

@@ -0,0 +1,3 @@
+"""SoupDB vector plugin package: IVF/HNSW indices and multimedia features."""
+
+__all__ = ["indexes", "multimedia", "plugin"]

soupdb_vector/indexes/__init__.py

@@ -0,0 +1,18 @@
+"""Vector indexes for Parte 4: paged vector store, brute force, IVF and HNSW."""
+
+from soupdb_vector.indexes.brute_force import BruteForceIndex
+from soupdb_vector.indexes.hnsw import HNSWIndex
+from soupdb_vector.indexes.ivf import IVFIndex
+from soupdb_vector.indexes.metrics import COSINE, EUCLIDEAN, INNER_PRODUCT, METRICS
+from soupdb_vector.indexes.store import VectorStore
+
+__all__ = [
+    "COSINE",
+    "EUCLIDEAN",
+    "INNER_PRODUCT",
+    "METRICS",
+    "BruteForceIndex",
+    "HNSWIndex",
+    "IVFIndex",
+    "VectorStore",
+]

soupdb_vector/indexes/_record_file.py

@@ -0,0 +1,154 @@
+"""Array of fixed-size records over pages of a DiskManager.
+
+Page 0 is a header (magic, record size, record count and a free area for the
+owner's metadata); record ``i`` lives at page ``1 + i // per_page``. Fixed
+size lets a record be addressed by arithmetic instead of a slot directory,
+which is what vectors and HNSW adjacency lists need.
+"""
+
+import struct
+from collections.abc import Iterator
+
+from engine.storage.buffer_manager import BufferManager
+from engine.storage.disk_manager import DiskManager
+
+_HEADER = struct.Struct("<4sHIQ")
+_VERSION = 1
+
+
+def peek_meta(buffer_manager: BufferManager, magic: bytes, layout: struct.Struct) -> tuple:
+    """Read the owner's metadata of an existing file before knowing its record size."""
+    frame = buffer_manager.pin(0)
+    try:
+        stored_magic = _HEADER.unpack_from(frame, 0)[0]
+        values = layout.unpack_from(frame, _HEADER.size)
+    finally:
+        buffer_manager.unpin(0)
+    if stored_magic != magic:
+        raise ValueError(f"not a {magic!r} file")
+    return values
+
+
+class FixedRecordFile:
+    """Append-only array of fixed-size records, read and written in place."""
+
+    def __init__(
+        self,
+        disk_manager: DiskManager,
+        buffer_manager: BufferManager,
+        *,
+        record_size: int,
+        magic: bytes,
+    ) -> None:
+        page_size = disk_manager.page_size
+        if record_size <= 0 or record_size > page_size:
+            raise ValueError(
+                f"record of {record_size} bytes does not fit in a {page_size}-byte page"
+            )
+        if len(magic) != 4:
+            raise ValueError("magic must be exactly 4 bytes")
+
+        self._disk = disk_manager
+        self._buffer = buffer_manager
+        self._record_size = record_size
+        self._per_page = page_size // record_size
+        self._magic = magic
+
+        if disk_manager.page_count == 0:
+            disk_manager.allocate_page()
+            self._count = 0
+            self._write_header()
+        else:
+            frame = buffer_manager.pin(0)
+            try:
+                stored_magic, version, stored_size, count = _HEADER.unpack_from(frame, 0)
+            finally:
+                buffer_manager.unpin(0)
+            if stored_magic != magic or version != _VERSION:
+                raise ValueError(f"{disk_manager.path} is not a {magic!r} file")
+            if stored_size != record_size:
+                raise ValueError(f"record size on disk is {stored_size}, expected {record_size}")
+            self._count = count
+
+    def __len__(self) -> int:
+        return self._count
+
+    @property
+    def per_page(self) -> int:
+        return self._per_page
+
+    @property
+    def meta_capacity(self) -> int:
+        return self._disk.page_size - _HEADER.size
+
+    def read_meta(self, layout: struct.Struct) -> tuple:
+        frame = self._buffer.pin(0)
+        try:
+            return layout.unpack_from(frame, _HEADER.size)
+        finally:
+            self._buffer.unpin(0)
+
+    def write_meta(self, layout: struct.Struct, *values: object) -> None:
+        if layout.size > self.meta_capacity:
+            raise ValueError("metadata does not fit in the header page")
+        frame = self._buffer.pin(0)
+        layout.pack_into(frame, _HEADER.size, *values)
+        self._buffer.unpin(0, dirty=True)
+
+    def append(self, layout: struct.Struct, *values: object) -> int:
+        index = self._count
+        page_id, offset = self._locate(index)
+        if page_id == self._disk.page_count:
+            self._disk.allocate_page()
+        frame = self._buffer.pin(page_id)
+        layout.pack_into(frame, offset, *values)
+        self._buffer.unpin(page_id, dirty=True)
+        self._count += 1
+        self._write_header()
+        return index
+
+    def unpack(self, index: int, layout: struct.Struct, extra: int = 0) -> tuple:
+        """Unpack ``layout`` at byte ``extra`` of record ``index``."""
+        page_id, offset = self._locate(self._check(index))
+        frame = self._buffer.pin(page_id)
+        try:
+            return layout.unpack_from(frame, offset + extra)
+        finally:
+            self._buffer.unpin(page_id)
+
+    def pack(self, index: int, layout: struct.Struct, *values: object, extra: int = 0) -> None:
+        page_id, offset = self._locate(self._check(index))
+        frame = self._buffer.pin(page_id)
+        layout.pack_into(frame, offset + extra, *values)
+        self._buffer.unpin(page_id, dirty=True)
+
+    def scan(self, layout: struct.Struct) -> Iterator[tuple[int, tuple]]:
+        """Yield (index, values) for every record, pinning each page once."""
+        index = 0
+        page_id = 1
+        while index < self._count:
+            frame = self._buffer.pin(page_id)
+            try:
+                in_page = min(self._per_page, self._count - index)
+                for slot in range(in_page):
+                    yield index + slot, layout.unpack_from(frame, slot * self._record_size)
+            finally:
+                self._buffer.unpin(page_id)
+            index += in_page
+            page_id += 1
+
+    def flush(self) -> None:
+        self._buffer.flush_all()
+
+    def _locate(self, index: int) -> tuple[int, int]:
+        return 1 + index // self._per_page, (index % self._per_page) * self._record_size
+
+    def _check(self, index: int) -> int:
+        if not 0 <= index < self._count:
+            raise IndexError(f"record {index} out of range (count {self._count})")
+        return index
+
+    def _write_header(self) -> None:
+        frame = self._buffer.pin(0)
+        _HEADER.pack_into(frame, 0, self._magic, _VERSION, self._record_size, self._count)
+        self._buffer.unpin(0, dirty=True)

soupdb_vector/indexes/brute_force.py

@@ -0,0 +1,68 @@
+"""Exact k-NN by scanning every stored vector.
+
+It is the sequential baseline of the comparison 2.4.4 and the oracle the
+recall of IVF and HNSW is measured against.
+"""
+
+import heapq
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from engine.common.rid import RID
+from engine.indexes.base import RankedIndex
+from soupdb_vector.indexes.constants import DEFAULT_METRIC, VECTOR_BUFFER_CAPACITY
+from soupdb_vector.indexes.metrics import norm, score_fn
+from soupdb_vector.indexes.store import VectorStore
+
+
+class BruteForceIndex(RankedIndex):
+    """RankedIndex that scores every vector; memory is O(k), not O(N)."""
+
+    def __init__(self, store: VectorStore) -> None:
+        self._store = store
+
+    @classmethod
+    def open(
+        cls,
+        path: str | Path,
+        dim: int | None = None,
+        *,
+        page_size: int = 4096,
+        buffer_capacity: int = VECTOR_BUFFER_CAPACITY,
+    ) -> "BruteForceIndex":
+        return cls(
+            VectorStore.open(path, dim, page_size=page_size, buffer_capacity=buffer_capacity)
+        )
+
+    @property
+    def store(self) -> VectorStore:
+        return self._store
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def insert(self, vector: Sequence[float], rid: RID) -> None:
+        self._store.append(vector, rid)
+
+    def top_k(
+        self, query: Sequence[float], k: int, *, metric: str = DEFAULT_METRIC
+    ) -> list[tuple[RID, float]]:
+        if k < 0:
+            raise ValueError(f"k must be non-negative, got {k}")
+        self._store.check_dim(query)
+        scorer = score_fn(metric)
+        query_norm = norm(query)
+        scored = (
+            (scorer(query, query_norm, vector, vector_norm), rid)
+            for _vid, rid, vector, vector_norm in self._store.scan()
+        )
+        best = heapq.nlargest(k, scored, key=_by_score)
+        return [(rid, value) for value, rid in best]
+
+    def close(self) -> None:
+        self._store.close()
+
+
+def _by_score(pair: tuple[float, Any]) -> float:
+    return pair[0]

soupdb_vector/indexes/constants.py

@@ -0,0 +1,34 @@
+"""Tunable parameters of the vector indexes.
+
+Kept apart so the benchmarks can vary them without touching the structures.
+"""
+
+#: Métrica por defecto cuando la consulta no indica ``metric``.
+DEFAULT_METRIC = "euclidean"
+
+#: Páginas en el buffer pool de cada archivo cuando el índice abre sus propios managers.
+VECTOR_BUFFER_CAPACITY = 64
+
+#: Número de listas (centroides) del IVF; la regla habitual es ~sqrt(N).
+IVF_NLIST = 32
+
+#: Listas que se visitan por consulta.
+IVF_NPROBE = 4
+
+#: Iteraciones de Lloyd al entrenar los centroides.
+IVF_KMEANS_ITERATIONS = 20
+
+#: Tope de la muestra con la que se entrena K-Means; acota la memoria del entrenamiento.
+IVF_TRAIN_SAMPLE = 10_000
+
+#: Vecinos por nodo en las capas superiores de HNSW; la capa 0 guarda el doble.
+HNSW_M = 16
+
+#: Tamaño de la lista de candidatos al insertar.
+HNSW_EF_CONSTRUCTION = 100
+
+#: Tamaño de la lista de candidatos al consultar.
+HNSW_EF_SEARCH = 64
+
+#: Tope de niveles del grafo; con M=16 un nivel 16 es astronómicamente improbable.
+HNSW_MAX_LEVEL = 16

soupdb_vector/indexes/hnsw.py

@@ -0,0 +1,372 @@
+"""Hierarchical Navigable Small World graph persisted on pages (Malkov & Yashunin).
+
+The index is a directory with three files:
+
+- ``vectors.db``: a :class:`VectorStore`; the vid of a vector is its node id.
+- ``layer0.db``: one fixed record per node with its top level, where its upper
+  layers start in ``upper.db`` and its up to ``2*M`` neighbors in layer 0.
+- ``upper.db``: one fixed record per (node, layer >= 1) with up to ``M``
+  neighbors. A node of level ``L`` owns ``L`` consecutive records.
+
+Fixed-size records make every neighbor list addressable by arithmetic, so a
+hop of the greedy search costs one page pin.
+"""
+
+import heapq
+import math
+import random
+import struct
+from collections.abc import Sequence
+from pathlib import Path
+
+from engine.common.rid import RID
+from engine.indexes.base import RankedIndex
+from engine.storage.buffer_manager import BufferManager
+from engine.storage.disk_manager import DiskManager
+from soupdb_vector.indexes._record_file import FixedRecordFile, peek_meta
+from soupdb_vector.indexes.constants import (
+    DEFAULT_METRIC,
+    HNSW_EF_CONSTRUCTION,
+    HNSW_EF_SEARCH,
+    HNSW_M,
+    HNSW_MAX_LEVEL,
+    VECTOR_BUFFER_CAPACITY,
+)
+from soupdb_vector.indexes.metrics import METRICS, ScoreFn, norm, score_fn
+from soupdb_vector.indexes.store import VectorStore
+
+_LAYER0_MAGIC = b"SHN0"
+_UPPER_MAGIC = b"SHNU"
+#: M, ef_construction, metric, seed, entry point, max level.
+_META = struct.Struct("<HHBQqi")
+_NODE_HEAD = struct.Struct("<BqH")
+_UPPER_HEAD = struct.Struct("<H")
+_NO_NODE = -1
+
+
+class HNSWIndex(RankedIndex):
+    """HNSW over pages; every node, vector and neighbor list lives on disk."""
+
+    def __init__(
+        self,
+        path: str | Path,
+        dim: int | None = None,
+        *,
+        m: int = HNSW_M,
+        ef_construction: int = HNSW_EF_CONSTRUCTION,
+        metric: str = DEFAULT_METRIC,
+        seed: int = 0,
+        page_size: int = 4096,
+        buffer_capacity: int = VECTOR_BUFFER_CAPACITY,
+    ) -> None:
+        """Open (or create) the index stored in the directory ``path``.
+
+        ``m``, ``ef_construction``, ``metric`` and ``seed`` only apply when the
+        index is created; reopening uses the values stored on disk.
+        """
+        directory = Path(path)
+        directory.mkdir(parents=True, exist_ok=True)
+        self._closed = False
+        self._disks: list[DiskManager] = []
+
+        vectors_disk = self._disk(directory / "vectors.db", page_size)
+        self._store = VectorStore(
+            vectors_disk, BufferManager(vectors_disk, capacity=buffer_capacity), dim
+        )
+        layer0_disk = self._disk(directory / "layer0.db", page_size)
+        layer0_buffer = BufferManager(layer0_disk, capacity=buffer_capacity)
+
+        if layer0_disk.page_count > 0:
+            m, ef_construction, metric_code, seed, entry, max_level = peek_meta(
+                layer0_buffer, _LAYER0_MAGIC, _META
+            )
+            metric = METRICS[metric_code]
+        else:
+            if m < 2:
+                raise ValueError(f"m must be at least 2, got {m}")
+            score_fn(metric)
+            entry, max_level = _NO_NODE, -1
+
+        self._m = m
+        self._m0 = 2 * m
+        self._ef_construction = ef_construction
+        self._metric = metric
+        self._seed = seed
+        self._entry = entry
+        self._max_level = max_level
+        self._insert_cache: dict[int, tuple[tuple[float, ...], float]] | None = None
+        # mL = 1/ln(M): así cada nivel tiene ~1/M de los nodos del nivel de abajo.
+        self._level_mult = 1.0 / math.log(m)
+
+        self._layer0_record = struct.Struct(f"{_NODE_HEAD.format}{self._m0}q")
+        self._upper_record = struct.Struct(f"{_UPPER_HEAD.format}{m}q")
+        self._layer0 = FixedRecordFile(
+            layer0_disk,
+            layer0_buffer,
+            record_size=self._layer0_record.size,
+            magic=_LAYER0_MAGIC,
+        )
+        upper_disk = self._disk(directory / "upper.db", page_size)
+        self._upper = FixedRecordFile(
+            upper_disk,
+            BufferManager(upper_disk, capacity=buffer_capacity),
+            record_size=self._upper_record.size,
+            magic=_UPPER_MAGIC,
+        )
+        self._write_meta()
+
+    @property
+    def dim(self) -> int:
+        return self._store.dim
+
+    @property
+    def metric(self) -> str:
+        return self._metric
+
+    @property
+    def max_level(self) -> int:
+        return self._max_level
+
+    @property
+    def disk_managers(self) -> tuple[DiskManager, ...]:
+        return tuple(self._disks)
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def insert(self, vector: Sequence[float], rid: RID) -> None:
+        # Una inserción vuelve a puntuar los mismos vecinos al elegir y al podar listas;
+        # la caché vive solo durante esta inserción, así que la memoria no crece con N.
+        self._insert_cache = {}
+        try:
+            self._insert(vector, rid)
+        finally:
+            self._insert_cache = None
+
+    def _insert(self, vector: Sequence[float], rid: RID) -> None:
+        node = self._store.append(vector, rid)
+        query, query_norm = self._vector(node)
+        level = self._random_level(node)
+
+        upper_start = len(self._upper)
+        for _ in range(level):
+            self._upper.append(self._upper_record, 0, *self._padded([], self._m))
+        self._layer0.append(self._layer0_record, level, upper_start, 0, *self._padded([], self._m0))
+
+        if self._entry == _NO_NODE:
+            self._entry, self._max_level = node, level
+            self._write_meta()
+            return
+
+        scorer = score_fn(self._metric)
+        entry_points = [(self._score(scorer, query, query_norm, self._entry), self._entry)]
+        for layer in range(self._max_level, level, -1):
+            entry_points = self._search_layer(scorer, query, query_norm, entry_points, 1, layer)
+
+        for layer in range(min(level, self._max_level), -1, -1):
+            candidates = self._search_layer(
+                scorer, query, query_norm, entry_points, self._ef_construction, layer
+            )
+            capacity = self._capacity(layer)
+            chosen = self._select_neighbors(scorer, candidates, capacity)
+            self._write_neighbors(node, layer, [neighbor for _s, neighbor in chosen])
+            for link_score, neighbor in chosen:
+                self._link(scorer, neighbor, node, link_score, layer)
+            entry_points = candidates
+
+        if level > self._max_level:
+            self._entry, self._max_level = node, level
+            self._write_meta()
+
+    def top_k(
+        self,
+        query: Sequence[float],
+        k: int,
+        *,
+        metric: str | None = None,
+        ef: int = HNSW_EF_SEARCH,
+    ) -> list[tuple[RID, float]]:
+        if k < 0:
+            raise ValueError(f"k must be non-negative, got {k}")
+        if ef <= 0:
+            raise ValueError(f"ef must be positive, got {ef}")
+        self._store.check_dim(query)
+        scorer = score_fn(metric or self._metric)
+        if k == 0 or self._entry == _NO_NODE:
+            return []
+
+        query_norm = norm(query)
+        entry_points = [(self._score(scorer, query, query_norm, self._entry), self._entry)]
+        for layer in range(self._max_level, 0, -1):
+            entry_points = self._search_layer(scorer, query, query_norm, entry_points, 1, layer)
+        found = self._search_layer(scorer, query, query_norm, entry_points, max(ef, k), 0)
+        return [(self._store.rid(node), value) for value, node in found[:k]]
+
+    def neighbors(self, node: int, layer: int) -> list[int]:
+        """Neighbor ids of a node in a layer; exposed for tests and the report."""
+        if layer == 0:
+            values = self._layer0.unpack(node, self._layer0_record)
+            return list(values[3 : 3 + values[2]])
+        values = self._upper.unpack(self._upper_index(node, layer), self._upper_record)
+        return list(values[1 : 1 + values[0]])
+
+    def level(self, node: int) -> int:
+        return self._layer0.unpack(node, _NODE_HEAD)[0]
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._write_meta()
+        self._store.flush()
+        self._layer0.flush()
+        self._upper.flush()
+        for disk in self._disks:
+            disk.close()
+
+    # --- internals ---------------------------------------------------------
+
+    def _disk(self, path: Path, page_size: int) -> DiskManager:
+        disk = DiskManager(path, page_size=page_size)
+        self._disks.append(disk)
+        return disk
+
+    def _write_meta(self) -> None:
+        self._layer0.write_meta(
+            _META,
+            self._m,
+            self._ef_construction,
+            METRICS.index(self._metric),
+            self._seed,
+            self._entry,
+            self._max_level,
+        )
+
+    def _random_level(self, node: int) -> int:
+        # Semilla derivada del nodo: el nivel es reproducible aunque se cierre y reabra.
+        uniform = random.Random(self._seed * 1_000_003 + node).random()
+        return min(int(-math.log(1.0 - uniform) * self._level_mult), HNSW_MAX_LEVEL)
+
+    def _capacity(self, layer: int) -> int:
+        return self._m0 if layer == 0 else self._m
+
+    def _upper_index(self, node: int, layer: int) -> int:
+        level, start, _count = self._layer0.unpack(node, _NODE_HEAD)
+        if not 1 <= layer <= level:
+            raise ValueError(f"node {node} has no layer {layer}")
+        return start + layer - 1
+
+    def _padded(self, neighbors: list[int], capacity: int) -> list[int]:
+        return neighbors + [_NO_NODE] * (capacity - len(neighbors))
+
+    def _write_neighbors(self, node: int, layer: int, neighbors: list[int]) -> None:
+        if layer == 0:
+            level, start, _count = self._layer0.unpack(node, _NODE_HEAD)
+            self._layer0.pack(
+                node,
+                self._layer0_record,
+                level,
+                start,
+                len(neighbors),
+                *self._padded(neighbors, self._m0),
+            )
+            return
+        record = self._upper_index(node, layer)
+        self._upper.pack(
+            record, self._upper_record, len(neighbors), *self._padded(neighbors, self._m)
+        )
+
+    def _vector(self, node: int) -> tuple[tuple[float, ...], float]:
+        cache = self._insert_cache
+        if cache is None:
+            return self._store.vector(node)
+        cached = cache.get(node)
+        if cached is None:
+            cached = cache[node] = self._store.vector(node)
+        return cached
+
+    def _score(
+        self, scorer: ScoreFn, query: Sequence[float], query_norm: float, node: int
+    ) -> float:
+        vector, vector_norm = self._vector(node)
+        return scorer(query, query_norm, vector, vector_norm)
+
+    def _search_layer(
+        self,
+        scorer: ScoreFn,
+        query: Sequence[float],
+        query_norm: float,
+        entry_points: list[tuple[float, int]],
+        ef: int,
+        layer: int,
+    ) -> list[tuple[float, int]]:
+        """Best-first search in one layer; returns up to ef (score, node), best first."""
+        visited = {node for _score, node in entry_points}
+        # candidates: max-heap por score (negado); results: min-heap con los ef mejores.
+        candidates = [(-value, node) for value, node in entry_points]
+        heapq.heapify(candidates)
+        results = list(entry_points)
+        heapq.heapify(results)
+        while len(results) > ef:
+            heapq.heappop(results)
+
+        while candidates:
+            negated, node = heapq.heappop(candidates)
+            if -negated < results[0][0] and len(results) >= ef:
+                break
+            for neighbor in self.neighbors(node, layer):
+                if neighbor in visited:
+                    continue
+                visited.add(neighbor)
+                value = self._score(scorer, query, query_norm, neighbor)
+                if len(results) < ef or value > results[0][0]:
+                    heapq.heappush(candidates, (-value, neighbor))
+                    heapq.heappush(results, (value, neighbor))
+                    if len(results) > ef:
+                        heapq.heappop(results)
+        return sorted(results, reverse=True)
+
+    def _select_neighbors(
+        self, scorer: ScoreFn, candidates: list[tuple[float, int]], capacity: int
+    ) -> list[tuple[float, int]]:
+        """Neighbor-selection heuristic (algorithm 4) with keepPrunedConnections.
+
+        A candidate is kept only if it is closer to the base node than to every
+        neighbor already kept; this favors links in different directions, which
+        is what keeps clustered data navigable. Pruned ones fill the remaining slots.
+        """
+        if len(candidates) <= capacity:
+            return list(candidates)
+        kept: list[tuple[float, int]] = []
+        kept_vectors: list[tuple[Sequence[float], float]] = []
+        pruned: list[tuple[float, int]] = []
+        for value, node in candidates:
+            if len(kept) >= capacity:
+                break
+            vector, vector_norm = self._vector(node)
+            if all(scorer(vector, vector_norm, o, o_norm) < value for o, o_norm in kept_vectors):
+                kept.append((value, node))
+                kept_vectors.append((vector, vector_norm))
+            else:
+                pruned.append((value, node))
+        for pair in pruned:
+            if len(kept) >= capacity:
+                break
+            kept.append(pair)
+        return kept
+
+    def _link(
+        self, scorer: ScoreFn, node: int, new_neighbor: int, link_score: float, layer: int
+    ) -> None:
+        """Add a back-link node -> new_neighbor, shrinking the list if it overflows."""
+        current = self.neighbors(node, layer)
+        capacity = self._capacity(layer)
+        if len(current) < capacity:
+            self._write_neighbors(node, layer, current + [new_neighbor])
+            return
+        base, base_norm = self._vector(node)
+        scored = [(self._score(scorer, base, base_norm, other), other) for other in current]
+        scored.append((link_score, new_neighbor))
+        scored.sort(reverse=True)
+        chosen = self._select_neighbors(scorer, scored, capacity)
+        self._write_neighbors(node, layer, [other for _s, other in chosen])