Support per-percentage ground truth datasets via filter_percentage param

Adds optional key_suffix to HDF5 dataset reads. When filter_percentage
is set (e.g. "10pct"), the neighbors and radial threshold datasets are
read from suffixed keys (neighbors_10pct, faiss_max_distance_10pct),
allowing one dataset file to hold precomputed answer sets for multiple
filter percentages.

Signed-off-by: Wenxin Li <liwenxin@amazon.com>

Author: Wenxin Li

osbenchmark/utils/dataset.py

@@ -72,17 +72,20 @@ def reset(self):
         """
 
 
-def get_data_set(data_set_format: str, path: str, context: Context):
+def get_data_set(data_set_format: str, path: str, context: Context, key_suffix: str = None):
     """
     Factory method to get instance of Dataset for given format.
     Args:
         data_set_format: File format like hdf5, bigann
         path: Data set file path
         context: Dataset Context Enum
+        key_suffix: Optional suffix appended to the HDF5 key name, e.g.
+            Context.NEIGHBORS with key_suffix "10pct" reads "neighbors_10pct".
+            Used by filtered benchmarks storing per-percentage ground truth.
     Returns: DataSet instance
     """
     if data_set_format == HDF5DataSet.FORMAT_NAME:
-        return HDF5DataSet(path, context)
+        return HDF5DataSet(path, context, key_suffix)
     if data_set_format == BigANNVectorDataSet.FORMAT_NAME:
         return create_big_ann_dataset(path)
     raise ConfigurationError("Invalid data set format")
@@ -95,9 +98,11 @@ class HDF5DataSet(DataSet):
 
     FORMAT_NAME = "hdf5"
 
-    def __init__(self, dataset_path: str, context: Context):
+    def __init__(self, dataset_path: str, context: Context, key_suffix: str = None):
         self.dataset_path = dataset_path
         self.context = self.parse_context(context)
+        if key_suffix:
+            self.context = f"{self.context}_{key_suffix}"
         self.current = self.BEGINNING
         self.data = None
 

osbenchmark/workload/params.py

@@ -1162,6 +1162,10 @@ def __init__(self, workloads, params, query_params, **kwargs):
         self.filter_type = self.query_params.get(self.PARAMS_NAME_FILTER_TYPE)
         self.filter_body = self.query_params.get(self.PARAMS_NAME_FILTER_BODY)
         self.space_type = params.get(self.PARAMS_NAME_SPACE_TYPE, "l2")
+        # Suffix selecting per-percentage ground truth datasets, e.g. "10pct"
+        # reads neighbors_10pct and faiss_max_distance_10pct instead of the
+        # unsuffixed keys. Used by filtered benchmarks.
+        self.filter_percentage = params.get("filter_percentage")
 
         if self.radial_search_type:
             self.radial_engine = params.get("radial_engine", "faiss")
@@ -1234,14 +1238,16 @@ def partition(self, partition_index, total_partitions):
             neighbors_context = Context.NEIGHBORS
 
         partition.neighbors_data_set = get_data_set(
-            self.neighbors_data_set_format, self.neighbors_data_set_path, neighbors_context)
+            self.neighbors_data_set_format, self.neighbors_data_set_path, neighbors_context,
+            self.filter_percentage)
         partition.neighbors_data_set.seek(partition.offset)
 
         if self.radial_search_type:
             threshold_context = self.RADIAL_THRESHOLD_CONTEXTS[
                 (self.radial_engine, self.radial_search_type)]
             partition.threshold_data_set = get_data_set(
-                self.neighbors_data_set_format, self.neighbors_data_set_path, threshold_context)
+                self.neighbors_data_set_format, self.neighbors_data_set_path, threshold_context,
+                self.filter_percentage)
             partition.threshold_data_set.seek(partition.offset)
 
         return partition