Scripts for Billion-Scale Synthetic Data Generation

dependencies.yaml

@@ -222,6 +222,15 @@ files:
       key: opensearch
     includes:
       - bench_python_opensearch
+  py_synthesize_dataset_cuvs_bench:
+    output: pyproject
+    pyproject_dir: python/cuvs_bench
+    extras:
+      table: project.optional-dependencies
+      key: synthesize_dataset
+    includes:
+      - bench_python_synthesize_dataset
+      - depends_on_cupy
 channels:
   - rapidsai-nightly
   - rapidsai
@@ -625,6 +634,12 @@ dependencies:
       - output_types: [conda, pyproject]
         packages:
           - opensearch-py>=2.4.0
+  bench_python_synthesize_dataset:
+    common:
+      - output_types: pyproject
+        packages:
+          - *numpy
+          - tqdm
   depends_on_cuda_python:
     specific:
       - output_types: [conda, requirements, pyproject]

fern/docs.yml

@@ -215,6 +215,8 @@ navigation:
                 path: "./pages/cuvs_bench/running.md"
               - page: "Datasets"
                 path: "./pages/cuvs_bench/datasets.md"
+              - page: "Synthetic Dataset Generation"
+                path: "./pages/cuvs_bench/synthesize_dataset.md"
               - page: "Backends"
                 path: "./pages/cuvs_bench/pluggable_backend.md"
               - page: "cuVS Bench Parameter Tuning Guide"

python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 import argparse
@@ -10,6 +10,7 @@
 import warnings
 
 from .utils import (
+    add_jitter,
     groundtruth_neighbors_filename,
     memmap_bin_file,
     offset_neighbor_indices,
@@ -111,6 +112,22 @@ def choose_random_queries(dataset, n_queries):
     return dataset[query_idx, :]
 
 
+def choose_random_queries_with_jitter(dataset, n_queries, seed=12345):
+    """Pick ``n_queries`` random rows from ``dataset`` and add Gaussian jitter
+    at scale ``0.1 * std(sample)``.
+    """
+    import numpy as _np
+
+    print("Choosing random vectors from dataset and jittering with noise")
+    rng = _np.random.default_rng(seed)
+    n_rows = dataset.shape[0]
+    # Sort indices so the memmap read is sequential rather than random-access.
+    query_idx = _np.sort(rng.choice(n_rows, size=n_queries, replace=False))
+    sampled = dataset[query_idx, :].astype(_np.float32, copy=True)
+
+    return add_jitter(sampled, rng, normalize=False)
+
+
 def cpu_search(dataset, queries, k, metric="squeclidean"):
     """
     Find the k nearest neighbors for each query point in the dataset using the
@@ -235,18 +252,22 @@ def main():
         "The input and output files are in big-ann-benchmark's binary format.",
         epilog="""Example usage
     # With existing query file
-    python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\
-fbin --output=groundtruth_dir --queries=/dataset/query.public.10K.fbin
+    python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \
+--output=groundtruth_dir --queries=/dataset/query.public.10K.fbin
 
     # With randomly generated queries
-    python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\
-fbin --output=groundtruth_dir --queries=random --n_queries=10000
+    python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \
+--output=groundtruth_dir --queries=random --n_queries=10000
 
     # Using only a subset of the dataset. Define queries by randomly
     # selecting vectors from the (subset of the) dataset.
-    python -m cuvs_bench.generate_groundtruth --dataset /dataset/base.\
-fbin --nrows=2000000 --cols=128 --output=groundtruth_dir \
+    python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \
+--rows=2000000 --cols=128 --output=groundtruth_dir \
 --queries=random-choice --n_queries=10000
+
+    # Jittered queries (following the logic of cuvs_bench.synthesize_dataset)
+    python -m cuvs_bench.generate_groundtruth /dataset/base.fbin \
+--output=groundtruth_dir --queries=random-jitter --n_queries=10000
     """,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
@@ -256,9 +277,11 @@ def main():
         "--queries",
         type=str,
         default="random",
-        help="Queries file name, or one of 'random-choice' or 'random' "
-        "(default). 'random-choice': select n_queries vectors from the input "
-        "dataset. 'random': generate n_queries as uniform random numbers.",
+        help="Queries file name, or one of 'random-choice', 'random-jitter', "
+        "or 'random' (default). 'random-choice': select n_queries vectors "
+        "from the input dataset. 'random-jitter': same as 'random-choice', "
+        "but add std-relative Gaussian noise to each query. 'random': generate "
+        "n_queries as uniform random numbers.",
     )
     parser.add_argument(
         "--output",
@@ -341,7 +364,7 @@ def main():
     if len(args.output) > 0:
         os.makedirs(args.output, exist_ok=True)
 
-    if args.queries == "random" or args.queries == "random-choice":
+    if args.queries in {"random", "random-choice", "random-jitter"}:
         if args.n_queries is None:
             raise RuntimeError(
                 "n_queries must be given to generate random queries"
@@ -352,9 +375,13 @@ def main():
             )
         elif args.queries == "random-choice":
             queries = choose_random_queries(dataset, args.n_queries)
+        elif args.queries == "random-jitter":
+            queries = choose_random_queries_with_jitter(
+                dataset, args.n_queries
+            )
 
         queries_filename = os.path.join(
-            args.output, "queries" + suffix_from_dtype(dtype)
+            args.output, "queries" + suffix_from_dtype(queries.dtype)
         )
         print("Writing queries file", queries_filename)
         write_bin(queries_filename, queries)

python/cuvs_bench/cuvs_bench/generate_groundtruth/utils.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -10,6 +10,22 @@
 from cuvs_bench._bin_format import read_bin_header, write_bin_header
 
 
+def add_jitter(
+    queries: np.ndarray,
+    rng: np.random.Generator,
+    normalize: bool,
+) -> np.ndarray:
+    """Add Gaussian jitter to query vectors and optionally re-normalize."""
+    noise_scale = float(np.std(queries)) * 0.1
+    queries = queries + rng.normal(0, noise_scale, queries.shape).astype(
+        np.float32
+    )
+    if normalize:
+        norms = np.linalg.norm(queries, axis=1, keepdims=True)
+        queries = queries / np.maximum(norms, 1e-8)
+    return queries.astype(np.float32)
+
+
 def dtype_from_filename(filename):
     ext = os.path.splitext(filename)[1]
     if ext == ".fbin":

python/cuvs_bench/cuvs_bench/synthesize_dataset/__init__.py

@@ -0,0 +1,42 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from ._fingerprint import Fingerprint
+from ._fit import fit_fingerprint
+from ._generate import (
+    gen_cluster_gpu,
+    generate_queries,
+    generate_synthetic_dataset,
+    generate_synthetic_dataset_to_file,
+    get_cluster_seed,
+    get_num_points_per_cluster,
+    resolve_norm_scheme,
+)
+from ._ground_truth import (
+    compute_groundtruth_exact,
+    compute_groundtruth_nprobe,
+)
+from ._io import (
+    load_fingerprint,
+    save_fingerprint,
+)
+from ._verify import verify_groundtruth
+
+__all__ = [
+    "Fingerprint",
+    "load_fingerprint",
+    "compute_groundtruth_nprobe",
+    "compute_groundtruth_exact",
+    "fit_fingerprint",
+    "gen_cluster_gpu",
+    "generate_queries",
+    "generate_synthetic_dataset",
+    "generate_synthetic_dataset_to_file",
+    "get_cluster_seed",
+    "get_num_points_per_cluster",
+    "resolve_norm_scheme",
+    "save_fingerprint",
+    "verify_groundtruth",
+]