Support per-percentage ground truth datasets via filter_percentage param

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,15 +98,22 @@ 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
 
     def _load(self):
         if self.data is None:
             file = h5py.File(self.dataset_path)
+            if self.context not in file:
+                raise ConfigurationError(
+                    f"Dataset key '{self.context}' does not exist in {self.dataset_path}. "
+                    f"Available keys: {sorted(file.keys())}. If filter_percentage is set, "
+                    f"ensure the dataset was generated with that percentage.")
             self.data = cast(h5py.Dataset, file[self.context])
 
     def read(self, chunk_size: int):

osbenchmark/workload/params.py

@@ -1162,6 +1162,12 @@ 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 = None
+        if "filter_percentage" in params:
+            self.filter_percentage = parse_string_parameter("filter_percentage", params)
 
         if self.radial_search_type:
             index_body = params.get("target_index_body", "")
@@ -1238,14 +1244,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

tests/utils/dataset_test.py

@@ -3,9 +3,13 @@
 # The OpenSearch Contributors require contributions made to
 # this file be licensed under the Apache-2.0 license or a
 # compatible open source license.
+import os
 import tempfile
 from unittest import TestCase
 
+import h5py
+import numpy as np
+
 from osbenchmark.utils.dataset import Context, get_data_set, HDF5DataSet, BigANNVectorDataSet
 from osbenchmark.utils.parse import ConfigurationError
 from tests.utils.dataset_helper import create_data_set, create_ground_truth
@@ -66,3 +70,39 @@ def testBigANNGroundTruthAsAcceptableDataSetFormat(self):
     def testUnSupportedDataSetFormat(self):
         with self.assertRaises(ConfigurationError) as _:
             get_data_set("random", "/some/path", Context.INDEX)
+
+    def testHDF5KeySuffix(self):
+        with tempfile.TemporaryDirectory() as data_set_dir:
+            data_set_path = os.path.join(data_set_dir, "key-suffix.hdf5")
+            neighbors = np.zeros((DEFAULT_NUM_VECTORS, DEFAULT_DIMENSION), dtype=np.int32)
+            neighbors_10pct = np.ones((DEFAULT_NUM_VECTORS, DEFAULT_DIMENSION), dtype=np.int32)
+            max_distance_10pct = np.full((DEFAULT_NUM_VECTORS, DEFAULT_DIMENSION), 2.0, dtype=np.float32)
+            with h5py.File(data_set_path, "w") as file:
+                file.create_dataset("neighbors", data=neighbors)
+                file.create_dataset("neighbors_10pct", data=neighbors_10pct)
+                file.create_dataset("faiss_max_distance_10pct", data=max_distance_10pct)
+
+            # without suffix, the unsuffixed key is read
+            data_set_instance = get_data_set("hdf5", data_set_path, Context.NEIGHBORS)
+            self.assertTrue(np.array_equal(data_set_instance.read(1)[0], neighbors[0]))
+
+            # with suffix, the suffixed keys are read
+            data_set_instance = get_data_set("hdf5", data_set_path, Context.NEIGHBORS, "10pct")
+            self.assertTrue(np.array_equal(data_set_instance.read(1)[0], neighbors_10pct[0]))
+
+            data_set_instance = get_data_set(
+                "hdf5", data_set_path, Context.FAISS_MAX_DISTANCE, "10pct")
+            self.assertTrue(np.array_equal(data_set_instance.read(1)[0], max_distance_10pct[0]))
+
+    def testHDF5MissingKeyRaisesConfigurationError(self):
+        with tempfile.TemporaryDirectory() as data_set_dir:
+            data_set_path = os.path.join(data_set_dir, "missing-key.hdf5")
+            with h5py.File(data_set_path, "w") as file:
+                file.create_dataset(
+                    "neighbors", data=np.zeros((DEFAULT_NUM_VECTORS, DEFAULT_DIMENSION), dtype=np.int32))
+
+            data_set_instance = get_data_set("hdf5", data_set_path, Context.NEIGHBORS, "5pct")
+            with self.assertRaises(ConfigurationError) as ctx:
+                data_set_instance.read(1)
+            self.assertIn("neighbors_5pct", str(ctx.exception))
+            self.assertIn("Available keys", str(ctx.exception))