Skip to content

Vector Search Study

vector_search_study

Public package interface for Vector Search Study.

ExactSearcher

Bases: Protocol

Common synchronous interface implemented by every exact searcher.

Source code in src/vector_search_study/api.py
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class ExactSearcher(Protocol):
    """Common synchronous interface implemented by every exact searcher."""

    @property
    def size(self) -> int:
        """Return the number of indexed corpus vectors."""
        ...

    @property
    def dimension(self) -> int:
        """Return the embedding dimension."""
        ...

    @property
    def dtype(self) -> FloatDType:
        """Return the indexed scalar dtype."""
        ...

    @property
    def objective(self) -> SearchObjective:
        """Return the search objective."""
        ...

    def prepare_queries(self, queries: FloatMatrix) -> PreparedQueries:
        """Prepare queries outside the timed search operation."""
        ...

    def search(self, queries: FloatMatrix, k: int) -> SearchResult:
        """Validate and search a raw query matrix."""
        ...

    def search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Search queries prepared outside the timed operation."""
        ...

size property

size: int

Return the number of indexed corpus vectors.

dimension property

dimension: int

Return the embedding dimension.

dtype property

dtype: FloatDType

Return the indexed scalar dtype.

objective property

objective: SearchObjective

Return the search objective.

prepare_queries

prepare_queries(queries: FloatMatrix) -> PreparedQueries

Prepare queries outside the timed search operation.

Source code in src/vector_search_study/api.py
167
168
169
def prepare_queries(self, queries: FloatMatrix) -> PreparedQueries:
    """Prepare queries outside the timed search operation."""
    ...

search

search(queries: FloatMatrix, k: int) -> SearchResult

Validate and search a raw query matrix.

Source code in src/vector_search_study/api.py
171
172
173
def search(self, queries: FloatMatrix, k: int) -> SearchResult:
    """Validate and search a raw query matrix."""
    ...

search_prepared

search_prepared(
    queries: PreparedQueries, k: int
) -> SearchResult

Search queries prepared outside the timed operation.

Source code in src/vector_search_study/api.py
175
176
177
def search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
    """Search queries prepared outside the timed operation."""
    ...

PreparedQueries dataclass

Validated, immutable queries ready for repeated timed searches.

Constructing this object validates normalization and copies the query matrix. Benchmarks can therefore prepare it outside the timed operation.

Parameters:

Name Type Description Default
values FloatMatrix

C-contiguous float32 or float64 query matrix.

required
objective SearchObjective

Score convention for which the queries were prepared.

NORMALIZED_COSINE
Source code in src/vector_search_study/api.py
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
@dataclass(frozen=True, slots=True)
class PreparedQueries:
    """Validated, immutable queries ready for repeated timed searches.

    Constructing this object validates normalization and copies the query
    matrix. Benchmarks can therefore prepare it outside the timed operation.

    Args:
        values: C-contiguous float32 or float64 query matrix.
        objective: Score convention for which the queries were prepared.
    """

    values: FloatMatrix
    objective: SearchObjective = SearchObjective.NORMALIZED_COSINE
    backend_name: str | None = None
    backend_payload: object | None = field(default=None, repr=False, compare=False)

    def __post_init__(self) -> None:
        """Validate, own, and freeze the query matrix."""
        objective = resolve_search_objective(self.objective)
        validated = validate_vector_matrix(
            self.values,
            name="queries",
            require_normalized=objective.requires_normalization,
        )
        owned = np.array(validated, dtype=validated.dtype, order="C", copy=True)
        owned.flags.writeable = False
        object.__setattr__(self, "values", owned)
        object.__setattr__(self, "objective", objective)
        if self.backend_name is not None and not self.backend_name:
            raise InvalidVectorDataError("backend_name must be non-empty when provided")

    @property
    def query_count(self) -> int:
        """Return the number of queries in the batch."""
        return self.values.shape[0]

    @property
    def dimension(self) -> int:
        """Return the embedding dimension."""
        return self.values.shape[1]

    @property
    def dtype(self) -> FloatDType:
        """Return the query scalar dtype."""
        return self.values.dtype

query_count property

query_count: int

Return the number of queries in the batch.

dimension property

dimension: int

Return the embedding dimension.

dtype property

dtype: FloatDType

Return the query scalar dtype.

__post_init__

__post_init__() -> None

Validate, own, and freeze the query matrix.

Source code in src/vector_search_study/api.py
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def __post_init__(self) -> None:
    """Validate, own, and freeze the query matrix."""
    objective = resolve_search_objective(self.objective)
    validated = validate_vector_matrix(
        self.values,
        name="queries",
        require_normalized=objective.requires_normalization,
    )
    owned = np.array(validated, dtype=validated.dtype, order="C", copy=True)
    owned.flags.writeable = False
    object.__setattr__(self, "values", owned)
    object.__setattr__(self, "objective", objective)
    if self.backend_name is not None and not self.backend_name:
        raise InvalidVectorDataError("backend_name must be non-empty when provided")

SearchObjective

Bases: StrEnum

Exact-search score convention used throughout the study.

Source code in src/vector_search_study/api.py
18
19
20
21
22
23
24
25
26
27
28
class SearchObjective(StrEnum):
    """Exact-search score convention used throughout the study."""

    SQUARED_L2 = "squared_l2"
    INNER_PRODUCT = "inner_product"
    NORMALIZED_COSINE = "normalized_cosine"

    @property
    def requires_normalization(self) -> bool:
        """Return whether corpus and query rows must have unit L2 norm."""
        return self is SearchObjective.NORMALIZED_COSINE

requires_normalization property

requires_normalization: bool

Return whether corpus and query rows must have unit L2 norm.

SearchResult dataclass

Ordered exact top-k indices and scores for a query batch.

Rows are ordered by decreasing score, with smaller corpus indices winning exact score ties. Scores are float64 and always use a higher-is-better convention: negative squared distance, inner product, or normalized cosine.

Parameters:

Name Type Description Default
indices NDArray[int64]

Corpus indices with shape (query_count, k).

required
scores NDArray[float64]

Objective scores with the same shape.

required
Source code in src/vector_search_study/api.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
@dataclass(frozen=True, slots=True)
class SearchResult:
    """Ordered exact top-k indices and scores for a query batch.

    Rows are ordered by decreasing score, with smaller corpus indices winning
    exact score ties. Scores are float64 and always use a higher-is-better
    convention: negative squared distance, inner product, or normalized cosine.

    Args:
        indices: Corpus indices with shape ``(query_count, k)``.
        scores: Objective scores with the same shape.
    """

    indices: NDArray[np.int64]
    scores: NDArray[np.float64]

    def __post_init__(self) -> None:
        """Enforce the public result representation and ordering contract."""
        if not isinstance(self.indices, np.ndarray) or self.indices.dtype != np.dtype(np.int64):
            raise InvalidVectorDataError("result indices must be an int64 NumPy array")
        if not isinstance(self.scores, np.ndarray) or self.scores.dtype != np.dtype(np.float64):
            raise InvalidVectorDataError("result scores must be a float64 NumPy array")
        if self.indices.ndim != 2 or self.scores.ndim != 2 or self.indices.shape != self.scores.shape:
            raise InvalidVectorDataError("result indices and scores must have the same two-dimensional shape")
        if self.indices.shape[0] == 0 or self.indices.shape[1] == 0:
            raise InvalidVectorDataError("result arrays must not have an empty axis")
        if not self.indices.flags.c_contiguous or not self.scores.flags.c_contiguous:
            raise InvalidVectorDataError("result arrays must be C-contiguous")
        if bool(np.any(self.indices < 0)):
            raise InvalidVectorDataError("result indices must be non-negative")
        if not bool(np.isfinite(self.scores).all()):
            raise InvalidVectorDataError("result scores must be finite")

        if self.scores.shape[1] > 1:
            left_scores = self.scores[:, :-1]
            right_scores = self.scores[:, 1:]
            if bool(np.any(left_scores < right_scores)):
                raise InvalidVectorDataError("result scores must be ordered from greatest to least")
            tied = left_scores == right_scores
            if bool(np.any(tied & (self.indices[:, :-1] > self.indices[:, 1:]))):
                raise InvalidVectorDataError("result index ties must be ordered from least to greatest")

        self.indices.flags.writeable = False
        self.scores.flags.writeable = False

__post_init__

__post_init__() -> None

Enforce the public result representation and ordering contract.

Source code in src/vector_search_study/api.py
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def __post_init__(self) -> None:
    """Enforce the public result representation and ordering contract."""
    if not isinstance(self.indices, np.ndarray) or self.indices.dtype != np.dtype(np.int64):
        raise InvalidVectorDataError("result indices must be an int64 NumPy array")
    if not isinstance(self.scores, np.ndarray) or self.scores.dtype != np.dtype(np.float64):
        raise InvalidVectorDataError("result scores must be a float64 NumPy array")
    if self.indices.ndim != 2 or self.scores.ndim != 2 or self.indices.shape != self.scores.shape:
        raise InvalidVectorDataError("result indices and scores must have the same two-dimensional shape")
    if self.indices.shape[0] == 0 or self.indices.shape[1] == 0:
        raise InvalidVectorDataError("result arrays must not have an empty axis")
    if not self.indices.flags.c_contiguous or not self.scores.flags.c_contiguous:
        raise InvalidVectorDataError("result arrays must be C-contiguous")
    if bool(np.any(self.indices < 0)):
        raise InvalidVectorDataError("result indices must be non-negative")
    if not bool(np.isfinite(self.scores).all()):
        raise InvalidVectorDataError("result scores must be finite")

    if self.scores.shape[1] > 1:
        left_scores = self.scores[:, :-1]
        right_scores = self.scores[:, 1:]
        if bool(np.any(left_scores < right_scores)):
            raise InvalidVectorDataError("result scores must be ordered from greatest to least")
        tied = left_scores == right_scores
        if bool(np.any(tied & (self.indices[:, :-1] > self.indices[:, 1:]))):
            raise InvalidVectorDataError("result index ties must be ordered from least to greatest")

    self.indices.flags.writeable = False
    self.scores.flags.writeable = False

BackendUnavailableError

Bases: VectorSearchStudyError

Raised when an optional search backend is not installed.

Source code in src/vector_search_study/exceptions.py
16
17
class BackendUnavailableError(VectorSearchStudyError):
    """Raised when an optional search backend is not installed."""

InvalidSearchParameterError

Bases: VectorSearchStudyError

Raised when a search parameter is outside its supported range.

Source code in src/vector_search_study/exceptions.py
12
13
class InvalidSearchParameterError(VectorSearchStudyError):
    """Raised when a search parameter is outside its supported range."""

InvalidVectorDataError

Bases: VectorSearchStudyError

Raised when a corpus or query matrix violates the vector contract.

Source code in src/vector_search_study/exceptions.py
8
9
class InvalidVectorDataError(VectorSearchStudyError):
    """Raised when a corpus or query matrix violates the vector contract."""

UnsupportedObjectiveError

Bases: VectorSearchStudyError

Raised when a backend cannot exactly implement a search objective.

Source code in src/vector_search_study/exceptions.py
20
21
class UnsupportedObjectiveError(VectorSearchStudyError):
    """Raised when a backend cannot exactly implement a search objective."""

VectorSearchStudyError

Bases: ValueError

Raised when Vector Search Study cannot complete an operation.

Source code in src/vector_search_study/exceptions.py
4
5
class VectorSearchStudyError(ValueError):
    """Raised when Vector Search Study cannot complete an operation."""

FaissFlatIPSearcher

Bases: BaseExactSearcher

Exact Faiss IndexFlatIP search for inner product or cosine.

Source code in src/vector_search_study/faiss_search.py
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class FaissFlatIPSearcher(BaseExactSearcher):
    """Exact Faiss IndexFlatIP search for inner product or cosine."""

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str = SearchObjective.INNER_PRODUCT,
    ) -> None:
        """Build a float32 flat inner-product index outside search timing."""
        resolved = resolve_search_objective(objective)
        if resolved not in {SearchObjective.INNER_PRODUCT, SearchObjective.NORMALIZED_COSINE}:
            raise UnsupportedObjectiveError("Faiss IndexFlatIP supports only inner product and normalized cosine")
        _require_float32(corpus)
        super().__init__(corpus, objective=resolved)
        faiss = import_optional("faiss", extra="benchmark-backends")
        faiss.omp_set_num_threads(1)
        self._index: Any = faiss.IndexFlatIP(self.dimension)
        self._index.add(self._corpus)

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Search the flat inner-product index."""
        scores, indices = self._index.search(queries.values, k)
        return canonical_result(indices, scores)

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective
    | str = SearchObjective.INNER_PRODUCT,
) -> None

Build a float32 flat inner-product index outside search timing.

Source code in src/vector_search_study/faiss_search.py
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str = SearchObjective.INNER_PRODUCT,
) -> None:
    """Build a float32 flat inner-product index outside search timing."""
    resolved = resolve_search_objective(objective)
    if resolved not in {SearchObjective.INNER_PRODUCT, SearchObjective.NORMALIZED_COSINE}:
        raise UnsupportedObjectiveError("Faiss IndexFlatIP supports only inner product and normalized cosine")
    _require_float32(corpus)
    super().__init__(corpus, objective=resolved)
    faiss = import_optional("faiss", extra="benchmark-backends")
    faiss.omp_set_num_threads(1)
    self._index: Any = faiss.IndexFlatIP(self.dimension)
    self._index.add(self._corpus)

FaissFlatL2Searcher

Bases: BaseExactSearcher

Exact Faiss IndexFlatL2 search with negative squared-distance scores.

Source code in src/vector_search_study/faiss_search.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class FaissFlatL2Searcher(BaseExactSearcher):
    """Exact Faiss IndexFlatL2 search with negative squared-distance scores."""

    def __init__(self, corpus: FloatMatrix) -> None:
        """Build a float32 flat L2 index outside search timing."""
        _require_float32(corpus)
        super().__init__(corpus, objective=SearchObjective.SQUARED_L2)
        faiss = import_optional("faiss", extra="benchmark-backends")
        faiss.omp_set_num_threads(1)
        self._index: Any = faiss.IndexFlatL2(self.dimension)
        self._index.add(self._corpus)

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Search the flat L2 index and invert its distances."""
        distances, indices = self._index.search(queries.values, k)
        return canonical_result(indices, -distances)

__init__

__init__(corpus: FloatMatrix) -> None

Build a float32 flat L2 index outside search timing.

Source code in src/vector_search_study/faiss_search.py
19
20
21
22
23
24
25
26
def __init__(self, corpus: FloatMatrix) -> None:
    """Build a float32 flat L2 index outside search timing."""
    _require_float32(corpus)
    super().__init__(corpus, objective=SearchObjective.SQUARED_L2)
    faiss = import_optional("faiss", extra="benchmark-backends")
    faiss.omp_set_num_threads(1)
    self._index: Any = faiss.IndexFlatL2(self.dimension)
    self._index.add(self._corpus)

NumpyArgpartitionSearcher

Bases: BaseExactSearcher

Score by matrix multiplication and partially select exact top-k rows.

Source code in src/vector_search_study/numpy_search.py
38
39
40
41
42
43
44
45
class NumpyArgpartitionSearcher(BaseExactSearcher):
    """Score by matrix multiplication and partially select exact top-k rows."""

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Materialize scores, partition, and repair deterministic ties."""
        scores_matrix = score_matrix(queries.values, self._corpus, self.objective)
        indices, scores = partition_top_k(scores_matrix, k)
        return SearchResult(indices=indices, scores=scores)

NumpyBlockedSearcher

Bases: BaseExactSearcher

Limit score-matrix memory by searching fixed-size corpus blocks.

Source code in src/vector_search_study/numpy_search.py
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class NumpyBlockedSearcher(BaseExactSearcher):
    """Limit score-matrix memory by searching fixed-size corpus blocks."""

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        block_size: int = 16_384,
        objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
    ) -> None:
        """Build an index with a fixed number of corpus rows per block.

        Args:
            corpus: Pre-normalized corpus vectors.
            block_size: Maximum number of corpus rows scored per matrix
                multiplication.
            objective: Exact-search score convention.
        """
        super().__init__(corpus, objective=objective)
        self._block_size = validate_positive_int(block_size, name="block_size")

    @property
    def block_size(self) -> int:
        """Return the configured corpus rows per score block."""
        return self._block_size

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Select block-local candidates and merge them into global top-k."""
        best_indices = np.empty((queries.query_count, 0), dtype=np.int64)
        best_scores = np.empty((queries.query_count, 0), dtype=np.float64)
        for start in range(0, self.size, self._block_size):
            stop = min(start + self._block_size, self.size)
            block_scores = score_matrix(queries.values, self._corpus[start:stop], self.objective)
            block_k = min(k, stop - start)
            block_indices, selected_scores = partition_top_k(block_scores, block_k, index_offset=start)
            best_indices, best_scores = merge_top_k(
                best_indices,
                best_scores,
                block_indices,
                selected_scores,
                k,
            )
        return SearchResult(indices=best_indices, scores=best_scores)

block_size property

block_size: int

Return the configured corpus rows per score block.

__init__

__init__(
    corpus: FloatMatrix,
    *,
    block_size: int = 16384,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> None

Build an index with a fixed number of corpus rows per block.

Parameters:

Name Type Description Default
corpus FloatMatrix

Pre-normalized corpus vectors.

required
block_size int

Maximum number of corpus rows scored per matrix multiplication.

16384
objective SearchObjective | str

Exact-search score convention.

NORMALIZED_COSINE
Source code in src/vector_search_study/numpy_search.py
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    block_size: int = 16_384,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> None:
    """Build an index with a fixed number of corpus rows per block.

    Args:
        corpus: Pre-normalized corpus vectors.
        block_size: Maximum number of corpus rows scored per matrix
            multiplication.
        objective: Exact-search score convention.
    """
    super().__init__(corpus, objective=objective)
    self._block_size = validate_positive_int(block_size, name="block_size")

NumpySortSearcher

Bases: BaseExactSearcher

Score by matrix multiplication and fully sort every score row.

Source code in src/vector_search_study/numpy_search.py
28
29
30
31
32
33
34
35
class NumpySortSearcher(BaseExactSearcher):
    """Score by matrix multiplication and fully sort every score row."""

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Materialize the full score matrix and fully sort it."""
        scores_matrix = score_matrix(queries.values, self._corpus, self.objective)
        indices, scores = full_sort_top_k(scores_matrix, k)
        return SearchResult(indices=indices, scores=scores)

PythonHeapSearcher

Bases: BaseExactSearcher

Stream exhaustive scalar scores through a bounded size-k heap.

Source code in src/vector_search_study/python_search.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class PythonHeapSearcher(BaseExactSearcher):
    """Stream exhaustive scalar scores through a bounded size-k heap."""

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
    ) -> None:
        """Build the scalar corpus representation outside search timing."""
        super().__init__(corpus, objective=objective)
        self._rows = tuple(tuple(float(value) for value in row) for row in self._corpus)

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Retain only the best k candidates while streaming scores."""
        all_indices: list[list[int]] = []
        all_scores: list[list[float]] = []
        for query_array in queries.values:
            query = tuple(float(value) for value in query_array)
            heap: list[tuple[float, int, int]] = []
            for index, vector in enumerate(self._rows):
                score = _scalar_score(vector, query, self.objective)
                item = (score, -index, index)
                if len(heap) < k:
                    heapq.heappush(heap, item)
                elif item[:2] > heap[0][:2]:
                    _ = heapq.heapreplace(heap, item)
            selected = sorted(heap, key=lambda item: (-item[0], item[2]))
            all_indices.append([index for _, _, index in selected])
            all_scores.append([score for score, _, _ in selected])
        return SearchResult(
            indices=np.asarray(all_indices, dtype=np.int64, order="C"),
            scores=np.asarray(all_scores, dtype=np.float64, order="C"),
        )

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> None

Build the scalar corpus representation outside search timing.

Source code in src/vector_search_study/python_search.py
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> None:
    """Build the scalar corpus representation outside search timing."""
    super().__init__(corpus, objective=objective)
    self._rows = tuple(tuple(float(value) for value in row) for row in self._corpus)

PythonSortSearcher

Bases: BaseExactSearcher

Exhaustively score into a Python list and fully sort it.

Source code in src/vector_search_study/python_search.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class PythonSortSearcher(BaseExactSearcher):
    """Exhaustively score into a Python list and fully sort it."""

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
    ) -> None:
        """Build the scalar corpus representation outside search timing."""
        super().__init__(corpus, objective=objective)
        self._rows = tuple(tuple(float(value) for value in row) for row in self._corpus)

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Score every vector, then sort all candidates."""
        all_indices: list[list[int]] = []
        all_scores: list[list[float]] = []
        for query_array in queries.values:
            query = tuple(float(value) for value in query_array)
            ranked: list[tuple[float, int]] = []
            for index, vector in enumerate(self._rows):
                score = _scalar_score(vector, query, self.objective)
                ranked.append((score, index))
            ranked.sort(key=lambda item: (-item[0], item[1]))
            selected = ranked[:k]
            all_indices.append([index for _, index in selected])
            all_scores.append([score for score, _ in selected])
        return SearchResult(
            indices=np.asarray(all_indices, dtype=np.int64, order="C"),
            scores=np.asarray(all_scores, dtype=np.float64, order="C"),
        )

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> None

Build the scalar corpus representation outside search timing.

Source code in src/vector_search_study/python_search.py
17
18
19
20
21
22
23
24
25
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> None:
    """Build the scalar corpus representation outside search timing."""
    super().__init__(corpus, objective=objective)
    self._rows = tuple(tuple(float(value) for value in row) for row in self._corpus)

ScipyCKDTreeSearcher

Bases: BaseExactSearcher

Exact one-worker SciPy cKDTree search for L2-derived objectives.

Source code in src/vector_search_study/scipy_search.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class ScipyCKDTreeSearcher(BaseExactSearcher):
    """Exact one-worker SciPy cKDTree search for L2-derived objectives."""

    _supported: ClassVar[frozenset[SearchObjective]] = frozenset(
        {SearchObjective.SQUARED_L2, SearchObjective.NORMALIZED_COSINE}
    )

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str,
        leaf_size: int = 16,
    ) -> None:
        """Build a balanced compact cKDTree outside search timing."""
        resolved = _require_supported(objective, self._supported, backend="SciPy cKDTree")
        super().__init__(corpus, objective=resolved)
        spatial = import_optional("scipy.spatial", extra="benchmark-backends")
        self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
        self._index: Any = spatial.cKDTree(
            self._corpus,
            leafsize=self._leaf_size,
            compact_nodes=True,
            copy_data=True,
            balanced_tree=True,
        )

    @property
    def leaf_size(self) -> int:
        """Return the configured tree leaf size."""
        return self._leaf_size

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Query the exact Euclidean tree with one worker."""
        distances, indices = self._index.query(queries.values, k=k, eps=0.0, p=2.0, workers=1)
        return canonical_result(indices, _euclidean_scores(distances, self.objective))

leaf_size property

leaf_size: int

Return the configured tree leaf size.

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 16,
) -> None

Build a balanced compact cKDTree outside search timing.

Source code in src/vector_search_study/scipy_search.py
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 16,
) -> None:
    """Build a balanced compact cKDTree outside search timing."""
    resolved = _require_supported(objective, self._supported, backend="SciPy cKDTree")
    super().__init__(corpus, objective=resolved)
    spatial = import_optional("scipy.spatial", extra="benchmark-backends")
    self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
    self._index: Any = spatial.cKDTree(
        self._corpus,
        leafsize=self._leaf_size,
        compact_nodes=True,
        copy_data=True,
        balanced_tree=True,
    )

SklearnBallTreeSearcher

Bases: BaseExactSearcher

Exact scikit-learn BallTree search for L2-derived objectives.

Source code in src/vector_search_study/sklearn_search.py
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
class SklearnBallTreeSearcher(BaseExactSearcher):
    """Exact scikit-learn BallTree search for L2-derived objectives."""

    _supported: ClassVar[frozenset[SearchObjective]] = SklearnKDTreeSearcher._supported

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str,
        leaf_size: int = 40,
    ) -> None:
        """Build a Euclidean BallTree outside search timing."""
        resolved = _require_supported(objective, self._supported, backend="scikit-learn BallTree")
        super().__init__(corpus, objective=resolved)
        neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
        self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
        self._index: Any = neighbors.BallTree(self._corpus, leaf_size=self._leaf_size, metric="euclidean")

    @property
    def leaf_size(self) -> int:
        """Return the configured tree leaf size."""
        return self._leaf_size

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Query the Euclidean tree and convert distances to objective scores."""
        distances, indices = self._index.query(queries.values, k=k, return_distance=True, dualtree=False)
        return canonical_result(indices, _euclidean_scores(distances, self.objective))

leaf_size property

leaf_size: int

Return the configured tree leaf size.

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 40,
) -> None

Build a Euclidean BallTree outside search timing.

Source code in src/vector_search_study/sklearn_search.py
76
77
78
79
80
81
82
83
84
85
86
87
88
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 40,
) -> None:
    """Build a Euclidean BallTree outside search timing."""
    resolved = _require_supported(objective, self._supported, backend="scikit-learn BallTree")
    super().__init__(corpus, objective=resolved)
    neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
    self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
    self._index: Any = neighbors.BallTree(self._corpus, leaf_size=self._leaf_size, metric="euclidean")

SklearnBruteSearcher

Bases: BaseExactSearcher

Exact scikit-learn brute-force L2 or cosine search.

Source code in src/vector_search_study/sklearn_search.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class SklearnBruteSearcher(BaseExactSearcher):
    """Exact scikit-learn brute-force L2 or cosine search."""

    _supported: ClassVar[frozenset[SearchObjective]] = frozenset(
        {SearchObjective.SQUARED_L2, SearchObjective.NORMALIZED_COSINE}
    )

    def __init__(self, corpus: FloatMatrix, *, objective: SearchObjective | str) -> None:
        """Build a one-thread brute-force nearest-neighbor index."""
        resolved = _require_supported(objective, self._supported, backend="scikit-learn brute")
        super().__init__(corpus, objective=resolved)
        neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
        metric = "cosine" if resolved is SearchObjective.NORMALIZED_COSINE else "euclidean"
        self._index: Any = neighbors.NearestNeighbors(algorithm="brute", metric=metric, n_jobs=1)
        self._index.fit(self._corpus)

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Query the fitted brute-force index."""
        distances, indices = self._index.kneighbors(queries.values, n_neighbors=k, return_distance=True)
        scores = -(distances * distances) if self.objective is SearchObjective.SQUARED_L2 else 1.0 - distances
        return canonical_result(indices, scores)

__init__

__init__(
    corpus: FloatMatrix, *, objective: SearchObjective | str
) -> None

Build a one-thread brute-force nearest-neighbor index.

Source code in src/vector_search_study/sklearn_search.py
23
24
25
26
27
28
29
30
def __init__(self, corpus: FloatMatrix, *, objective: SearchObjective | str) -> None:
    """Build a one-thread brute-force nearest-neighbor index."""
    resolved = _require_supported(objective, self._supported, backend="scikit-learn brute")
    super().__init__(corpus, objective=resolved)
    neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
    metric = "cosine" if resolved is SearchObjective.NORMALIZED_COSINE else "euclidean"
    self._index: Any = neighbors.NearestNeighbors(algorithm="brute", metric=metric, n_jobs=1)
    self._index.fit(self._corpus)

SklearnKDTreeSearcher

Bases: BaseExactSearcher

Exact scikit-learn KDTree search for L2-derived objectives.

Source code in src/vector_search_study/sklearn_search.py
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class SklearnKDTreeSearcher(BaseExactSearcher):
    """Exact scikit-learn KDTree search for L2-derived objectives."""

    _supported: ClassVar[frozenset[SearchObjective]] = frozenset(
        {SearchObjective.SQUARED_L2, SearchObjective.NORMALIZED_COSINE}
    )

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str,
        leaf_size: int = 40,
    ) -> None:
        """Build a Euclidean KDTree outside search timing."""
        resolved = _require_supported(objective, self._supported, backend="scikit-learn KDTree")
        super().__init__(corpus, objective=resolved)
        neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
        self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
        self._index: Any = neighbors.KDTree(self._corpus, leaf_size=self._leaf_size, metric="euclidean")

    @property
    def leaf_size(self) -> int:
        """Return the configured tree leaf size."""
        return self._leaf_size

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Query the Euclidean tree and convert distances to objective scores."""
        distances, indices = self._index.query(queries.values, k=k, return_distance=True, dualtree=False)
        return canonical_result(indices, _euclidean_scores(distances, self.objective))

leaf_size property

leaf_size: int

Return the configured tree leaf size.

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 40,
) -> None

Build a Euclidean KDTree outside search timing.

Source code in src/vector_search_study/sklearn_search.py
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str,
    leaf_size: int = 40,
) -> None:
    """Build a Euclidean KDTree outside search timing."""
    resolved = _require_supported(objective, self._supported, backend="scikit-learn KDTree")
    super().__init__(corpus, objective=resolved)
    neighbors = import_optional("sklearn.neighbors", extra="benchmark-backends")
    self._leaf_size = validate_positive_int(leaf_size, name="leaf_size")
    self._index: Any = neighbors.KDTree(self._corpus, leaf_size=self._leaf_size, metric="euclidean")

SyntheticDataset dataclass

An immutable generated corpus/query pair and its provenance.

Attributes:

Name Type Description
corpus FloatMatrix

Pre-normalized corpus vectors.

queries FloatMatrix

Pre-normalized query vectors.

distribution str

Generator family name.

seed int

PCG64 seed.

objective SearchObjective

Search objective for which vectors were generated.

Source code in src/vector_search_study/synthetic.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
@dataclass(frozen=True, slots=True)
class SyntheticDataset:
    """An immutable generated corpus/query pair and its provenance.

    Attributes:
        corpus: Pre-normalized corpus vectors.
        queries: Pre-normalized query vectors.
        distribution: Generator family name.
        seed: PCG64 seed.
        objective: Search objective for which vectors were generated.
    """

    corpus: FloatMatrix
    queries: FloatMatrix
    distribution: str
    seed: int
    objective: SearchObjective = SearchObjective.NORMALIZED_COSINE

TorchTopKSearcher

Bases: BaseExactSearcher

Exact CPU PyTorch matmul/topk search for every study objective.

Source code in src/vector_search_study/torch_search.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class TorchTopKSearcher(BaseExactSearcher):
    """Exact CPU PyTorch matmul/topk search for every study objective."""

    _backend_name = "torch_cpu"

    def __init__(
        self,
        corpus: FloatMatrix,
        *,
        objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
    ) -> None:
        """Materialize the corpus tensor outside search timing."""
        super().__init__(corpus, objective=objective)
        self._torch: Any = import_optional("torch", extra="benchmark-backends")
        self._torch.set_num_threads(1)
        self._corpus_tensor: Any = self._torch.from_numpy(np.array(self._corpus, copy=True, order="C"))
        self._corpus_norms: Any | None = None
        if self.objective is SearchObjective.SQUARED_L2:
            self._corpus_norms = (self._corpus_tensor * self._corpus_tensor).sum(dim=1).unsqueeze(0)

    def prepare_queries(self, queries: FloatMatrix) -> PreparedQueries:
        """Validate queries and materialize their CPU tensor outside timing."""
        prepared = super().prepare_queries(queries)
        tensor = self._torch.from_numpy(np.array(prepared.values, copy=True, order="C"))
        return PreparedQueries(
            prepared.values,
            objective=self.objective,
            backend_name=self._backend_name,
            backend_payload=tensor,
        )

    def _search_prepared(self, queries: PreparedQueries, k: int) -> SearchResult:
        """Run matmul and topk on CPU, then canonicalize tied candidates."""
        if queries.backend_name != self._backend_name or queries.backend_payload is None:
            raise InvalidVectorDataError("queries must be prepared by this PyTorch searcher")
        query_tensor = cast(Any, queries.backend_payload)
        scores = self._torch.matmul(query_tensor, self._corpus_tensor.transpose(0, 1))
        if self.objective is SearchObjective.SQUARED_L2:
            query_norms = (query_tensor * query_tensor).sum(dim=1).unsqueeze(1)
            scores = -(query_norms + self._corpus_norms - 2.0 * scores).clamp_min(0.0)
        selected_scores, selected_indices = self._torch.topk(scores, k, dim=1, largest=True, sorted=False)
        return canonical_result(
            selected_indices.detach().cpu().numpy(),
            selected_scores.detach().cpu().numpy(),
        )

__init__

__init__(
    corpus: FloatMatrix,
    *,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> None

Materialize the corpus tensor outside search timing.

Source code in src/vector_search_study/torch_search.py
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def __init__(
    self,
    corpus: FloatMatrix,
    *,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> None:
    """Materialize the corpus tensor outside search timing."""
    super().__init__(corpus, objective=objective)
    self._torch: Any = import_optional("torch", extra="benchmark-backends")
    self._torch.set_num_threads(1)
    self._corpus_tensor: Any = self._torch.from_numpy(np.array(self._corpus, copy=True, order="C"))
    self._corpus_norms: Any | None = None
    if self.objective is SearchObjective.SQUARED_L2:
        self._corpus_norms = (self._corpus_tensor * self._corpus_tensor).sum(dim=1).unsqueeze(0)

prepare_queries

prepare_queries(queries: FloatMatrix) -> PreparedQueries

Validate queries and materialize their CPU tensor outside timing.

Source code in src/vector_search_study/torch_search.py
36
37
38
39
40
41
42
43
44
45
def prepare_queries(self, queries: FloatMatrix) -> PreparedQueries:
    """Validate queries and materialize their CPU tensor outside timing."""
    prepared = super().prepare_queries(queries)
    tensor = self._torch.from_numpy(np.array(prepared.values, copy=True, order="C"))
    return PreparedQueries(
        prepared.values,
        objective=self.objective,
        backend_name=self._backend_name,
        backend_payload=tensor,
    )

normalize_rows

normalize_rows(values: FloatMatrix) -> FloatMatrix

Return a C-contiguous copy with every row L2-normalized.

Parameters:

Name Type Description Default
values FloatMatrix

A finite, non-empty float32 or float64 matrix.

required

Returns:

Type Description
FloatMatrix

A normalized matrix with the input dtype.

Raises:

Type Description
InvalidVectorDataError

If the input violates the matrix contract or contains a zero row.

Source code in src/vector_search_study/api.py
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def normalize_rows(values: FloatMatrix) -> FloatMatrix:
    """Return a C-contiguous copy with every row L2-normalized.

    Args:
        values: A finite, non-empty float32 or float64 matrix.

    Returns:
        A normalized matrix with the input dtype.

    Raises:
        InvalidVectorDataError: If the input violates the matrix contract or
            contains a zero row.
    """
    matrix = validate_vector_matrix(values, name="values", require_normalized=False)
    scales = np.max(np.abs(matrix), axis=1, keepdims=True)
    if bool(np.any(scales == 0.0)):
        raise InvalidVectorDataError("values must not contain a zero row")
    scaled = matrix / scales
    norms = np.sqrt(np.sum(scaled * scaled, axis=1, keepdims=True))
    return np.asarray(scaled / norms, dtype=matrix.dtype, order="C")

prepare_queries

prepare_queries(
    queries: FloatMatrix,
    *,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> PreparedQueries

Validate and copy queries for repeated search.

Parameters:

Name Type Description Default
queries FloatMatrix

C-contiguous float32 or float64 matrix. Rows must be normalized for normalized cosine.

required
objective SearchObjective | str

Score convention for the prepared queries.

NORMALIZED_COSINE

Returns:

Type Description
PreparedQueries

An immutable prepared query batch.

Source code in src/vector_search_study/api.py
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def prepare_queries(
    queries: FloatMatrix,
    *,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> PreparedQueries:
    """Validate and copy queries for repeated search.

    Args:
        queries: C-contiguous float32 or float64 matrix. Rows must be
            normalized for normalized cosine.
        objective: Score convention for the prepared queries.

    Returns:
        An immutable prepared query batch.
    """
    return PreparedQueries(queries, objective=resolve_search_objective(objective))
reference_search(
    corpus: FloatMatrix,
    queries: FloatMatrix,
    k: int,
    *,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> SearchResult

Compute canonical exact top-k results with accurate scalar summation.

This intentionally slow implementation is designed for correctness tests and untimed benchmark validation, not performance measurement.

Parameters:

Name Type Description Default
corpus FloatMatrix

Corpus matrix with shape (N, D).

required
queries FloatMatrix

Query matrix with shape (Q, D).

required
k int

Number of ordered neighbors to return.

required
objective SearchObjective | str

Exact-search score convention.

NORMALIZED_COSINE

Returns:

Type Description
SearchResult

Canonically ordered exact results.

Raises:

Type Description
InvalidVectorDataError

If corpus and query contracts do not match.

Source code in src/vector_search_study/reference.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def reference_search(
    corpus: FloatMatrix,
    queries: FloatMatrix,
    k: int,
    *,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> SearchResult:
    """Compute canonical exact top-k results with accurate scalar summation.

    This intentionally slow implementation is designed for correctness tests
    and untimed benchmark validation, not performance measurement.

    Args:
        corpus: Corpus matrix with shape ``(N, D)``.
        queries: Query matrix with shape ``(Q, D)``.
        k: Number of ordered neighbors to return.
        objective: Exact-search score convention.

    Returns:
        Canonically ordered exact results.

    Raises:
        InvalidVectorDataError: If corpus and query contracts do not match.
    """
    resolved_objective = resolve_search_objective(objective)
    validated_corpus = validate_vector_matrix(
        corpus,
        name="corpus",
        require_normalized=resolved_objective.requires_normalization,
    )
    prepared = PreparedQueries(queries, objective=resolved_objective)
    if prepared.dimension != validated_corpus.shape[1]:
        raise InvalidVectorDataError(
            f"query dimension {prepared.dimension} does not match corpus dimension {validated_corpus.shape[1]}"
        )
    if prepared.dtype != validated_corpus.dtype:
        raise InvalidVectorDataError(
            f"query dtype {prepared.dtype} does not match corpus dtype {validated_corpus.dtype}"
        )
    resolved_k = validate_search_k(k, corpus_size=validated_corpus.shape[0])

    indices = np.empty((prepared.query_count, resolved_k), dtype=np.int64)
    scores = np.empty((prepared.query_count, resolved_k), dtype=np.float64)
    for query_index, query in enumerate(prepared.values):
        ranked: list[tuple[float, int]] = []
        for corpus_index, vector in enumerate(validated_corpus):
            score = _reference_score(vector, query, resolved_objective)
            ranked.append((score, corpus_index))
        ranked.sort(key=lambda item: (-item[0], item[1]))
        for result_index, (score, corpus_index) in enumerate(ranked[:resolved_k]):
            indices[query_index, result_index] = corpus_index
            scores[query_index, result_index] = score
    return SearchResult(indices=indices, scores=scores)

make_clustered_dataset

make_clustered_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    cluster_count: int = 8,
    noise: float = 0.15,
    dtype: object = np.float32,
    seed: int = 20260801,
    objective: SearchObjective
    | str = SearchObjective.NORMALIZED_COSINE,
) -> SyntheticDataset

Generate normalized vectors around shared random cluster centroids.

Parameters:

Name Type Description Default
corpus_size int

Number of corpus vectors.

required
dimension int

Embedding dimension.

required
query_count int

Number of query vectors.

required
cluster_count int

Number of latent centroids.

8
noise float

Positive standard deviation around each centroid.

0.15
dtype object

Either float32 or float64.

float32
seed int

Non-negative PCG64 seed.

20260801
objective SearchObjective | str

Exact-search score convention. Cosine output is normalized; L2 and inner-product output is not.

NORMALIZED_COSINE

Returns:

Type Description
SyntheticDataset

A deterministic normalized clustered dataset.

Source code in src/vector_search_study/synthetic.py
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def make_clustered_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    cluster_count: int = 8,
    noise: float = 0.15,
    dtype: object = np.float32,
    seed: int = 20_260_801,
    objective: SearchObjective | str = SearchObjective.NORMALIZED_COSINE,
) -> SyntheticDataset:
    """Generate normalized vectors around shared random cluster centroids.

    Args:
        corpus_size: Number of corpus vectors.
        dimension: Embedding dimension.
        query_count: Number of query vectors.
        cluster_count: Number of latent centroids.
        noise: Positive standard deviation around each centroid.
        dtype: Either float32 or float64.
        seed: Non-negative PCG64 seed.
        objective: Exact-search score convention. Cosine output is normalized;
            L2 and inner-product output is not.

    Returns:
        A deterministic normalized clustered dataset.
    """
    size, dimensions, queries, resolved_dtype, resolved_seed = _validate_generator_inputs(
        corpus_size,
        dimension,
        query_count,
        dtype=dtype,
        seed=seed,
    )
    clusters = validate_positive_int(cluster_count, name="cluster_count")
    if isinstance(noise, bool) or not isinstance(noise, (int, float)) or not np.isfinite(noise) or noise <= 0:
        raise InvalidSearchParameterError("noise must be a positive finite number")

    resolved_objective = resolve_search_objective(objective)
    generator = np.random.Generator(np.random.PCG64(resolved_seed))
    centroids = generator.standard_normal((clusters, dimensions)).astype(resolved_dtype, copy=False)
    corpus_assignments = generator.integers(0, clusters, size=size)
    query_assignments = generator.integers(0, clusters, size=queries)
    corpus_noise = generator.standard_normal((size, dimensions)).astype(resolved_dtype, copy=False)
    query_noise = generator.standard_normal((queries, dimensions)).astype(resolved_dtype, copy=False)
    corpus = np.asarray(centroids[corpus_assignments] + noise * corpus_noise, dtype=resolved_dtype, order="C")
    query_matrix = np.asarray(
        centroids[query_assignments] + noise * query_noise,
        dtype=resolved_dtype,
        order="C",
    )
    if resolved_objective.requires_normalization:
        corpus = normalize_rows(corpus)
        query_matrix = normalize_rows(query_matrix)
    return _dataset(
        corpus,
        query_matrix,
        distribution="clustered",
        seed=resolved_seed,
        objective=resolved_objective,
    )

make_gaussian_dataset

make_gaussian_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    objective: SearchObjective | str,
    dtype: object = np.float32,
    seed: int = 20260801,
) -> SyntheticDataset

Generate deterministic unnormalized Gaussian embeddings.

Parameters:

Name Type Description Default
corpus_size int

Number of corpus vectors.

required
dimension int

Embedding dimension.

required
query_count int

Number of query vectors.

required
objective SearchObjective | str

Squared L2 or inner-product search.

required
dtype object

Either float32 or float64.

float32
seed int

Non-negative PCG64 seed.

20260801

Returns:

Type Description
SyntheticDataset

A deterministic unnormalized synthetic dataset.

Raises:

Type Description
InvalidSearchParameterError

If normalized cosine is requested.

Source code in src/vector_search_study/synthetic.py
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def make_gaussian_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    objective: SearchObjective | str,
    dtype: object = np.float32,
    seed: int = 20_260_801,
) -> SyntheticDataset:
    """Generate deterministic unnormalized Gaussian embeddings.

    Args:
        corpus_size: Number of corpus vectors.
        dimension: Embedding dimension.
        query_count: Number of query vectors.
        objective: Squared L2 or inner-product search.
        dtype: Either float32 or float64.
        seed: Non-negative PCG64 seed.

    Returns:
        A deterministic unnormalized synthetic dataset.

    Raises:
        InvalidSearchParameterError: If normalized cosine is requested.
    """
    resolved_objective = resolve_search_objective(objective)
    if resolved_objective.requires_normalization:
        raise InvalidSearchParameterError("use make_uniform_sphere_dataset for normalized cosine")
    size, dimensions, queries, resolved_dtype, resolved_seed = _validate_generator_inputs(
        corpus_size,
        dimension,
        query_count,
        dtype=dtype,
        seed=seed,
    )
    generator = np.random.Generator(np.random.PCG64(resolved_seed))
    corpus = generator.standard_normal((size, dimensions)).astype(resolved_dtype, copy=False)
    query_matrix = generator.standard_normal((queries, dimensions)).astype(resolved_dtype, copy=False)
    return _dataset(
        corpus,
        query_matrix,
        distribution="gaussian",
        seed=resolved_seed,
        objective=resolved_objective,
    )

make_uniform_sphere_dataset

make_uniform_sphere_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    dtype: object = np.float32,
    seed: int = 20260801,
) -> SyntheticDataset

Generate independent corpus and query vectors on the unit sphere.

Parameters:

Name Type Description Default
corpus_size int

Number of corpus vectors.

required
dimension int

Embedding dimension.

required
query_count int

Number of query vectors.

required
dtype object

Either float32 or float64.

float32
seed int

Non-negative PCG64 seed.

20260801

Returns:

Type Description
SyntheticDataset

A deterministic normalized synthetic dataset.

Source code in src/vector_search_study/synthetic.py
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def make_uniform_sphere_dataset(
    corpus_size: int,
    dimension: int,
    query_count: int,
    *,
    dtype: object = np.float32,
    seed: int = 20_260_801,
) -> SyntheticDataset:
    """Generate independent corpus and query vectors on the unit sphere.

    Args:
        corpus_size: Number of corpus vectors.
        dimension: Embedding dimension.
        query_count: Number of query vectors.
        dtype: Either float32 or float64.
        seed: Non-negative PCG64 seed.

    Returns:
        A deterministic normalized synthetic dataset.
    """
    size, dimensions, queries, resolved_dtype, resolved_seed = _validate_generator_inputs(
        corpus_size,
        dimension,
        query_count,
        dtype=dtype,
        seed=seed,
    )
    generator = np.random.Generator(np.random.PCG64(resolved_seed))
    corpus = _normal_matrix(generator, size, dimensions, resolved_dtype)
    query_matrix = _normal_matrix(generator, queries, dimensions, resolved_dtype)
    return _dataset(
        corpus,
        query_matrix,
        distribution="uniform_sphere",
        seed=resolved_seed,
        objective=SearchObjective.NORMALIZED_COSINE,
    )