Add benchmark support for radial search

Signed-off-by: Junqiu Lei <[email protected]>
opensearch-project · Jun 4, 2024 · 3aa080e · 3aa080e
1 parent f194726
commit 3aa080e
Show file tree

Hide file tree

Showing 4 changed files with 443 additions and 122 deletions.
diff --git a/osbenchmark/utils/dataset.py b/osbenchmark/utils/dataset.py
@@ -24,6 +24,8 @@ class Context(Enum):
     INDEX = 1
     QUERY = 2
     NEIGHBORS = 3
+    MAX_DISTANCE_NEIGHBORS = 4
+    MIN_SCORE_NEIGHBORS = 5
 
 
 class DataSet(ABC):
@@ -141,6 +143,12 @@ def parse_context(context: Context) -> str:
         if context == Context.QUERY:
             return "test"
 
+        if context == Context.MAX_DISTANCE_NEIGHBORS:
+            return "max_distance_neighbors"
+
+        if context == Context.MIN_SCORE_NEIGHBORS:
+            return "min_score_neighbors"
+
         raise Exception("Unsupported context")