Skip to content

Sklearn Search

Exact scikit-learn nearest-neighbor adapters.

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

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