Merge remote-tracking branch 'origin/upstreamchanges2502'

Author: randersenYB

.github/workflows/benchmarks.yml

@@ -55,6 +55,7 @@ jobs:
           - onng_ngt
           - opensearchknn
           - panng_ngt
+          - parlayann
           - pg_embedding
           - pgvector
           - pgvecto_rs
@@ -70,6 +71,7 @@ jobs:
           - vearch
           - vespa
           - voyager
+          - vsag
           - weaviate
         include:
           - library: pynndescent

README.md

@@ -71,6 +71,8 @@ We have a number of precomputed data sets in HDF5 format. All data sets have bee
 | [NYTimes](https://archive.ics.uci.edu/ml/datasets/bag+of+words)   |        256 |    290,000 |    10,000 |       100 | Angular   | [HDF5](http://ann-benchmarks.com/nytimes-256-angular.hdf5) (301MB)         |
 | [SIFT](http://corpus-texmex.irisa.fr/)                           |        128 |  1,000,000 |    10,000 |       100 | Euclidean | [HDF5](http://ann-benchmarks.com/sift-128-euclidean.hdf5) (501MB)          |
 | [Last.fm](https://github.com/erikbern/ann-benchmarks/pull/91)     |         65 |    292,385 |    50,000 |       100 | Angular   | [HDF5](http://ann-benchmarks.com/lastfm-64-dot.hdf5) (135MB)               |
+| [COCO-I2I](https://cocodataset.org/)                              |        512 |    113,287 |    10,000 |       100 | Angular   | [HDF5](https://github.com/fabiocarrara/str-encoders/releases/download/v0.1.3/coco-i2i-512-angular.hdf5) (136MB) |
+| [COCO-T2I](https://cocodataset.org/)                              |        512 |    113,287 |    10,000 |       100 | Angular   | [HDF5](https://github.com/fabiocarrara/str-encoders/releases/download/v0.1.3/coco-t2i-512-angular.hdf5) (136MB) |
 
 Results
 =======

ann_benchmarks/algorithms/elastiknn/Dockerfile

@@ -13,13 +13,13 @@ WORKDIR /home/elasticsearch
 USER elasticsearch
 
 # Install elasticsearch.
-RUN curl -o elasticsearch.tar.gz https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-8.12.2-linux-x86_64.tar.gz
+RUN curl -o elasticsearch.tar.gz https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-8.15.0-linux-x86_64.tar.gz
 RUN tar xzf elasticsearch.tar.gz
 RUN mv elasticsearch-* elasticsearch && rm elasticsearch.tar.gz
 
 # Install plugin.
 RUN /home/elasticsearch/elasticsearch/bin/elasticsearch-plugin install --batch \
-    https://github.com/alexklibisz/elastiknn/releases/download/8.12.2.1/elastiknn-8.12.2.1.zip
+    https://github.com/alexklibisz/elastiknn/releases/download/8.15.0.1/elastiknn-8.15.0.1.zip
 
 # Configuration
 # Backup the original configurations, which can be useful for comparing.
@@ -66,10 +66,7 @@ USER root
 WORKDIR /home/app
 
 # Install python client.
-# Using no-deps because scipy (1.7.0) is incompatible with the container version of Python (3.6).
-# Then we need to install the deps manually.
-RUN python3 -m pip install --no-deps elastiknn-client==8.6.2.0
-RUN python3 -m pip install elasticsearch==8.6.2 dataclasses-json==0.3.7 tqdm==4.61.1
+RUN python3 -m pip install elastiknn-client==8.15.0.1
 
 # Custom entrypoint that also starts the Elasticsearch server in the background
 RUN echo '\

ann_benchmarks/algorithms/elastiknn/config.yml

@@ -18,5 +18,5 @@ float:
     name: elastiknn-l2lsh
     run_groups:
       elastiknn-l2lsh:
-        args: [[100], [4], [1024, 2048]]
-        query_args: [[500, 1000], [0, 3]]
+        args: [[175], [7], [3900]]
+        query_args: [[100,500,1000], [0]]

ann_benchmarks/algorithms/hnswlib/module.py

@@ -9,8 +9,6 @@ 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
@@ -24,6 +22,7 @@ def fit(self, X):
 
     def set_query_arguments(self, ef):
         self.p.set_ef(ef)
+        self.name = "hnswlib (%s, 'efQuery': %s)" % (self.method_param, ef)
 
     def query(self, v, n):
         # print(np.expand_dims(v,axis=0).shape)

ann_benchmarks/algorithms/parlayann/Dockerfile

@@ -0,0 +1,20 @@
+FROM ann-benchmarks
+
+RUN apt update
+RUN apt install -y software-properties-common
+RUN add-apt-repository -y ppa:git-core/ppa
+RUN apt update
+RUN DEBIAN_FRONTEND=noninteractive apt install -y git make cmake g++ libaio-dev libgoogle-perftools-dev libunwind-dev clang-format libboost-dev libboost-program-options-dev libmkl-full-dev libcpprest-dev python3.10
+
+#RUN apt-get update 
+#RUN apt-get install -y g++ software-properties-common
+RUN pip3 install pybind11 numpy
+
+ARG CACHEBUST=1
+RUN git clone -b annbench https://github.com/cmuparlay/ParlayANN.git 
+RUN cd ParlayANN && git submodule update --init --recursive
+RUN cd ParlayANN/python && bash compile.sh
+#RUN cd ParlayANN/python && pip install -e .
+#RUN python3 -c 'import parlaypy'
+ENV PYTHONPATH=$PYTHONPATH:/home/app/ParlayANN/python
+WORKDIR /home/app

ann_benchmarks/algorithms/parlayann/config.yml

@@ -0,0 +1,43 @@
+float:
+  euclidean:
+  - base_args: ['@metric']
+    constructor: ParlayANN
+    disabled: false
+    docker_tag: ann-benchmarks-parlayann
+    module: ann_benchmarks.algorithms.parlayann
+    name: parlayann
+    run_groups:
+      parlay_80:
+        args: [{alpha: 1.15, R: 80, L: 160, two_pass: True}]
+        query_args: [[{Q: 20}, {Q: 22}, {Q: 25}, {Q: 30}, {Q: 40}, {Q: 50}, {Q: 60}, {Q: 80}, {Q: 100}, {Q: 125}, {Q: 150}, {Q: 200}, {Q: 300}, {Q: 400}, {Q: 600}, {Q: 800}]]
+      parlay_64:
+        args: [{alpha: 1.1, R: 64, L: 128, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 14}, {Q: 16}, {Q: 18}, {Q: 20}, {Q: 22}, {Q: 25}, {Q: 30}, {Q: 40}, {Q: 50}, {Q: 60}, {Q: 80}, {Q: 100}, {Q: 125}, {Q: 150}, {Q: 200}, {Q: 300}, {Q: 400}, {Q: 600}, {Q: 800}]]
+      parlay_40:
+        args: [{alpha: 1.08, R: 40, L: 80, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 14}, {Q: 16}, {Q: 18}, {Q: 20}]]
+      parlay_32_05:  
+        args: [{alpha: 1.05, R: 32, L: 64, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 14}, {Q: 16}, {Q: 18}, {Q: 20}]]
+  angular:
+  - base_args: ['@metric']
+    constructor: ParlayANN
+    disabled: false
+    docker_tag: ann-benchmarks-parlayann
+    module: ann_benchmarks.algorithms.parlayann
+    name: parlayann
+    run_groups:
+      parlay_130:
+        args: [{alpha: .85, R: 130, L: 260, two_pass: True}]
+        query_args: [[{Q: 20}, {Q: 22}, {Q: 25}, {Q: 30}, {Q: 40}, {Q: 50}, {Q: 60}, {Q: 80}, {Q: 100}, {Q: 125}, {Q: 150}, {Q: 200}, {Q: 400}, {Q: 800}]]
+      parlay_100:
+        args: [{alpha: .85, R: 100, L: 200, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {limit: 25}, {limit: 30}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 13}, {Q: 14}, {Q: 15}, {Q: 16}, {Q: 17}, {Q: 18}, {Q: 20}, {Q: 22}, {Q: 25}, {Q: 30}, {Q: 40}, {Q: 50}, {Q: 60}, {Q: 80}, {Q: 100}, {Q: 125}, {Q: 150}, {Q: 200}, {Q: 400}, {Q: 800}]]
+      parlay_80:
+        args: [{alpha: .90, R: 80, L: 160, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {limit: 25}, {limit: 30}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 13}, {Q: 14}, {Q: 15}, {Q: 16}, {Q: 17}, {Q: 18}, {Q: 20}]]
+      parlay_50:
+        args: [{alpha: .95, R: 50, L: 100, two_pass: True}]
+        query_args: [[{limit: 10}, {limit: 11}, {limit: 12}, {limit: 13}, {limit: 14}, {limit: 15}, {limit: 16}, {limit: 18}, {limit: 20}, {limit: 22}, {limit: 25}, {limit: 30}, {Q: 10}, {Q: 11}, {Q: 12}, {Q: 14}, {Q: 16}, {Q: 18}, {Q: 20}]]
+
+

ann_benchmarks/algorithms/parlayann/module.py

@@ -0,0 +1,84 @@
+from __future__ import absolute_import
+import psutil
+import os
+import struct
+import time
+import numpy as np
+import wrapper as pann
+
+from ..base.module import BaseANN
+
+class ParlayANN(BaseANN):
+    def __init__(self, metric, index_params):
+        self.name = "parlayann_(" + str(index_params) + ")"
+        self._index_params = index_params
+        self._metric = self.translate_dist_fn(metric)
+
+        self.R = int(index_params.get("R", 50))
+        self.L = int(index_params.get("L", 100))
+        self.alpha = float(index_params.get("alpha", 1.15))
+        self.two_pass = bool(index_params.get("two_pass", False))
+
+    def translate_dist_fn(self, metric):
+        if metric == 'euclidean':
+            return 'Euclidian'
+        elif metric == 'ip':
+            return 'mips'
+        elif metric == 'angular':
+            return 'mips'
+        else:
+            raise Exception('Invalid metric')
+        
+    def translate_dtype(self, dtype:str):
+        if dtype == 'float32':
+            return 'float'
+        else:
+            return dtype
+
+    def fit(self, X):
+        def bin_to_float(binary):
+            return struct.unpack("!f", struct.pack("!I", int(binary, 2)))[0]
+
+        print("Vamana: Starting Fit...")
+        index_dir = "indices"
+
+        if not os.path.exists(index_dir):
+            os.makedirs(index_dir)
+
+        data_path = os.path.join(index_dir, "base.bin")
+        save_path = os.path.join(index_dir, self.name)
+        print("parlayann: Index Stored At: " + save_path)
+        nb, dims = X.shape
+        shape = [
+            np.float32(bin_to_float("{:032b}".format(nb))),
+            np.float32(bin_to_float("{:032b}".format(dims))),
+        ]
+        X = X.flatten()
+        X = np.insert(X, 0, shape)
+        X.tofile(data_path)
+
+        if not os.path.exists(save_path):
+            print("parlayann: Creating Index")
+            start = time.time()
+            self.params = pann.build_vamana_index(self._metric, "float", data_path, save_path,
+                                                  self.R, self.L, self.alpha, self.two_pass)
+            end = time.time()
+            print("Indexing time: ", end - start)
+            print(f"Wrote index to {save_path}")
+        self.index = pann.load_index(self._metric, "float", data_path, save_path)
+        print("Index loaded")
+
+    def query(self, X, k):
+        return self.index.single_search(X, k, self.Q, True, self.limit)
+
+    def batch_query(self, X, k):
+        print("running batch")
+        nq, dims = X.shape
+        self.res, self.distances = self.index.batch_search(X, k, self.Q, True, self.limit)
+        return self.res
+
+    def set_query_arguments(self, query_args):
+        self.name = "parlayann_(" + str(self._index_params) + "," + str(query_args) + ")"
+        print(query_args)
+        self.limit = 1000 if query_args.get("limit") is None else query_args.get("limit")
+        self.Q = 10 if query_args.get("Q") is None else query_args.get("Q")

ann_benchmarks/algorithms/pgvector/module.py

@@ -1,10 +1,52 @@
+"""
+This module supports connecting to a PostgreSQL instance and performing vector
+indexing and search using the pgvector extension. The default behavior uses
+the "ann" value of PostgreSQL user name, password, and database name, as well
+as the default host and port values of the psycopg driver.
+
+If PostgreSQL is managed externally, e.g. in a cloud DBaaS environment, the
+environment variable overrides listed below are available for setting PostgreSQL
+connection parameters:
+
+ANN_BENCHMARKS_PG_USER
+ANN_BENCHMARKS_PG_PASSWORD
+ANN_BENCHMARKS_PG_DBNAME
+ANN_BENCHMARKS_PG_HOST
+ANN_BENCHMARKS_PG_PORT
+
+This module starts the PostgreSQL service automatically using the "service"
+command. The environment variable ANN_BENCHMARKS_PG_START_SERVICE could be set
+to "false" (or e.g. "0" or "no") in order to disable this behavior.
+
+This module will also attempt to create the pgvector extension inside the
+target database, if it has not been already created.
+"""
+
 import subprocess
 import sys
+import os
 
 import pgvector.psycopg
 import psycopg
 
+from typing import Dict, Any, Optional
+
 from ..base.module import BaseANN
+from ...util import get_bool_env_var
+
+
+def get_pg_param_env_var_name(pg_param_name: str) -> str:
+    return f'ANN_BENCHMARKS_PG_{pg_param_name.upper()}'
+
+
+def get_pg_conn_param(
+        pg_param_name: str,
+        default_value: Optional[str] = None) -> Optional[str]:
+    env_var_name = get_pg_param_env_var_name(pg_param_name)
+    env_var_value = os.getenv(env_var_name, default_value)
+    if env_var_value is None or len(env_var_value.strip()) == 0:
+        return default_value
+    return env_var_value
 
 
 class PGVector(BaseANN):
@@ -21,9 +63,61 @@ def __init__(self, metric, method_param):
         else:
             raise RuntimeError(f"unknown metric {metric}")
 
+    def ensure_pgvector_extension_created(self, conn: psycopg.Connection) -> None:
+        """
+        Ensure that `CREATE EXTENSION vector` has been executed.
+        """
+        with conn.cursor() as cur:
+            # We have to use a separate cursor for this operation.
+            # If we reuse the same cursor for later operations, we might get
+            # the following error:
+            # KeyError: "couldn't find the type 'vector' in the types registry"
+            cur.execute(
+                "SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector')")
+            pgvector_exists = cur.fetchone()[0]
+            if pgvector_exists:
+                print("vector extension already exists")
+            else:
+                print("vector extension does not exist, creating")
+                cur.execute("CREATE EXTENSION vector")
+
     def fit(self, X):
-        subprocess.run("service postgresql start", shell=True, check=True, stdout=sys.stdout, stderr=sys.stderr)
-        conn = psycopg.connect(user="ann", password="ann", dbname="ann", autocommit=True)
+        psycopg_connect_kwargs: Dict[str, Any] = dict(
+            autocommit=True,
+        )
+        for arg_name in ['user', 'password', 'dbname']:
+            # The default value is "ann" for all of these parameters.
+            psycopg_connect_kwargs[arg_name] = get_pg_conn_param(
+                arg_name, 'ann')
+
+        # If host/port are not specified, leave the default choice to the
+        # psycopg driver.
+        pg_host: Optional[str] = get_pg_conn_param('host')
+        if pg_host is not None:
+            psycopg_connect_kwargs['host'] = pg_host
+
+        pg_port_str: Optional[str] = get_pg_conn_param('port')
+        if pg_port_str is not None:
+            psycopg_connect_kwargs['port'] = int(pg_port_str)
+
+        should_start_service = get_bool_env_var(
+            get_pg_param_env_var_name('start_service'),
+            default_value=True)
+        if should_start_service:
+            subprocess.run(
+                "service postgresql start",
+                shell=True,
+                check=True,
+                stdout=sys.stdout,
+                stderr=sys.stderr)
+        else:
+            print(
+                "Assuming that PostgreSQL service is managed externally. "
+                "Not attempting to start the service.")
+
+        conn = psycopg.connect(**psycopg_connect_kwargs)
+        self.ensure_pgvector_extension_created(conn)
+
         pgvector.psycopg.register_vector(conn)
         cur = conn.cursor()
         cur.execute("DROP TABLE IF EXISTS items")
@@ -46,6 +140,7 @@ def fit(self, X):
         print("done!")
         self._cur = cur
 
+
     def set_query_arguments(self, ef_search):
         self._ef_search = ef_search
         self._cur.execute("SET hnsw.ef_search = %d" % ef_search)

ann_benchmarks/algorithms/qsg_ngt/Dockerfile

@@ -4,8 +4,9 @@ RUN apt update
 RUN apt install -y git cmake g++ python3 python3-setuptools python3-pip libblas-dev liblapack-dev
 RUN pip3 install wheel pybind11==2.5.0
 RUN git clone https://github.com/WPJiang/HWTL_SDU-ANNS.git
-RUN cp HWTL_SDU-ANNS/lib/* /usr/local/lib/
-RUN cp HWTL_SDU-ANNS/bin/* /usr/local/bin/
+RUN cp HWTL_SDU-ANNS/create.py /home/app/
+RUN cp -r HWTL_SDU-ANNS/lib/* /usr/local/lib/
+RUN cp -r HWTL_SDU-ANNS/bin/* /usr/local/bin/
 RUN chmod a+x /usr/local/bin/* && chmod a+x HWTL_SDU-ANNS/*
 RUN ldconfig
 RUN pip3 install HWTL_SDU-ANNS/qsgngt-*-linux_x86_64.whl