clean Cargo.toml and add CLI for stand-alone-test

Signed-off-by: whateveraname <[email protected]>

Cargo.toml

@@ -10,31 +10,11 @@ crate-type = ["cdylib", "lib"]
 name = "pgrx_embed_vectors"
 path = "./src/bin/pgrx_embed.rs"
 
-[[bin]]
-name = "prepare_dataset"
-path = "./crates/stand-alone-test/src/bin/prepare_dataset.rs"
-required-features = ["stand-alone-test"]
-
-[[bin]]
-name = "make_hnsw"
-path = "./crates/stand-alone-test/src/bin/make_hnsw.rs"
-required-features = ["stand-alone-test"]
-
-[[bin]]
-name = "search_hnsw"
-path = "./crates/stand-alone-test/src/bin/search_hnsw.rs"
-required-features = ["stand-alone-test"]
-
 [features]
 default = ["pg15"]
 pg14 = ["pgrx/pg14"]
 pg15 = ["pgrx/pg15"]
 pg16 = ["pgrx/pg16"]
-stand-alone-test = [
-    "hnsw/stand-alone-test",
-    "storage/stand-alone-test",
-    "stand_alone_test/stand-alone-test",
-]
 
 [dependencies]
 arrayvec.workspace = true
@@ -69,7 +49,6 @@ memfd = { path = "crates/memfd" }
 scopeguard = "1.2.0"
 send_fd = { path = "crates/send_fd" }
 service = { path = "crates/service" }
-stand_alone_test = { path = "crates/stand-alone-test" }
 storage = { path = "crates/storage" }
 
 [patch.crates-io]

crates/hnsw/src/lib.rs

@@ -760,15 +760,15 @@ impl ElementHeap {
 }
 
 #[cfg(feature = "stand-alone-test")]
-fn mock_make(path: &Path, options: IndexOptions) -> HnswRam<Vecf32L2> {
+fn mock_make(path: &Path, data_path: &Path, options: IndexOptions) -> HnswRam<Vecf32L2> {
     let HnswIndexingOptions {
         m,
         ef_construction,
         quantization: quantization_opts,
     } = options.indexing.clone().unwrap_hnsw();
-    let vectors = MmapArray::open(Path::new("/home/yanqi/stand-alone-test/data/sift_vectors"));
-    let payload = MmapArray::open(Path::new("/home/yanqi/stand-alone-test/data/sift_payload"));
-    let dims = 128;
+    let vectors = MmapArray::open(&Path::new(data_path).join("vectors"));
+    let payload = MmapArray::open(&Path::new(data_path).join("payload"));
+    let dims = options.vector.dims as u16;
     let storage = Arc::new(StorageCollection::<Vecf32L2>::new(VecStorage::<F32>::new(
         vectors, payload, dims,
     )));
@@ -968,20 +968,20 @@ fn mock_make(path: &Path, options: IndexOptions) -> HnswRam<Vecf32L2> {
 }
 
 #[cfg(feature = "stand-alone-test")]
-pub fn mock_create(path: &Path, options: IndexOptions) -> Hnsw<Vecf32L2> {
+pub fn mock_create(path: &Path, data_path: &Path, options: IndexOptions) -> Hnsw<Vecf32L2> {
     create_dir(path).unwrap();
-    let ram = mock_make(path, options);
+    let ram = mock_make(path, data_path, options);
     let mmap = save(ram, path);
     sync_dir(path);
     Hnsw { mmap }
 }
 
 #[cfg(feature = "stand-alone-test")]
-pub fn mock_open(path: &Path, options: IndexOptions) -> Hnsw<Vecf32L2> {
+pub fn mock_open(path: &Path, data_path: &Path, options: IndexOptions) -> Hnsw<Vecf32L2> {
     let idx_opts = options.indexing.clone().unwrap_hnsw();
-    let vectors = MmapArray::open(Path::new("/home/yanqi/stand-alone-test/data/sift_vectors"));
-    let payload = MmapArray::open(Path::new("/home/yanqi/stand-alone-test/data/sift_payload"));
-    let dims = 128;
+    let vectors = MmapArray::open(&Path::new(data_path).join("vectors"));
+    let payload = MmapArray::open(&Path::new(data_path).join("payload"));
+    let dims = options.vector.dims as u16;
     let storage = Arc::new(StorageCollection::<Vecf32L2>::new(VecStorage::<F32>::new(
         vectors, payload, dims,
     )));

crates/stand-alone-test/Cargo.toml

@@ -4,13 +4,16 @@ version.workspace = true
 edition.workspace = true
 
 [features]
-stand-alone-test = []
+default = ["stand-alone-test"]
+stand-alone-test = ["hnsw/stand-alone-test", "storage/stand-alone-test"]
 
 [dependencies]
 base = { path = "../base" }
+clap = { version = "4.5.4", features = ["derive"] }
 common = { path = "../common" }
 hnsw = { path = "../hnsw" }
 rayon = { path = "../rayon" }
+storage = { path = "../storage" }
 thiserror = "1.0.58"
 
 [lints]

crates/stand-alone-test/src/bin/main.rs

@@ -0,0 +1,103 @@
+use clap::{Parser, Subcommand};
+use stand_alone_test::{make_hnsw, prepare_dataset, search_hnsw};
+
+#[derive(Parser)]
+struct Args {
+    #[command(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// Prepare dataset for testing
+    PrepareDataset {
+        /// path to the data file in fvecs format
+        #[arg(short, long)]
+        data_file: String,
+
+        /// directory to save the prepared dataset
+        #[arg(short, long)]
+        output_dir: String,
+    },
+
+    /// Make HNSW index
+    MakeHnsw {
+        /// directory of the prepared dataset
+        #[arg(short, long)]
+        data_dir: String,
+
+        /// dimension of the data
+        #[arg(long)]
+        dims: u32,
+
+        /// m parameter for HNSW index
+        #[arg(short, long)]
+        m: u32,
+
+        /// ef_construction parameter for HNSW index
+        #[arg(short, long)]
+        ef_construction: usize,
+
+        /// directory to save the hnsw index
+        #[arg(short, long)]
+        output_dir: String,
+    },
+
+    /// Search HNSW index
+    SearchHnsw {
+        /// directory of the prepared dataset
+        #[arg(short, long)]
+        data_dir: String,
+
+        /// dimension of the data
+        #[arg(long)]
+        dims: u32,
+
+        /// directory of the hnsw index
+        #[arg(short, long)]
+        hnsw_dir: String,
+
+        /// path to the query file in fvecs format
+        #[arg(short, long)]
+        query_file: String,
+
+        /// path to the ground truth file in ivecs format
+        #[arg(short, long)]
+        gt_file: String,
+
+        /// ef_search parameter for HNSW search
+        #[arg(long)]
+        ef: usize,
+    },
+}
+
+fn main() {
+    let args = Args::parse();
+    match &args.command {
+        Commands::PrepareDataset {
+            data_file,
+            output_dir,
+        } => {
+            prepare_dataset(data_file, output_dir);
+        }
+        Commands::MakeHnsw {
+            data_dir,
+            dims,
+            m,
+            ef_construction,
+            output_dir,
+        } => {
+            make_hnsw(data_dir, *dims, *m, *ef_construction, output_dir);
+        }
+        Commands::SearchHnsw {
+            data_dir,
+            dims,
+            hnsw_dir,
+            query_file,
+            gt_file,
+            ef,
+        } => {
+            search_hnsw(data_dir, *dims, hnsw_dir, query_file, gt_file, *ef);
+        }
+    }
+}