Merge pull request #29 from qdrant/add-filtering-support

Add filtering support

Author: generall

.gitignore

@@ -1,4 +1,5 @@
 .idea/
+.pytest_cache/
 __pycache__
 *.pyc
 NOTES.md

benchmark/dataset.py

@@ -2,10 +2,11 @@
 import shutil
 import tarfile
 import urllib.request
-from dataclasses import dataclass
-from typing import Optional
+from dataclasses import dataclass, field
+from typing import Dict, Optional
 
 from benchmark import DATASETS_DIR
+from dataset_reader.ann_compound_reader import AnnCompoundReader
 from dataset_reader.ann_h5_reader import AnnH5Reader
 from dataset_reader.base_reader import BaseReader
 from dataset_reader.json_reader import JSONReader
@@ -19,9 +20,10 @@ class DatasetConfig:
     type: str
     path: str
     link: Optional[str] = None
+    schema: Optional[Dict[str, str]] = field(default_factory=dict)
 
 
-READER_TYPE = {"h5": AnnH5Reader, "jsonl": JSONReader}
+READER_TYPE = {"h5": AnnH5Reader, "jsonl": JSONReader, "tar": AnnCompoundReader}
 
 
 class Dataset:
@@ -39,9 +41,11 @@ def download(self):
             print(f"Downloading {self.config.link}...")
             tmp_path, _ = urllib.request.urlretrieve(self.config.link)
 
-            if tmp_path.endswith(".tgz") or tmp_path.endswith(".tar.gz"):
+            if self.config.link.endswith(".tgz") or self.config.link.endswith(
+                ".tar.gz"
+            ):
                 print(f"Extracting: {tmp_path} -> {target_path}")
-                (DATASETS_DIR / self.config.path).mkdir(exist_ok=True)
+                (DATASETS_DIR / self.config.path).mkdir(exist_ok=True, parents=True)
                 file = tarfile.open(tmp_path)
                 file.extractall(target_path)
                 file.close()

dataset_reader/ann_compound_reader.py

@@ -0,0 +1,39 @@
+import json
+from typing import Iterator, List
+
+import numpy as np
+
+from dataset_reader.base_reader import Query
+from dataset_reader.json_reader import JSONReader
+
+
+class AnnCompoundReader(JSONReader):
+    """
+    A reader created specifically to read the format used in
+    https://github.com/qdrant/ann-filtering-benchmark-datasets, in which vectors
+    and their metadata are stored in separate files.
+    """
+
+    VECTORS_FILE = "vectors.npy"
+    QUERIES_FILE = "tests.jsonl"
+
+    def read_vectors(self) -> Iterator[List[float]]:
+        vectors = np.load(self.path / self.VECTORS_FILE)
+        for vector in vectors:
+            if self.normalize:
+                vector = vector / np.linalg.norm(vector)
+            yield vector.tolist()
+
+    def read_queries(self) -> Iterator[Query]:
+        with open(self.path / self.QUERIES_FILE) as payloads_fp:
+            for idx, row in enumerate(payloads_fp):
+                row_json = json.loads(row)
+                vector = np.array(row_json["query"])
+                if self.normalize:
+                    vector /= np.linalg.norm(vector)
+                yield Query(
+                    vector=vector.tolist(),
+                    meta_conditions=row_json["conditions"],
+                    expected_result=row_json["closest_ids"],
+                    expected_scores=row_json["closest_scores"],
+                )

dataset_reader/json_reader.py

@@ -6,46 +6,46 @@
 
 from dataset_reader.base_reader import BaseReader, Query, Record
 
-VECTORS_FILE = "vectors.jsonl"
-PAYLOADS_FILE = "payloads.jsonl"
-QUERIES_FILE = "queries.jsonl"
-NEIGHBOURS_FILE = "neighbours.jsonl"
-
 
 class JSONReader(BaseReader):
+    VECTORS_FILE = "vectors.jsonl"
+    PAYLOADS_FILE = "payloads.jsonl"
+    QUERIES_FILE = "queries.jsonl"
+    NEIGHBOURS_FILE = "neighbours.jsonl"
+
     def __init__(self, path: Path, normalize=False):
         self.path = path
         self.normalize = normalize
 
     def read_payloads(self) -> Iterator[dict]:
-        if not (self.path / PAYLOADS_FILE).exists():
+        if not (self.path / self.PAYLOADS_FILE).exists():
             while True:
                 yield {}
-        with open(self.path / PAYLOADS_FILE, "r") as json_fp:
+        with open(self.path / self.PAYLOADS_FILE, "r") as json_fp:
             for json_line in json_fp:
                 line = json.loads(json_line)
                 yield line
 
     def read_vectors(self) -> Iterator[List[float]]:
-        with open(self.path / VECTORS_FILE, "r") as json_fp:
+        with open(self.path / self.VECTORS_FILE, "r") as json_fp:
             for json_line in json_fp:
                 vector = json.loads(json_line)
                 if self.normalize:
                     vector = vector / np.linalg.norm(vector)
                 yield vector
 
     def read_neighbours(self) -> Iterator[Optional[List[int]]]:
-        if not (self.path / NEIGHBOURS_FILE).exists():
+        if not (self.path / self.NEIGHBOURS_FILE).exists():
             while True:
                 yield None
 
-        with open(self.path / NEIGHBOURS_FILE, "r") as json_fp:
+        with open(self.path / self.NEIGHBOURS_FILE, "r") as json_fp:
             for json_line in json_fp:
                 line = json.loads(json_line)
                 yield line
 
     def read_query_vectors(self) -> Iterator[List[float]]:
-        with open(self.path / QUERIES_FILE, "r") as json_fp:
+        with open(self.path / self.QUERIES_FILE, "r") as json_fp:
             for json_line in json_fp:
                 vector = json.loads(json_line)
                 if self.normalize:

datasets/datasets.json

@@ -39,6 +39,149 @@
     "path": "gist-960-angular/gist-960-angular.hdf5",
     "link": "http://ann-benchmarks.com/gist-960-euclidean.hdf5"
   },
+  {
+    "name": "h-and-m-2048-angular-filters",
+    "vector_size": 2048,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "h-and-m-2048-angular/hnm",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/hnm.tgz",
+    "schema": {
+      "product_code": "int",
+      "prod_name": "keyword",
+      "product_type_no": "int",
+      "product_type_name": "keyword",
+      "product_group_name": "keyword",
+      "graphical_appearance_no": "int",
+      "graphical_appearance_name": "keyword",
+      "colour_group_code": "int",
+      "colour_group_name": "keyword",
+      "perceived_colour_value_id": "int",
+      "perceived_colour_value_name": "keyword",
+      "perceived_colour_master_id": "int",
+      "perceived_colour_master_name": "keyword",
+      "department_no": "int",
+      "department_name": "keyword",
+      "index_code": "keyword",
+      "index_name": "keyword",
+      "index_group_no": "int",
+      "index_group_name": "keyword",
+      "section_no": "int",
+      "section_name": "keyword",
+      "garment_group_no": "int",
+      "garment_group_name": "keyword",
+      "detail_desc": "text"
+    }
+  },
+  {
+    "name": "arxiv-titles-384-angular-filters",
+    "vector_size": 384,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "arxiv-titles-384-angular/arxiv",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/arxiv.tar.gz",
+    "schema": {
+      "update_date_ts": "int",
+      "labels": "keyword",
+      "submitter": "keyword"
+    }
+  },
+  {
+    "name": "random-match-keyword-100-angular-filters",
+    "vector_size": 100,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-match-keyword-100-angular/random_keywords_1m",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_keywords_1m.tgz",
+    "schema": {
+      "a": "keyword",
+      "b": "keyword"
+    }
+  },
+  {
+    "name": "random-match-int-100-angular-filters",
+    "vector_size": 100,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-match-int-100-angular/random_ints_1m",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_ints_1m.tgz",
+    "schema": {
+      "a": "int",
+      "b": "int"
+    }
+  },
+  {
+    "name": "random-range-100-angular-filters",
+    "vector_size": 100,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-range-100-angular/random_float_1m",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_float_1m.tgz",
+    "schema": {
+      "a": "float",
+      "b": "float"
+    }
+  },
+  {
+    "name": "random-geo-radius-100-angular-filters",
+    "vector_size": 100,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-geo-radius-100-angular/random_geo_1m",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_geo_1m.tgz",
+    "schema": {
+      "a": "geo",
+      "b": "geo"
+    }
+  },
+  {
+    "name": "random-match-keyword-2048-angular-filters",
+    "vector_size": 2048,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-match-keyword-2048-angular/random_keywords_100k",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_keywords_100k.tgz",
+    "schema": {
+      "a": "keyword",
+      "b": "keyword"
+    }
+  },
+  {
+    "name": "random-match-int-2048-angular-filters",
+    "vector_size": 2048,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-match-int-2048-angular/random_ints_100k",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_ints_100k.tgz",
+    "schema": {
+      "a": "int",
+      "b": "int"
+    }
+  },
+  {
+    "name": "random-range-2048-angular-filters",
+    "vector_size": 2048,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-range-2048-angular/random_float_100k",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_float_100k.tgz",
+    "schema": {
+      "a": "float",
+      "b": "float"
+    }
+  },
+  {
+    "name": "random-geo-radius-2048-angular-filters",
+    "vector_size": 2048,
+    "distance": "cosine",
+    "type": "tar",
+    "path": "random-geo-radius-2048-angular/random_geo_100k",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_geo_100k.tgz",
+    "schema": {
+      "a": "geo",
+      "b": "geo"
+    }
+  },
   {
     "name": "random-100",
     "vector_size": 100,

engine/base_client/client.py

@@ -33,10 +33,12 @@ def save_search_results(
         experiments_file = (
             f"{self.name}-{dataset_name}-search-{search_id}-{timestamp}.json"
         )
-        with open(RESULTS_DIR / experiments_file, "w") as out:
+        result_path = RESULTS_DIR / experiments_file
+        with open(result_path, "w") as out:
             out.write(
                 json.dumps({"params": search_params, "results": results}, indent=2)
             )
+        return result_path
 
     def save_upload_results(
         self, dataset_name: str, results: dict, upload_params: dict
@@ -60,10 +62,7 @@ def run_experiment(self, dataset: Dataset, skip_upload: bool = False):
 
         if not skip_upload:
             print("Experiment stage: Configure")
-            self.configurator.configure(
-                distance=dataset.config.distance,
-                vector_size=dataset.config.vector_size,
-            )
+            self.configurator.configure(dataset)
 
             print("Experiment stage: Upload")
             upload_stats = self.uploader.upload(
@@ -88,3 +87,4 @@ def run_experiment(self, dataset: Dataset, skip_upload: bool = False):
                 dataset.config.name, search_stats, search_id, search_params
             )
         print("Experiment stage: Done")
+        print("Results saved to: ", RESULTS_DIR)

engine/base_client/configure.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from benchmark.dataset import Dataset
+
 
 class BaseConfigurator:
     DISTANCE_MAPPING = {}
@@ -12,12 +14,12 @@ def __init__(self, host, collection_params: dict, connection_params: dict):
     def clean(self):
         raise NotImplementedError()
 
-    def recreate(self, distance, vector_size, collection_params):
+    def recreate(self, dataset: Dataset, collection_params):
         raise NotImplementedError()
 
-    def configure(self, distance, vector_size) -> Optional[dict]:
+    def configure(self, dataset: Dataset) -> Optional[dict]:
         self.clean()
-        return self.recreate(distance, vector_size, self.collection_params) or {}
+        return self.recreate(dataset, self.collection_params) or {}
 
     def execution_params(self, distance, vector_size) -> dict:
         return {}