mojolearn.refine#

mojolearn.refine(dataset, queries, candidates, k, numeric_mode=None, metric='sqeuclidean')[source]#

Re-rank candidate neighbors by exact distance.

Follows cuVS’s refine. For each query, computes the exact distance to each of its candidate rows of dataset and keeps the k smallest, ties broken by row id. Use it after an approximate search (for example IVFRaBitQIndex or IVFPQIndex with a larger n_neighbors) to get exact distances and a better ordering.

Parameters:

  • dataset (array-like of shape (n_samples, n_features)) – The indexed rows, converted to float32.

  • queries (array-like of shape (n_queries, n_features)) – Query rows, converted to float32.

  • candidates (array-like of int, shape (n_queries, n_candidates)) – Candidate row ids for each query. Negative ids are padding and are skipped, as are ids outside [0, n_samples).

  • k (int) – Number of neighbors to return per query.

  • numeric_mode (str or None, default None) – Numeric mode to run in; None uses the package default set with mojolearn.set_numeric_mode.

  • metric ({‘sqeuclidean’, ‘euclidean’}, default ‘sqeuclidean’) – ‘sqeuclidean’ returns squared Euclidean distances (cuVS L2Expanded); ‘euclidean’ returns their square roots (cuVS L2SqrtExpanded). The ids and their order are the same under both. Inner product is not supported.

Returns:

  • distances (ndarray of shape (n_queries, k), float32) – Exact distances, ascending.

  • indices (ndarray of shape (n_queries, k), int32) – Row ids of the neighbors. A query with fewer than k valid candidates is padded with distance inf and id -1.

Raises:

  • ValueError – If metric is not supported, the feature counts differ, candidates is not an integer array of shape (n_queries, n_candidates), or a candidate id does not fit in int32.

  • TypeError – If k is not an int.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> candidates = np.tile(np.arange(50), (10, 1))
>>> distances, indices = mojolearn.refine(X, X[:10], candidates, k=5)
>>> indices.shape
(10, 5)
On this page