redisbench-admin 0.11.63__py3-none-any.whl → 0.11.65__py3-none-any.whl

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

@@ -0,0 +1,107 @@
+import logging
+from time import sleep
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import bulk
+
+from ann_benchmarks.algorithms.base import BaseANN
+
+from .elasticsearch import es_wait
+
+from tqdm import tqdm
+
+# Configure the logger.
+logging.getLogger("elasticsearch").setLevel(logging.WARN)
+
+class OpenSearchKNN(BaseANN):
+    def __init__(self, metric, dimension, method_param):
+        self.metric = {"angular": "cosinesimil", "euclidean": "l2"}[metric]
+        self.dimension = dimension
+        self.method_param = method_param
+        self.param_string = "-".join(k+"-"+str(v) for k,v in self.method_param.items()).lower()
+        self.name = f"os-{self.param_string}"
+        self.es = Elasticsearch(["http://localhost:9200"])
+        es_wait()
+
+    def fit(self, X):
+        body = {
+            "settings": {
+                "index": {"knn": True},
+                "number_of_shards": 1, 
+                "number_of_replicas": 0,
+                "refresh_interval": -1
+            }
+        }
+
+        mapping = {
+            "properties": {
+                "id": {"type": "keyword", "store": True},
+                "vec": {
+                    "type": "knn_vector", 
+                    "dimension": self.dimension,
+                    "method": {
+                        "name": "hnsw",
+                        "space_type": self.metric,
+                        "engine": "nmslib",
+                        "parameters": {
+                            "ef_construction": self.method_param["efConstruction"],
+                            "m": self.method_param["M"]
+                        }
+                    }
+                }
+            }
+        }
+            
+        self.es.indices.create(self.name, body=body)
+        self.es.indices.put_mapping(mapping, self.name)
+
+        print("Uploading data to the Index:", self.name)
+        def gen():
+            for i, vec in enumerate(tqdm(X)):
+                yield { "_op_type": "index", "_index": self.name, "vec": vec.tolist(), 'id': str(i + 1) }
+
+        (_, errors) = bulk(self.es, gen(), chunk_size=500, max_retries=2, request_timeout=10)
+        assert len(errors) == 0, errors
+          
+        print("Force Merge...")
+        self.es.indices.forcemerge(self.name, max_num_segments=1, request_timeout=1000)
+               
+        print("Refreshing the Index...")
+        self.es.indices.refresh(self.name, request_timeout=1000)
+       
+        print("Running Warmup API...")
+        res = urlopen(Request("http://localhost:9200/_plugins/_knn/warmup/"+self.name+"?pretty"))
+        print(res.read().decode("utf-8"))
+
+    def set_query_arguments(self, ef):
+        body = {
+            "settings": {
+                "index": {"knn.algo_param.ef_search": ef}
+            }
+        }
+        self.es.indices.put_settings(body=body)
+
+    def query(self, q, n):
+        body = {
+            "query": {
+                "knn": {
+                    "vec": {"vector": q.tolist(), "k": n}
+                }
+            }
+        }
+
+        res = self.es.search(index=self.name, body=body, size=n, _source=False, docvalue_fields=['id'],
+                             stored_fields="_none_", filter_path=["hits.hits.fields.id"], request_timeout=10)
+        
+        return [int(h['fields']['id'][0]) - 1 for h in res['hits']['hits']]
+
+    def batch_query(self, X, n):
+        self.batch_res = [self.query(q, n) for q in X]
+
+    def get_batch_results(self):
+        return self.batch_res
+    
+    def freeIndex(self):
+        self.es.indices.delete(index=self.name)

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

@@ -0,0 +1,79 @@
+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 PANNG(BaseANN):
+    def __init__(self, metric, object_type, param):
+        metrics = {'euclidean': 'L2', 'angular': 'Cosine'}
+        self._edge_size = int(param['edge'])
+        self._pathadj_size = int(param['pathadj'])
+        self._edge_size_for_search = int(param['searchedge'])
+        self._metric = metrics[metric]
+        self._object_type = object_type
+        print('PANNG: edge_size=' + str(self._edge_size))
+        print('PANNG: pathadj_size=' + str(self._pathadj_size))
+        print('PANNG: edge_size_for_search=' + str(self._edge_size_for_search))
+        print('PANNG: metric=' + metric)
+        print('PANNG: object_type=' + object_type)
+
+    def fit(self, X):
+        print('PANNG: start indexing...')
+        dim = len(X[0])
+        print('PANNG: # of data=' + str(len(X)))
+        print('PANNG: Dimensionality=' + str(dim))
+        index_dir = 'indexes'
+        if not os.path.exists(index_dir):
+            os.makedirs(index_dir)
+        index = os.path.join(
+            index_dir,
+            'PANNG-' + str(self._edge_size) + '-' + str(self._pathadj_size))
+        print(index)
+        if os.path.exists(index):
+            print('PANNG: index already exists! ' + str(index))
+        else:
+            t0 = time.time()
+            ngtpy.create(path=index, dimension=dim,
+                         edge_size_for_creation=self._edge_size,
+                         distance_type=self._metric,
+                         object_type=self._object_type)
+            idx = ngtpy.Index(path=index)
+            idx.batch_insert(X, num_threads=24, debug=False)
+            idx.save()
+            idx.close()
+            if self._pathadj_size > 0:
+                print('PANNG: path adjustment')
+                args = ['ngt', 'prune', '-s ' + str(self._pathadj_size),
+                        index]
+                subprocess.call(args)
+            indexingtime = time.time() - t0
+            print('PANNG: indexing, adjustment and saving time(sec)={}'
+                  .format(indexingtime))
+        t0 = time.time()
+        self.index = ngtpy.Index(path=index, read_only=True)
+        opentime = time.time() - t0
+        print('PANNG: open time(sec)=' + str(opentime))
+
+    def set_query_arguments(self, epsilon):
+        print("PANNG: epsilon=" + str(epsilon))
+        self._epsilon = epsilon - 1.0
+        self.name = 'PANNG-NGT(%d, %d, %d, %1.3f)' % (
+            self._edge_size,
+            self._pathadj_size,
+            self._edge_size_for_search,
+            self._epsilon + 1.0)
+
+    def query(self, v, n):
+        results = self.index.search(
+            v, n, self._epsilon, self._edge_size_for_search,
+            with_distance=False)
+        return results
+
+    def freeIndex(self):
+        print('PANNG: free')

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

@@ -0,0 +1,39 @@
+from __future__ import absolute_import
+from sqlite3 import paramstyle
+from ann_benchmarks.algorithms.base import BaseANN
+import sys
+import pinecone
+
+class Pinecone(BaseANN):
+    def __init__(self, metric, dim, conn_params, type):
+        pinecone.init(api_key=conn_params['auth'])
+        m = {'angular': 'cosine', 'euclidean': 'euclidean'}[metric]
+        self.name = 'ann-benchmark'
+        if self.name not in pinecone.list_indexes():
+            pinecone.create_index(self.name, dimension=dim, metric=m,
+                                  index_type=type, shards=int(conn_params["shards"]), )
+        self.index = pinecone.Index(self.name)
+
+    def fit(self, X, offset=0, limit=None):
+        limit = limit if limit else len(X)
+
+        bulk = [(str(i), X[i].tolist()) for i in range(offset, limit)]
+        # approximation for pinecone insert limit (2MB or 1000 vectors)
+        batch_size = min(1000, 2 * 1024 * 1024 // (sys.getsizeof(bulk[-1]))) # bulk[-1] should be the largest (longest name)
+
+        for batch in [bulk[i: i+batch_size] for i in range(0, len(bulk), batch_size)]:
+            # print(f'inserting vectors {batch[0][0]} to {batch[-1][0]}')
+            self.index.upsert(batch)
+        
+        # print(self.index.describe_index_stats())
+        # print(pinecone.describe_index(self.name))
+    
+    def query(self, v, n):
+        res = self.index.query(v.tolist(), top_k=n)
+        return [int(e['id']) for e in res['matches']]
+    
+    def freeIndex(self):
+        pinecone.delete_index(self.name)
+
+    def __str__(self):
+        return f'Pinecone({pinecone.describe_index(self.name)})'

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

@@ -0,0 +1,45 @@
+from __future__ import absolute_import
+import puffinn
+from ann_benchmarks.algorithms.base import BaseANN
+import numpy
+
+class Puffinn(BaseANN):
+    def __init__(self, metric, space=10**6, hash_function="fht_crosspolytope", hash_source='pool', hash_args=None):
+        if metric not in ['jaccard', 'angular']:
+            raise NotImplementedError(
+                    "Puffinn doesn't support metric %s" % metric)
+        self.metric = metric
+        self.space = space
+        self.hash_function = hash_function
+        self.hash_source = hash_source
+        self.hash_args = hash_args
+
+    def fit(self, X):
+        if self.metric == 'angular':
+            dimensions = len(X[0])
+        else:
+            dimensions = 0
+            for x in X:
+                dimensions = max(dimensions, max(x)+1)
+
+        if self.hash_args:
+            self.index = puffinn.Index(self.metric, dimensions, self.space,\
+                    hash_function=self.hash_function, hash_source=self.hash_source,\
+                    hash_args=self.hash_args)
+        else:
+            self.index = puffinn.Index(self.metric, dimensions, self.space,\
+                    hash_function=self.hash_function, hash_source=self.hash_source)
+        for i, x in enumerate(X):
+            x = x.tolist()
+            self.index.insert(x)
+        self.index.rebuild()
+
+    def set_query_arguments(self, recall):
+        self.recall = recall
+
+    def query(self, v, n):
+        v = v.tolist()
+        return self.index.search(v, n, self.recall)
+
+    def __str__(self):
+        return 'PUFFINN(space=%d, recall=%f, hf=%s, hashsource=%s)' % (self.space, self.recall, self.hash_function, self.hash_source)

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

@@ -0,0 +1,115 @@
+from __future__ import absolute_import
+import pynndescent
+from ann_benchmarks.algorithms.base import BaseANN
+import numpy as np
+import scipy.sparse
+
+
+class PyNNDescent(BaseANN):
+    def __init__(self, metric, index_param_dict, n_search_trees=1):
+        if "n_neighbors" in index_param_dict:
+            self._n_neighbors = int(index_param_dict["n_neighbors"])
+        else:
+            self._n_neighbors = 30
+
+        if "pruning_degree_multiplier" in index_param_dict:
+            self._pruning_degree_multiplier = float(
+                index_param_dict["pruning_degree_multiplier"]
+            )
+        else:
+            self._pruning_degree_multiplier = 1.5
+
+        if "diversify_prob" in index_param_dict:
+            self._diversify_prob = float(index_param_dict["diversify_prob"])
+        else:
+            self._diversify_prob = 1.0
+
+        if "leaf_size" in index_param_dict:
+            self._leaf_size = int(index_param_dict["leaf_size"])
+        else:
+            leaf_size = 32
+
+        self._n_search_trees = int(n_search_trees)
+
+        self._pynnd_metric = {
+            "angular": "dot",
+            # 'angular': 'cosine',
+            "euclidean": "euclidean",
+            "hamming": "hamming",
+            "jaccard": "jaccard",
+        }[metric]
+
+    def _sparse_convert_for_fit(self, X):
+        lil_data = []
+        self._n_cols = 1
+        self._n_rows = len(X)
+        for i in range(self._n_rows):
+            lil_data.append([1] * len(X[i]))
+            if max(X[i]) + 1 > self._n_cols:
+                self._n_cols = max(X[i]) + 1
+
+        result = scipy.sparse.lil_matrix(
+            (self._n_rows, self._n_cols), dtype=np.int
+        )
+        result.rows[:] = list(X)
+        result.data[:] = lil_data
+        return result.tocsr()
+
+    def _sparse_convert_for_query(self, v):
+        result = scipy.sparse.csr_matrix((1, self._n_cols), dtype=np.int)
+        result.indptr = np.array([0, len(v)])
+        result.indices = np.array(v).astype(np.int32)
+        result.data = np.ones(len(v), dtype=np.int)
+        return result
+
+    def fit(self, X):
+        if self._pynnd_metric == "jaccard":
+            # Convert to sparse matrix format
+            X = self._sparse_convert_for_fit(X)
+
+        self._index = pynndescent.NNDescent(
+            X,
+            n_neighbors=self._n_neighbors,
+            metric=self._pynnd_metric,
+            low_memory=True,
+            leaf_size=self._leaf_size,
+            pruning_degree_multiplier=self._pruning_degree_multiplier,
+            diversify_prob=self._diversify_prob,
+            n_search_trees=self._n_search_trees,
+            compressed=True,
+            verbose=True,
+        )
+        if hasattr(self._index, "prepare"):
+            self._index.prepare()
+        else:
+            self._index._init_search_graph()
+            if self._index._is_sparse:
+                if hasattr(self._index, "_init_sparse_search_function"):
+                    self._index._init_sparse_search_function()
+            else:
+                if hasattr(self._index, "_init_search_function"):
+                    self._index._init_search_function()
+
+    def set_query_arguments(self, epsilon=0.1):
+        self._epsilon = float(epsilon)
+
+    def query(self, v, n):
+        if self._pynnd_metric == "jaccard":
+            # convert index array to sparse matrix format and query
+            v = self._sparse_convert_for_query(v)
+            ind, dist = self._index.query(v, k=n, epsilon=self._epsilon)
+        else:
+            ind, dist = self._index.query(
+                v.reshape(1, -1).astype("float32"), k=n, epsilon=self._epsilon
+            )
+        return ind[0]
+
+    def __str__(self):
+        str_template = "PyNNDescent(n_neighbors=%d, pruning_mult=%.2f, diversify_prob=%.3f, epsilon=%.3f, leaf_size=%02d)"
+        return str_template % (
+            self._n_neighbors,
+            self._pruning_degree_multiplier,
+            self._diversify_prob,
+            self._epsilon,
+            self._leaf_size,
+        )

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

@@ -0,0 +1,102 @@
+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 QG(BaseANN):
+    def __init__(self, metric, object_type, epsilon, param):
+        metrics = {'euclidean': '2', 'angular': 'E'}
+        self._edge_size = int(param['edge'])
+        self._outdegree = int(param['outdegree'])
+        self._indegree = int(param['indegree'])
+        self._max_edge_size = int(param['max_edge']) if 'max_edge' in param.keys() else 128
+        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 -2
+        self._tree_disabled = (param['tree'] == False) if 'tree' in param.keys() else False
+        self._build_time_limit = 4
+        self._epsilon = epsilon
+        print('QG: edge_size=' + str(self._edge_size))
+        print('QG: outdegree=' + str(self._outdegree))
+        print('QG: indegree=' + str(self._indegree))
+        print('QG: edge_size_for_search=' + str(self._edge_size_for_search))
+        print('QG: epsilon=' + str(self._epsilon))
+        print('QG: metric=' + metric)
+        print('QG: object_type=' + object_type)
+
+    def fit(self, X):
+        print('QG: start indexing...')
+        dim = len(X[0])
+        print('QG: # of data=' + str(len(X)))
+        print('QG: 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('QG: index=' + index)
+        if (not os.path.exists(index)) and (not os.path.exists(anngIndex)):
+            print('QG: create ANNG')
+            t = time.time()
+            args = ['ngt', 'create', '-it', '-p8', '-b500', '-ga', '-of',
+                    '-D' + self._metric, '-d' + str(dim),
+                    '-E' + str(self._edge_size), '-S40',
+                    '-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)
+            idx.save()
+            idx.close()
+            print('QG: ANNG construction time(sec)=' + str(time.time() - t))
+        if not os.path.exists(index):
+            print('QG: degree adjustment')
+            t = time.time()
+            args = ['ngt', 'reconstruct-graph', '-mS',
+                    '-E ' + str(self._outdegree),
+                    '-o ' + str(self._outdegree),
+                    '-i ' + str(self._indegree), anngIndex, index]
+            subprocess.call(args)
+            print('QG: degree adjustment time(sec)=' + str(time.time() - t))
+        if not os.path.exists(index + '/qg'):
+            print('QG: quantization')
+            t = time.time()
+            args = ['ngtqg', 'quantize', index]
+            subprocess.call(args)
+            print('QG: quantization time(sec)=' + str(time.time() - t))
+        if os.path.exists(index):
+            print('QG: index already exists! ' + str(index))
+            t = time.time()
+            self.index = ngtpy.QuantizedIndex(index, self._max_edge_size)
+            self.index.set_with_distance(False)
+            self.indexName = index
+            print('QG: open time(sec)=' + str(time.time() - t))
+        else:
+            print('QG: something wrong.')
+        print('QG: end of fit')
+
+    def set_query_arguments(self, parameters):
+        result_expansion, epsilon = parameters
+        print("QG: result_expansion=" + str(result_expansion))
+        print("QG: epsilon=" + str(epsilon))
+        self.name = 'QG-NGT(%s, %s, %s, %s, %s, %1.3f)' % (
+            self._edge_size, self._outdegree,
+            self._indegree, self._max_edge_size,
+            epsilon,
+            result_expansion)
+        epsilon = epsilon - 1.0
+        self.index.set(epsilon=epsilon, result_expansion=result_expansion)
+
+    def query(self, v, n):
+        return self.index.search(v, n)
+
+    def freeIndex(self):
+        print('QG: free')

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

@@ -0,0 +1,90 @@
+from __future__ import absolute_import
+from optparse import Values
+from redis import Redis
+from redis.cluster import RedisCluster
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+import math
+
+
+class RediSearch(BaseANN):
+    def __init__(self, algo, metric, conn_params, method_param):
+        self.metric = {'angular': 'cosine', 'euclidean': 'l2'}[metric]
+        self.method_param = method_param
+        self.algo = algo
+        self.name = 'redisearch-%s (%s)' % (self.algo, self.method_param)
+        self.index_name = "ann_benchmark"
+        self.text = None
+        
+        redis = RedisCluster if conn_params['cluster'] else Redis
+        host = conn_params["host"] if conn_params["host"] else 'localhost'
+        port = conn_params["port"] if conn_params["port"] else 6379
+        self.redis = redis(host=host, port=port, decode_responses=False,
+                           password=conn_params["auth"], username=conn_params["user"])
+        self.shards = int(conn_params["shards"])
+        if conn_params['cluster']:
+            self.shards = len(self.redis.get_primaries())
+
+    def fit(self, X, offset=0, limit=None, hybrid_buckets = None):
+        limit = limit if limit else len(X)
+        try:
+            args = [self.index_name, 'SCHEMA']
+            if hybrid_buckets:
+                args.extend(['n', 'NUMERIC', 't', 'TEXT'])
+            # https://oss.redis.com/redisearch/master/Commands/#ftcreate
+            if self.algo == "HNSW":
+                args.extend(['vector', 'VECTOR', self.algo, '10', 'TYPE', 'FLOAT32', 'DIM', len(X[0]), 'DISTANCE_METRIC', self.metric, 'M', self.method_param['M'], 'EF_CONSTRUCTION', self.method_param["efConstruction"]])
+            elif self.algo == "FLAT":
+                args.extend(['vector', 'VECTOR', self.algo, '6', 'TYPE', 'FLOAT32', 'DIM', len(X[0]), 'DISTANCE_METRIC', self.metric])
+            print("Calling FT.CREATE", *args)
+            self.redis.execute_command('FT.CREATE', *args,  target_nodes='random')
+        except Exception as e:
+            if 'Index already exists' not in str(e):
+                raise
+        p = self.redis.pipeline(transaction=False)
+        count = 0
+        if hybrid_buckets:
+            print('running hybrid')
+            for bucket in hybrid_buckets.values():
+                ids = bucket['ids']
+                text = bucket['text'].decode()
+                number = bucket['number']
+                print('calling HSET', f'<id>', 'vector', '<vector blob>', 't', text, 'n', number)
+                for id in ids:
+                    if id >= offset and id < limit:
+                        p.execute_command('HSET', int(id), 'vector', X[id].tobytes(), 't', text, 'n', int(number))
+                        count+=1
+                        if count % 1000 == 0:
+                            p.execute()
+                            p.reset()
+            p.execute()
+        else:
+            for i in range(offset, limit):
+                p.execute_command('HSET', i, 'vector', X[i].tobytes())
+                count+=1
+                if count % 1000 == 0:
+                    p.execute()
+                    p.reset()
+            p.execute()
+
+    def set_query_arguments(self, ef):
+        self.ef = ef
+
+    def set_hybrid_query(self, text):
+        self.text = text
+
+    def query(self, v, k):
+        # https://oss.redis.com/redisearch/master/Commands/#ftsearch
+        qparams = f' EF_RUNTIME {self.ef}' if self.algo == 'HNSW' else ''
+        if self.text:
+            vq = f'(@t:{self.text})=>[KNN {k} @vector $BLOB {qparams}]'
+        else:
+            vq = f'*=>[KNN {k} @vector $BLOB {qparams}]'
+        q = ['FT.SEARCH', self.index_name, vq, 'NOCONTENT', 'SORTBY', '__vector_score', 'LIMIT', '0', str(k), 'PARAMS', '2', 'BLOB', v.tobytes(), 'DIALECT', '2']
+        return [int(doc) for doc in self.redis.execute_command(*q, target_nodes='random')[1:]]
+
+    def freeIndex(self):
+        self.redis.execute_command("FLUSHALL")
+
+    def __str__(self):
+        return self.name + f", efRuntime: {self.ef}"

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

@@ -0,0 +1,20 @@
+from __future__ import absolute_import
+import rpforest
+import numpy
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class RPForest(BaseANN):
+    def __init__(self, leaf_size, n_trees):
+        self.name = 'RPForest(leaf_size=%d, n_trees=%d)' % (leaf_size, n_trees)
+        self._model = rpforest.RPForest(leaf_size=leaf_size, no_trees=n_trees)
+
+    def fit(self, X):
+        if X.dtype != numpy.double:
+            X = numpy.array(X).astype(numpy.double)
+        self._model.fit(X)
+
+    def query(self, v, n):
+        if v.dtype != numpy.double:
+            v = numpy.array(v).astype(numpy.double)
+        return self._model.query(v, n)

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

@@ -0,0 +1,34 @@
+from __future__ import absolute_import
+import os
+import numpy as np
+import scann
+from ann_benchmarks.algorithms.base import BaseANN
+
+class Scann(BaseANN):
+
+  def __init__(self, n_leaves, avq_threshold, dims_per_block, dist):
+    self.name = "scann n_leaves={} avq_threshold={:.02f} dims_per_block={}".format(
+            n_leaves, avq_threshold, dims_per_block)
+    self.n_leaves = n_leaves
+    self.avq_threshold = avq_threshold
+    self.dims_per_block = dims_per_block
+    self.dist = dist
+
+  def fit(self, X):
+    if self.dist == "dot_product":
+      spherical = True
+      X[np.linalg.norm(X, axis=1) == 0] = 1.0 / np.sqrt(X.shape[1])
+      X /= np.linalg.norm(X, axis=1)[:, np.newaxis]
+    else:
+      spherical = False
+
+    self.searcher = scann.scann_ops_pybind.builder(X, 10, self.dist).tree(
+        self.n_leaves, 1, training_sample_size=len(X), spherical=spherical, quantize_centroids=True).score_ah(
+            self.dims_per_block, anisotropic_quantization_threshold=self.avq_threshold).reorder(
+                1).build()
+
+  def set_query_arguments(self, leaves_reorder):
+      self.leaves_to_search, self.reorder = leaves_reorder
+
+  def query(self, v, n):
+    return self.searcher.search(v, n, self.reorder, self.leaves_to_search)[0]

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

@@ -0,0 +1,28 @@
+from __future__ import absolute_import
+import SPTAG
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class Sptag(BaseANN):
+    def __init__(self, metric, algo):
+        self._algo = str(algo)
+        self._metric = {
+            'angular': 'Cosine', 'euclidean': 'L2'}[metric]
+
+    def fit(self, X):
+        self._sptag = SPTAG.AnnIndex(self._algo, 'Float', X.shape[1])
+        self._sptag.SetBuildParam("NumberOfThreads", '32')
+        self._sptag.SetBuildParam("DistCalcMethod", self._metric)
+        self._sptag.Build(X, X.shape[0])
+
+    def set_query_arguments(self, MaxCheck):
+        self._maxCheck = MaxCheck
+        self._sptag.SetSearchParam("MaxCheck", str(self._maxCheck))
+
+    def query(self, v, k):
+        return self._sptag.Search(v, k)[0]
+
+    def __str__(self):
+        return 'Sptag(metric=%s, algo=%s, check=%d)' % (self._metric,
+                                              self._algo, self._maxCheck)
+