Skip to content

Numpy Search

Vectorized NumPy exact search implementations.

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)

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")

score_matrix 

score_matrix(
    queries: FloatMatrix,
    corpus: FloatMatrix,
    objective: SearchObjective,
) -> FloatMatrix

Return a higher-is-better objective score matrix.

Source code in src/vector_search_study/numpy_search.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
def score_matrix(
    queries: FloatMatrix,
    corpus: FloatMatrix,
    objective: SearchObjective,
) -> FloatMatrix:
    """Return a higher-is-better objective score matrix."""
    products = queries @ corpus.T
    if objective is not SearchObjective.SQUARED_L2:
        return products
    query_norms = np.sum(queries * queries, axis=1, keepdims=True)
    corpus_norms = np.sum(corpus * corpus, axis=1, keepdims=True).T
    distances = np.maximum(query_norms + corpus_norms - 2.0 * products, 0.0)
    return np.asarray(-distances, dtype=queries.dtype, order="C")