Adding HannoyTransformer

src/sklearn_ann/kneighbors/hannoy.py

@@ -0,0 +1,118 @@
+# hannoy needs a filesystem path (LMDB-backed)
+import tempfile
+from itertools import count
+
+import numpy as np
+from hannoy import Database, Metric
+from scipy.sparse import csr_matrix
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.utils import Tags, TargetTags, TransformerTags
+from sklearn.utils.validation import validate_data
+
+from ..utils import TransformerChecksMixin
+
+# guards against panic in Rust, i.e. unrecoverable errors
+SUPPORTED_M = frozenset({4, 8, 12, 16, 24, 32})
+
+# due to hannoy sharing one global LMDB environement
+# using to avoid colliding
+_index_counter = count()
+
+class HannoyTransformer(TransformerChecksMixin, TransformerMixin, BaseEstimator):
+    # known issue where multiple Database instances silently share the first one's LMDB env
+
+    def __init__(
+        self,
+        n_neighbors=5,
+        *,
+        metric=None,
+        path=None,
+        m=16,
+        ef_construction=96,
+        ef_search=200,
+    ):
+        self.n_neighbors = n_neighbors
+        self.metric = metric
+        # LMDB directory for the index; if None = auto-create a temp dir
+        self.path = path
+        self.m = m
+        self.ef_construction = ef_construction
+        self.ef_search = ef_search
+
+    def fit(self, X, y=None):
+        X = validate_data(self, X, dtype=np.float32)
+        # guard to avoid panic abort
+        if self.m not in SUPPORTED_M:
+            raise ValueError(
+                f"m={self.m!r} is not supported by hannoy; "
+                f"choose one of {sorted(SUPPORTED_M)}."
+            )
+        self.n_samples_fit_ = X.shape[0]
+        metric = Metric.EUCLIDEAN if self.metric is None else self.metric
+        path = (
+            self.path if self.path is not None else tempfile.mkdtemp(prefix="hannoy_")
+        )
+        self._index_ = next(_index_counter) % 2**16
+
+        # metric is fixed for the entire database
+        self.hannoy_db_ = Database(path, metric)
+        with self.hannoy_db_.writer(
+            X.shape[1], index=self._index_, m=self.m, ef=self.ef_construction
+        ) as writer:
+            writer.add_items(list(range(self.n_samples_fit_)), X)
+        self.hannoy_reader_ = self.hannoy_db_.reader(index=self._index_)
+        return self
+
+    def transform(self, X):
+        # verify that fit was called and + that X has the right number of features
+        X = self._transform_checks(X, "hannoy_reader_", dtype=np.float32)
+        return self._transform(X)
+
+    def fit_transform(self, X, y=None):
+        return self.fit(X)._transform(X=None)
+
+    def _transform(self, X):
+        # how many points
+        n_samples_transform = self.n_samples_fit_ if X is None else X.shape[0]
+        n_neighbors = self.n_neighbors + 1
+        # pre allocating indicies for which points are neighbots
+        indices = np.empty((n_samples_transform, n_neighbors), dtype=np.int32)
+        distances = np.empty((n_samples_transform, n_neighbors), dtype=np.float32)
+
+
+        if X is None:
+            # by_item path
+            for i in range(n_samples_transform):
+                neighbours = self.hannoy_reader_.by_item(
+                    i, n=self.n_neighbors, ef_search=self.ef_search
+                )
+                indices[i, 0] = i
+                distances[i, 0] = 0.0
+                indices[i, 1:] = [j for j, _ in neighbours]
+                distances[i, 1:] = [d for _, d in neighbours]
+        else:
+            # by_vec path
+            for i, x in enumerate(X):
+                neighbours = self.hannoy_reader_.by_vec(
+                    x.tolist(), n=n_neighbors, ef_search=self.ef_search
+                )
+                indices[i] = [j for j, _ in neighbours]
+                distances[i] = [d for _, d in neighbours]
+
+        metric = Metric.EUCLIDEAN if self.metric is None else self.metric
+        if metric == Metric.EUCLIDEAN:
+            np.sqrt(distances, out=distances)
+
+        indptr = np.arange(0, n_samples_transform * n_neighbors + 1, n_neighbors)
+        return csr_matrix(
+            (distances.ravel(), indices.ravel(), indptr),
+            shape=(n_samples_transform, self.n_samples_fit_),
+        )
+
+    def __sklearn_tags__(self) -> Tags:
+        # metadata
+        return Tags(
+            estimator_type="transformer",
+            target_tags=TargetTags(required=False),
+            transformer_tags=TransformerTags(preserves_dtype=[np.float32]),
+        )

src/sklearn_ann/test_utils.py

@@ -24,6 +24,7 @@ class needs(Enum):
 
     annoy = ("annoy",)
     faiss = ("faiss-cpu", "faiss-gpu")
+    hannoy = ("hannoy",) 
     nmslib = ("nmslib",)
     pynndescent = ("pynndescent",)
 

tests/test_kneighbors/test_hannoy.py

@@ -0,0 +1,41 @@
+import numpy as np
+import pytest
+
+from sklearn_ann.test_utils import assert_row_close, needs
+
+try:
+    from hannoy import Metric
+
+    from sklearn_ann.kneighbors.hannoy import HannoyTransformer
+except ImportError:
+    pass
+
+
+@needs.hannoy
+def test_euclidean(random_small, random_small_pdists):
+    trans = HannoyTransformer(metric=Metric.EUCLIDEAN)
+    mat = trans.fit_transform(random_small)
+    euclidean_dist = random_small_pdists["euclidean"]
+    assert_row_close(mat, euclidean_dist)
+
+
+@needs.hannoy
+def test_cosine(random_small, random_small_pdists):
+    trans = HannoyTransformer(metric=Metric.COSINE)
+    mat = trans.fit_transform(random_small)
+    # hannoy's cosine metric returns (1 - cos_sim) / 2
+    cosine_dist = random_small_pdists["cosine"] / 2
+    assert_row_close(mat, cosine_dist)
+
+
+@needs.hannoy
+def test_transform_matches_fit_transform(random_small):
+    # fit_transform uses the by_item path; fit().transform() uses by_vec
+    # comparing the two paths
+    trans = HannoyTransformer(metric=Metric.EUCLIDEAN)
+    by_item = trans.fit_transform(random_small)
+    by_vec = trans.fit(random_small).transform(random_small)
+    by_item.sort_indices()
+    by_vec.sort_indices()
+    np.testing.assert_array_equal(by_item.indices, by_vec.indices)
+    np.testing.assert_allclose(by_item.data, by_vec.data, atol=1e-5)