Fix search shuffle

pyproject.toml

@@ -8,7 +8,7 @@ authors = [
 dynamic = ["version"]
 requires-python = ">=3.12"
 dependencies = [
-    "numpy>=1.10",
+    "numpy>=2.4",
     "scalable-vs>=0.0.7",
     "tqdm>=4.67",
     "typer-slim>=0.15.2",

src/svsbench/build.py

@@ -73,7 +73,11 @@ def _read_args(argv: list[str] | None = None) -> argparse.Namespace:
         action="store_true",
     )
     parser.add_argument(
-        "--shuffle", help="Shuffle order of vectors", action="store_true"
+        "--shuffle", help="Shuffle order of vectors."
+        " Either search with svsbench.search --shuffle"
+        " or regenerate the ground truth with"
+        " svsbench.generate_ground_truth --shuffle",
+        action="store_true",
     )
     parser.add_argument(
         "--static", help="Index is static", action="store_true"
@@ -469,7 +473,7 @@ def save(
     index: svs.Vamana | svs.DynamicVamana,
     out_dir: Path = Path("out"),
     name: str = "index",
-) -> None:
+) -> Path:
     idx_dir = out_dir / name
     idx_dir.mkdir(exist_ok=True)
     index.save(
@@ -478,6 +482,7 @@ def save(
         str(idx_dir / "data"),
     )
     logger.info({"index_saved": idx_dir})
+    return idx_dir
 
 
 if __name__ == "__main__":

src/svsbench/generate_ground_truth.py

@@ -65,7 +65,7 @@ def generate_ground_truth(
     vecs_path: Path,
     query_file: Path,
     distance: svs.DistanceType,
-    num_vectors: int | None,
+    num_vectors: int | None = None,
     k: int = 100,
     num_threads: int = 1,
     out_file: Path | None = None,

src/svsbench/search.py

@@ -128,8 +128,8 @@ def _read_args(argv: list[str] | None = None) -> argparse.Namespace:
         default="all",
     )
     parser.add_argument(
-        "--shuffle_ground_truth",
-        help="Shuffle ground truth. See also svsbench.build --shuffle",
+        "--shuffle",
+        help="Search an index built with svsbench.build --shuffle",
         action="store_true",
     )
     return parser.parse_args(argv)
@@ -148,7 +148,7 @@ def search(
     max_threads=255,
     search_window_sizes: list[int] | None = None,
     recall: float = 0.9,
-    compress: bool = True,
+    compress: bool = False,
     leanvec_dims: int | None = -4,
     leanvec_alignment: int | None = 32,
     num_rep: int = 5,
@@ -164,9 +164,9 @@ def search(
     lvq_strategy: svs.LVQStrategy | None = None,
     train_prefetchers: bool = True,
     search_buffer_optimization: svs.VamanaSearchBufferOptimization = svs.VamanaSearchBufferOptimization.All,
-    shuffle_ground_truth: bool = False,
+    shuffle: bool = False,
     seed: int = 42,
-):
+) -> tuple[np.ndarray, np.ndarray, float]:
     logger.info({"search_args": locals()})
     logger.info(utils.read_system_config())
     if query_path is None:
@@ -216,8 +216,10 @@ def search(
 
     query = svs.read_vecs(str(query_path))
     ground_truth = svs.read_vecs(str(ground_truth_path))
-    if shuffle_ground_truth:
-        np.random.default_rng(seed).shuffle(ground_truth)
+    if shuffle:
+        dtype = ground_truth.dtype
+        permutation = np.random.default_rng(seed).permutation(index.size)
+        ground_truth = np.argsort(permutation)[ground_truth].astype(dtype=dtype, casting="same_value")
 
     for batch_size_idx, batch_size in enumerate(batch_sizes):
         index.num_threads = min(max_threads, batch_size)
@@ -300,18 +302,23 @@ def search(
         p95s = []
         for _ in tqdm(range(num_rep)):
             total_time = 0
-            results = np.empty((0, count), np.int32)
+            results = None
             batch_times = []
             for batch_idx in tqdm(range(num_batches)):
                 init_batch = batch_idx * batch_size
                 end_batch = min(init_batch + batch_size, query_size)
 
                 start = time.perf_counter()
-                result, _ = index.search(query[init_batch:end_batch], count)
+                result, batch_distances = index.search(query[init_batch:end_batch], count)
                 batch_time = time.perf_counter() - start
                 total_time += batch_time
                 batch_times.append(batch_time)
-                results = np.append(results, result, axis=0)
+                if results is None:
+                    results = result
+                    distances = batch_distances
+                else:
+                    results = np.append(results, result, axis=0)
+                    distances = np.append(distances, batch_distances, axis=0)
 
             qps.append(query_size / total_time)
             p95s.append(np.percentile(batch_times, 95))
@@ -338,6 +345,7 @@ def search(
                 },
             }
         )
+        return results, distances, recall
 
 
 def main(argv: str | None = None) -> None:
@@ -381,7 +389,7 @@ def main(argv: str | None = None) -> None:
         search_buffer_optimization=STR_TO_CALIBRATE_SEARCH_BUFFER[
             args.calibrate_search_buffer
         ],
-        shuffle_ground_truth=args.shuffle_ground_truth,
+        shuffle=args.shuffle,
         seed=args.seed,
     )
 

tests/test_search.py

@@ -3,7 +3,9 @@
 import pytest
 import svs
 
+from svsbench.build import build_dynamic, save
 from svsbench.consts import SVS_TYPES
+from svsbench.generate_ground_truth import generate_ground_truth
 from svsbench.search import search
 
 
@@ -30,7 +32,7 @@ def test_search(
             pytest.skip("Not supported")
     if svs_type != index_svs_type:
         compress = True
-    search(
+    _, _, recall = search(
         idx_dir=index_dir,
         svs_type=svs_type,
         distance=svs.DistanceType.L2,
@@ -40,7 +42,43 @@ def test_search(
         static=static,
         load_from_static=not index_dynamic,
     )
+    # Search parameters are calibrated to recall 0.9
+    assert recall > 0.8
 
 
 def test_search_with_separate_data_dir():
     pytest.xfail("TODO: Implement")
+
+
+@pytest.mark.parametrize("tmp_vecs", [".fvecs"], indirect=True)
+def test_search_with_shuffle(tmp_vecs, query_path, tmp_path):
+    seed = 123
+    svs_type = "float32"
+    distance = svs.DistanceType.L2
+
+    ground_truth_path = tmp_path / "ground_truth.ivecs"
+    generate_ground_truth(
+        vecs_path=tmp_vecs,
+        query_file=query_path,
+        distance=distance,
+        out_file=ground_truth_path,
+    )
+    build_result = build_dynamic(
+        vecs_path=tmp_vecs,
+        svs_type=svs_type,
+        distance=distance,
+        shuffle=True,
+        seed=seed,
+    )
+    idx_dir = save(build_result[0], tmp_path)
+    _, _, recall = search(
+        idx_dir=idx_dir,
+        svs_type=svs_type,
+        distance=distance,
+        ground_truth_path=ground_truth_path,
+        query_path=query_path,
+        shuffle=True,
+        seed=seed,
+    )
+    # Search parameters are calibrated to recall 0.9
+    assert recall > 0.8