Update gen_fbin_file.cpp to support binary data type (#211)

Signed-off-by: Yudong Cai <[email protected]>
zilliztech · Nov 24, 2023 · 439a0dd · 439a0dd
1 parent 9534d56
commit 439a0dd
Showing 1 changed file with 58 additions and 0 deletions.
diff --git a/benchmark/hdf5/gen_fbin_file.cpp b/benchmark/hdf5/gen_fbin_file.cpp
@@ -52,6 +52,14 @@ class Create_FBIN : public Benchmark_hdf5, public ::testing::Test {
         writer((void*)data, rows * dim * sizeof(float));
     }
 
+    void
+    fbin_write_binary(const std::string& filename, const uint32_t rows, const uint32_t dim, const void* data) {
+        FileIOWriter writer(filename);
+        writer((void*)&rows, sizeof(rows));
+        writer((void*)&dim, sizeof(dim));
+        writer((void*)data, rows * (dim / 8) * sizeof(uint8_t));
+    }
+
     void
     fbin_read(const std::string& filename, uint32_t& rows, uint32_t& dim, void* data) {
         FileIOReader reader(filename);
@@ -60,6 +68,14 @@ class Create_FBIN : public Benchmark_hdf5, public ::testing::Test {
         reader((void*)data, rows * dim * sizeof(float));
     }
 
+    void
+    fbin_read_binary(const std::string& filename, uint32_t& rows, uint32_t& dim, void* data) {
+        FileIOReader reader(filename);
+        reader((void*)&rows, sizeof(rows));
+        reader((void*)&dim, sizeof(dim));
+        reader((void*)data, rows * (dim / 8) * sizeof(uint8_t));
+    }
+
     void
     fbin_result_write(const std::string& filename, const uint32_t rows, const uint32_t topk, const uint32_t* ids,
                       const float* dist) {
@@ -173,3 +189,45 @@ TEST_F(Create_FBIN, HDF5_RANGE_TO_FBIN) {
 
     free_all();
 }
+
+TEST_F(Create_FBIN, HDF5_BIN_TO_FBIN) {
+    set_ann_test_name("rand-1024-hamming");
+    parse_ann_test_name();
+    load_hdf5_data<true>();
+
+    std::string prefix = dataset_name_ + "-" + std::to_string(dim_) + "-";
+    std::string postfix = ".fbin";
+    std::string filename;
+
+    filename = prefix + "base" + postfix;
+    fbin_write_binary(filename, nb_, dim_, xb_);
+
+    filename = prefix + "query" + postfix;
+    fbin_write_binary(filename, nq_, dim_, xq_);
+
+    filename = prefix + metric_str_ + "-gt" + postfix;
+    fbin_result_write(filename, nq_, gt_k_, (uint32_t*)gt_ids_, gt_dist_);
+
+    free_all();
+}
+
+TEST_F(Create_FBIN, HDF5_BIN_RANGE_TO_FBIN) {
+    set_ann_test_name("rand-1024-hamming-range");
+    parse_ann_test_name_with_range();
+    load_hdf5_data_range<true>();
+
+    std::string prefix = dataset_name_ + "-" + std::to_string(dim_) + "-range-";
+    std::string postfix = ".fbin";
+    std::string filename;
+
+    filename = prefix + "base" + postfix;
+    fbin_write_binary(filename, nb_, dim_, xb_);
+
+    filename = prefix + "query" + postfix;
+    fbin_write_binary(filename, nq_, dim_, xq_);
+
+    filename = prefix + metric_str_ + "-gt" + postfix;
+    fbin_range_result_write(filename, nq_, *gt_radius_, (uint32_t*)gt_lims_, (uint32_t*)gt_ids_, gt_dist_);
+
+    free_all();
+}