redisbench-admin 0.11.64__py3-none-any.whl → 0.11.66__py3-none-any.whl

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/faiss_hnsw.py

@@ -0,0 +1,39 @@
+from __future__ import absolute_import
+import os
+import faiss
+import numpy as np
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+from ann_benchmarks.algorithms.faiss import Faiss
+
+
+class FaissHNSW(Faiss):
+    def __init__(self, metric, method_param):
+        self._metric = metric
+        self.method_param = method_param
+
+    def fit(self, X):
+        self.index = faiss.IndexHNSWFlat(len(X[0]), self.method_param["M"])
+        self.index.hnsw.efConstruction = self.method_param["efConstruction"]
+        self.index.verbose = True
+
+        if self._metric == 'angular':
+            X = X / np.linalg.norm(X, axis=1)[:, np.newaxis]
+        if X.dtype != np.float32:
+            X = X.astype(np.float32)
+
+        self.index.add(X)
+        faiss.omp_set_num_threads(1)
+
+    def set_query_arguments(self, ef):
+        faiss.cvar.hnsw_stats.reset()
+        self.index.hnsw.efSearch = ef
+
+    def get_additional(self):
+        return {"dist_comps": faiss.cvar.hnsw_stats.ndis}
+
+    def __str__(self):
+        return 'faiss (%s, ef: %d)' % (self.method_param, self.index.hnsw.efSearch)
+
+    def freeIndex(self):
+        del self.p

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/flann.py

@@ -0,0 +1,27 @@
+from __future__ import absolute_import
+import pyflann
+import numpy
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class FLANN(BaseANN):
+    def __init__(self, metric, target_precision):
+        self._target_precision = target_precision
+        self.name = 'FLANN(target_precision=%f)' % self._target_precision
+        self._metric = metric
+
+    def fit(self, X):
+        self._flann = pyflann.FLANN(
+            target_precision=self._target_precision,
+            algorithm='autotuned', log_level='info')
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+        self._flann.build_index(X)
+
+    def query(self, v, n):
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize([v], axis=1, norm='l2')[0]
+        if v.dtype != numpy.float32:
+            v = v.astype(numpy.float32)
+        return self._flann.nn_index(v, n)[0][0]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/hnswlib.py

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+import os
+import hnswlib
+import numpy as np
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class HnswLib(BaseANN):
+    def __init__(self, metric, method_param):
+        self.metric = {'angular': 'cosine', 'euclidean': 'l2'}[metric]
+        self.method_param = method_param
+        # print(self.method_param,save_index,query_param)
+        # self.ef=query_param['ef']
+        self.name = 'hnswlib (%s)' % (self.method_param)
+
+    def fit(self, X):
+        # Only l2 is supported currently
+        self.p = hnswlib.Index(space=self.metric, dim=len(X[0]))
+        self.p.init_index(max_elements=len(X),
+                          ef_construction=self.method_param["efConstruction"],
+                          M=self.method_param["M"])
+        data_labels = np.arange(len(X))
+        self.p.add_items(np.asarray(X), data_labels)
+        self.p.set_num_threads(1)
+
+    def set_query_arguments(self, ef):
+        self.p.set_ef(ef)
+
+    def query(self, v, n):
+        # print(np.expand_dims(v,axis=0).shape)
+        # print(self.p.knn_query(np.expand_dims(v,axis=0), k = n)[0])
+        return self.p.knn_query(np.expand_dims(v, axis=0), k=n)[0][0]
+
+    def freeIndex(self):
+        del self.p

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/kdtree.py

@@ -0,0 +1,22 @@
+from __future__ import absolute_import
+import sklearn.neighbors
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class KDTree(BaseANN):
+    def __init__(self, metric, leaf_size=20):
+        self._leaf_size = leaf_size
+        self._metric = metric
+        self.name = 'KDTree(leaf_size=%d)' % self._leaf_size
+
+    def fit(self, X):
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+        self._tree = sklearn.neighbors.KDTree(X, leaf_size=self._leaf_size)
+
+    def query(self, v, n):
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize([v], axis=1, norm='l2')[0]
+        dist, ind = self._tree.query([v], k=n)
+        return ind[0]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/kgraph.py

@@ -0,0 +1,39 @@
+from __future__ import absolute_import
+import os
+import numpy
+import pykgraph
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class KGraph(BaseANN):
+    def __init__(self, metric, index_params, save_index):
+        metric = str(metric)
+        self.name = 'KGraph(%s)' % (metric)
+        self._metric = metric
+        self._index_params = index_params
+        self._save_index = save_index
+
+    def fit(self, X):
+        if X.dtype != numpy.float32:
+            X = X.astype(numpy.float32)
+        self._kgraph = pykgraph.KGraph(X, self._metric)
+        path = os.path.join(INDEX_DIR, 'kgraph-index-%s' % self._metric)
+        if os.path.exists(path):
+            self._kgraph.load(path)
+        else:
+            # iterations=30, L=100, delta=0.002, recall=0.99, K=25)
+            self._kgraph.build(**self._index_params)
+            if not os.path.exists(INDEX_DIR):
+                os.makedirs(INDEX_DIR)
+            self._kgraph.save(path)
+
+    def set_query_arguments(self, P):
+        self._P = P
+
+    def query(self, v, n):
+        if v.dtype != numpy.float32:
+            v = v.astype(numpy.float32)
+        result = self._kgraph.search(
+            numpy.array([v]), K=n, threads=1, P=self._P)
+        return result[0]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/lshf.py

@@ -0,0 +1,25 @@
+from __future__ import absolute_import
+import sklearn.neighbors
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class LSHF(BaseANN):
+    def __init__(self, metric, n_estimators=10, n_candidates=50):
+        self.name = 'LSHF(n_est=%d, n_cand=%d)' % (n_estimators, n_candidates)
+        self._metric = metric
+        self._n_estimators = n_estimators
+        self._n_candidates = n_candidates
+
+    def fit(self, X):
+        self._lshf = sklearn.neighbors.LSHForest(
+            n_estimators=self._n_estimators, n_candidates=self._n_candidates)
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+        self._lshf.fit(X)
+
+    def query(self, v, n):
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize([v], axis=1, norm='l2')[0]
+        return self._lshf.kneighbors([v], return_distance=False,
+                                     n_neighbors=n)[0]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/milvus.py

@@ -0,0 +1,99 @@
+from __future__ import absolute_import
+from sqlite3 import paramstyle
+from pymilvus import (
+    connections,
+    utility,
+    FieldSchema,
+    CollectionSchema,
+    DataType,
+    IndexType,
+    Collection,
+)
+import numpy
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+import sys
+
+
+class Milvus(BaseANN):
+    def __init__(self, metric, dim, conn_params, index_type, method_params):
+        self._host = conn_params['host']
+        self._port = conn_params['port'] # 19530
+        self._index_type = index_type
+        self._method_params = method_params
+        self._metric = {'angular': 'IP', 'euclidean': 'L2'}[metric]
+        self._query_params = dict()
+        connections.connect(host=conn_params['host'], port=conn_params['port'])
+        try:
+            fields = [
+                FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
+                FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
+            ]
+            schema = CollectionSchema(fields)
+            if utility.has_collection('milvus'):
+                self._milvus = Collection('milvus')
+            else:
+                self._milvus = Collection('milvus', schema)
+        except:
+            self._milvus = Collection('milvus')
+        print('initialization completed!')
+    
+    def fit(self, X, offset=0, limit=None):
+        limit = limit if limit else len(X)
+        X = X[offset:limit]
+        if self._metric == 'IP':
+            X = sklearn.preprocessing.normalize(X)
+
+        X = X.tolist()
+        bulk_size = 1000 * 1024 * 1024 // (sys.getsizeof(X[0])) # approximation for milvus insert limit (1024MB)
+        for bulk in [X[i: i+bulk_size] for i in range(0, len(X), bulk_size)]:
+            print(f'inserting vectors {offset} to {offset + len(bulk) - 1}')
+            self._milvus.insert([list(range(offset, offset + len(bulk))), bulk])
+            offset += len(bulk)
+
+        if not self._milvus.has_index():
+            print('indexing...', end=' ')
+            try:
+                self._milvus.create_index('vector', {'index_type': self._index_type, 'metric_type':self._metric, 'params':self._method_params})
+                print('done!')
+            except:
+                print('failed!')
+        
+
+    def set_query_arguments(self, param):
+        if self._milvus.has_index():
+            print('waiting for index... ', end='')
+            if utility.wait_for_index_building_complete('milvus', 'vector'):
+                print('done!')
+                self._milvus.load()
+                print('waiting for data to be loaded... ', end='')
+                utility.wait_for_loading_complete('milvus')
+                print('done!')
+            else: raise Exception('index has error')
+        else: raise Exception('index is missing')
+        if 'IVF_' in self._index_type:
+            if param > self._method_params['nlist']:
+                print('warning! nprobe > nlist')
+                param = self._method_params['nlist']
+            self._query_params['nprobe'] = param
+        if 'HNSW' in self._index_type:
+            self._query_params['ef'] = param
+
+    def query(self, v, n):
+        if self._metric == 'IP':
+            v /= numpy.linalg.norm(v)
+        v = v.tolist()
+        results = self._milvus.search([v], 'vector', {'metric_type':self._metric, 'params':self._query_params}, limit=n)
+        if not results:
+            return []  # Seems to happen occasionally, not sure why
+        result_ids = [result.id for result in results[0]]
+        return result_ids
+
+    def __str__(self):
+        return 'Milvus(index_type=%s, method_params=%s, query_params=%s)' % (self._index_type, str(self._method_params), str(self._query_params))
+
+    def freeIndex(self):
+        utility.drop_collection("mlivus")
+
+    def done(self):
+        connections.disconnect('default')

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/mrpt.py

@@ -0,0 +1,41 @@
+from __future__ import absolute_import
+import numpy
+import sklearn.preprocessing
+import mrpt
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class MRPT(BaseANN):
+    def __init__(self, metric, count):
+        self._metric = metric
+        self._k = count
+
+    def fit(self, X):
+        if X.dtype != numpy.float32:
+            X = X.astype(numpy.float32)
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+
+        self._index_autotuned = mrpt.MRPTIndex(X)
+        self._index_autotuned.build_autotune_sample(
+            target_recall=None, k=self._k, n_test=1000)
+
+    def set_query_arguments(self, target_recall):
+        self._target_recall = target_recall
+        self._index = self._index_autotuned.subset(target_recall)
+        self._par = self._index.parameters()
+
+    def query(self, v, n):
+        if v.dtype != numpy.float32:
+            v = v.astype(numpy.float32)
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize(
+                v.reshape(1, -1), axis=1, norm='l2').flatten()
+        return self._index.ann(v)
+
+    def __str__(self):
+        str_template = ('MRPT(target recall=%.3f, trees=%d, depth=%d, vote '
+                        'threshold=%d, estimated recall=%.3f)')
+        return str_template % (self._target_recall, self._par['n_trees'],
+                               self._par['depth'], self._par['votes'],
+                               self._par['estimated_recall'])

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/n2.py

@@ -0,0 +1,28 @@
+from __future__ import absolute_import
+import n2
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class N2(BaseANN):
+    def __init__(self, metric, method_param):
+        self._metric = metric
+        self._m = method_param['M']
+        self._m0 = self._m * 2
+        self._ef_construction = method_param['efConstruction']
+        self._n_threads = 1
+        self._ef_search = -1
+
+    def fit(self, X):
+        self._n2 = n2.HnswIndex(X.shape[1], self._metric)
+        for x in X:
+            self._n2.add_data(x)
+        self._n2.build(m=self._m, max_m0=self._m0, ef_construction=self._ef_construction, n_threads=self._n_threads, graph_merging='merge_level0')
+
+    def set_query_arguments(self, ef):
+        self._ef_search = ef
+
+    def query(self, v, n):
+        return self._n2.search_by_vector(v, n, self._ef_search)
+
+    def __str__(self):
+        return "N2 (M%d_efCon%d)" % (self._m, self._ef_construction)

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/nearpy.py

@@ -0,0 +1,48 @@
+from __future__ import absolute_import
+import nearpy
+from nearpy.filters import NearestFilter
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class NearPy(BaseANN):
+    def __init__(self, metric, n_bits, hash_counts):
+        self._n_bits = n_bits
+        self._hash_counts = hash_counts
+        self._metric = metric
+        self._filter = NearestFilter(10)
+        self.name = 'NearPy(n_bits=%d, hash_counts=%d)' % (
+            self._n_bits, self._hash_counts)
+
+    def fit(self, X):
+        hashes = []
+
+        for k in range(self._hash_counts):
+            nearpy_rbp = nearpy.hashes.RandomBinaryProjections(
+                'rbp_%d' % k, self._n_bits)
+            hashes.append(nearpy_rbp)
+
+        if self._metric == 'euclidean':
+            dist = nearpy.distances.EuclideanDistance()
+            self._nearpy_engine = nearpy.Engine(
+                X.shape[1],
+                lshashes=hashes,
+                distance=dist)
+        else:  # Default (angular) = Cosine distance
+            self._nearpy_engine = nearpy.Engine(
+                X.shape[1],
+                lshashes=hashes,
+                vector_filters=[self._filter])
+
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+        for i, x in enumerate(X):
+            self._nearpy_engine.store_vector(x, i)
+
+    def query(self, v, n):
+        # XXX: This feels like an unpleasant hack, but it's not clear how to do
+        # better without making changes to NearPy
+        self._filter.N = n
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize([v], axis=1, norm='l2')[0]
+        return [y for x, y, z in self._nearpy_engine.neighbours(v)]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/nmslib.py

@@ -0,0 +1,74 @@
+from __future__ import absolute_import
+import os
+import nmslib
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class NmslibReuseIndex(BaseANN):
+    @staticmethod
+    def encode(d):
+        return ["%s=%s" % (a, b) for (a, b) in d.items()]
+
+    def __init__(self, metric, method_name, index_param, query_param):
+        self._nmslib_metric = {
+            'angular': 'cosinesimil', 'euclidean': 'l2'}[metric]
+        self._method_name = method_name
+        self._save_index = False
+        self._index_param = NmslibReuseIndex.encode(index_param)
+        if query_param is not False:
+            self._query_param = NmslibReuseIndex.encode(query_param)
+            self.name = ('Nmslib(method_name={}, index_param={}, '
+                         'query_param={})'.format(self._method_name,
+                                                  self._index_param,
+                                                  self._query_param))
+        else:
+            self._query_param = None
+            self.name = 'Nmslib(method_name=%s, index_param=%s)' % (
+                self._method_name, self._index_param)
+
+        self._index_name = os.path.join(INDEX_DIR, "nmslib_%s_%s_%s" % (
+            self._method_name, metric, '_'.join(self._index_param)))
+
+        d = os.path.dirname(self._index_name)
+        if not os.path.exists(d):
+            os.makedirs(d)
+
+    def fit(self, X):
+        if self._method_name == 'vptree':
+            # To avoid this issue: terminate called after throwing an instance
+            # of 'std::runtime_error'
+            # what():  The data size is too small or the bucket size is too
+            # big. Select the parameters so that <total # of records> is NOT
+            # less than <bucket size> * 1000
+            # Aborted (core dumped)
+            self._index_param.append('bucketSize=%d' %
+                                     min(int(X.shape[0] * 0.0005), 1000))
+
+        self._index = nmslib.init(
+            space=self._nmslib_metric, method=self._method_name)
+        self._index.addDataPointBatch(X)
+
+        if os.path.exists(self._index_name):
+            print('Loading index from file')
+            self._index.loadIndex(self._index_name)
+        else:
+            self._index.createIndex(self._index_param)
+            if self._save_index:
+                self._index.saveIndex(self._index_name)
+        if self._query_param is not None:
+            self._index.setQueryTimeParams(self._query_param)
+
+    def set_query_arguments(self, ef):
+        if self._method_name == 'hnsw' or self._method_name == 'sw-graph':
+            self._index.setQueryTimeParams(["efSearch=%s" % (ef)])
+
+    def query(self, v, n):
+        ids, distances = self._index.knnQuery(v, n)
+        return ids
+
+    def batch_query(self, X, n):
+        self.res = self._index.knnQueryBatch(X, n)
+
+    def get_batch_results(self):
+        return [x for x, _ in self.res]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/onng_ngt.py

@@ -0,0 +1,100 @@
+from __future__ import absolute_import
+import sys
+import os
+import ngtpy
+import numpy as np
+import subprocess
+import time
+from ann_benchmarks.algorithms.base import BaseANN
+from ann_benchmarks.constants import INDEX_DIR
+
+
+class ONNG(BaseANN):
+    def __init__(self, metric, object_type, epsilon, param):
+        metrics = {'euclidean': '2', 'angular': 'C'}
+        self._edge_size = int(param['edge'])
+        self._outdegree = int(param['outdegree'])
+        self._indegree = int(param['indegree'])
+        self._metric = metrics[metric]
+        self._object_type = object_type
+        self._edge_size_for_search = int(param['search_edge']) if 'search_edge' in param.keys() else 0
+        self._tree_disabled = (param['tree'] == False) if 'tree' in param.keys() else False
+        self._refine_enabled = (param['refine'] == True) if 'refine' in param.keys() else False
+        self._build_time_limit = 4
+        self._epsilon = epsilon
+        print('ONNG: edge_size=' + str(self._edge_size))
+        print('ONNG: outdegree=' + str(self._outdegree))
+        print('ONNG: indegree=' + str(self._indegree))
+        print('ONNG: edge_size_for_search=' + str(self._edge_size_for_search))
+        print('ONNG: epsilon=' + str(self._epsilon))
+        print('ONNG: metric=' + metric)
+        print('ONNG: object_type=' + object_type)
+
+    def fit(self, X):
+        print('ONNG: start indexing...')
+        dim = len(X[0])
+        print('ONNG: # of data=' + str(len(X)))
+        print('ONNG: dimensionality=' + str(dim))
+        index_dir = 'indexes'
+        if not os.path.exists(index_dir):
+            os.makedirs(index_dir)
+        index = os.path.join(
+            index_dir,
+            'ONNG-{}-{}-{}'.format(self._edge_size, self._outdegree,
+                                   self._indegree))
+        anngIndex = os.path.join(index_dir, 'ANNG-' + str(self._edge_size))
+        print('ONNG: index=' + index)
+        if (not os.path.exists(index)) and (not os.path.exists(anngIndex)):
+            print('ONNG: create ANNG')
+            t = time.time()
+            args = ['ngt', 'create', '-it', '-p8', '-b500', '-ga', '-of',
+                    '-D' + self._metric, '-d' + str(dim),
+                    '-E' + str(self._edge_size),
+                    '-S' + str(self._edge_size_for_search),
+                    '-e' + str(self._epsilon), '-P0', '-B30',
+                    '-T' + str(self._build_time_limit), anngIndex]
+            subprocess.call(args)
+            idx = ngtpy.Index(path=anngIndex)
+            idx.batch_insert(X, num_threads=24, debug=False)
+            print('ONNG: ANNG construction time(sec)=' + str(time.time() - t))
+            t = time.time()
+            if self._refine_enabled:
+                idx.refine_anng(epsilon=self._epsilon, num_of_edges=self._edge_size,
+                                num_of_explored_edges=self._edge_size_for_search)
+            print('ONNG: RNNG construction time(sec)=' + str(time.time() - t))
+            idx.save()
+            idx.close()
+        if not os.path.exists(index):
+            print('ONNG: degree adjustment')
+            t = time.time()
+            args = ['ngt', 'reconstruct-graph', '-mS',
+                    '-o ' + str(self._outdegree),
+                    '-i ' + str(self._indegree), anngIndex, index]
+            subprocess.call(args)
+            print('ONNG: degree adjustment time(sec)=' + str(time.time() - t))
+        if os.path.exists(index):
+            print('ONNG: index already exists! ' + str(index))
+            t = time.time()
+            print(self._tree_disabled)
+            self.index = ngtpy.Index(index, read_only=True, tree_disabled=self._tree_disabled)
+            self.indexName = index
+            print('ONNG: open time(sec)=' + str(time.time() - t))
+        else:
+            print('ONNG: something wrong.')
+        print('ONNG: end of fit')
+
+    def set_query_arguments(self, parameters):
+        epsilon, edge_size = parameters
+        print("ONNG: edge_size=" + str(edge_size))
+        print("ONNG: epsilon=" + str(epsilon))
+        self.name = 'ONNG-NGT(%s, %s, %s, %s, %1.3f)' % (
+            self._edge_size, self._outdegree,
+            self._indegree, edge_size, epsilon)
+        epsilon = epsilon - 1.0
+        self.index.set(epsilon=epsilon, edge_size=edge_size)
+
+    def query(self, v, n):
+        return self.index.search(v, n, with_distance=False)
+
+    def freeIndex(self):
+        print('ONNG: free')