Basic Weaviate implementation

algos.yaml

@@ -595,6 +595,15 @@ float:
           arg-groups: [[4, 8, 12, 16, 24, 36, 48, 64, 96]]
           query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
 
+    weaviate:
+      docker-tag: ann-benchmarks-weaviate
+      module: ann_benchmarks.algorithms.weaviate
+      constructor: Weaviate
+      base-args: ["@metric"]
+      run-groups:
+        weaviate:
+          args: [[20, 40, 100, 200, 400, 1000]]
+
   euclidean:
     vamana(diskann):
       docker-tag: ann-benchmarks-diskann

ann_benchmarks/algorithms/weaviate.py

@@ -0,0 +1,66 @@
+import subprocess
+import sys
+
+import weaviate
+from weaviate.embedded import EmbeddedOptions
+from weaviate.util import generate_uuid5
+
+from .base import BaseANN
+
+
+class Weaviate(BaseANN):
+    def __init__(self, metric, ef_construction):
+        self.client = weaviate.Client(
+            embedded_options=EmbeddedOptions()
+        )
+        self.ef_construction = ef_construction
+        self.distance = {
+            "angular": "cosine",
+            "euclidean": "l2-squared",
+        }[metric]
+
+    def fit(self, X):
+        self.client.schema.create({
+            "classes": [
+                {
+                    "class": "Vector",
+                    "properties": [
+                        {
+                            "name": "i",
+                            "dataType": ["int"],
+                        }
+                    ],
+                    "vectorIndexConfig": {
+                        "distance": self.distance,
+                        "efConstruction": self.ef_construction,
+                    },
+                }
+            ]
+        })
+        with self.client.batch as batch:
+            batch.batch_size = 100
+            for i, x in enumerate(X):
+                properties = { "i": i }
+                uuid = generate_uuid5(properties, "Vector")
+                self.client.batch.add_data_object(
+                    data_object=properties,
+                    class_name="Vector",
+                    uuid=uuid,
+                    vector=x
+                )
+
+    def query(self, v, n):
+        ret = (
+            self.client.query
+            .get("Vector", ["i"])
+            .with_near_vector({
+                "vector": v,
+            })
+            .with_limit(n)
+            .do()
+        )
+        # {'data': {'Get': {'Vector': [{'i': 3618}, {'i': 8213}, {'i': 4462}, {'i': 6709}, {'i': 3975}, {'i': 3129}, {'i': 5120}, {'i': 2979}, {'i': 6319}, {'i': 3244}]}}}
+        return [d["i"] for d in ret["data"]["Get"]["Vector"]]
+
+    def __str__(self):
+        return "Weaviate()"

install/Dockerfile.weaviate

@@ -1,14 +1,2 @@
-# Build Weviate in the official golang image (which uses Debian)
-# See https://github.com/weaviate/weaviate/blob/master/Dockerfile
-FROM golang as weaviate-go
-RUN apt install -y git
-RUN git clone --depth=1 https://github.com/weaviate/weaviate
-WORKDIR weaviate
-RUN GOOS=linux \
-      go build -ldflags '-w -extldflags "-static"' \
-      -o /weaviate-server ./cmd/weaviate-server
-
-# Copy it into a new image & install the client
 FROM ann-benchmarks
-COPY --from=weaviate-go /weaviate-server /bin/weaviate
-RUN pip install weaviate-client
+RUN pip install weaviate-client