From 3eb00546ba5c4be8a7218affcd475a1c5584c0b1 Mon Sep 17 00:00:00 2001
From: Neelam Mahapatro <nmahapatro@microsoft.com>
Date: Mon, 6 Feb 2023 10:07:45 +0530
Subject: [PATCH 01/13] Fix filter diskann issue

- Fix sharded filter build
- fix search label parsing
---
 src/disk_utils.cpp     | 2 +-
 src/pq_flash_index.cpp | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp
index 1e1286418..9b7c3ca24 100644
--- a/src/disk_utils.cpp
+++ b/src/disk_utils.cpp
@@ -1282,7 +1282,7 @@ namespace diskann {
     gen_random_slice<T>(data_file_to_use.c_str(), sample_base_prefix,
                         sample_sampling_rate);
     if (use_filters) {
-      copy_file(mem_labels_file, disk_labels_file);
+      copy_file(labels_file_to_use, disk_labels_file);
       std::remove(mem_labels_file.c_str());
       if (universal_label != "") {
         copy_file(mem_univ_label_file, disk_univ_label_file);
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
index 1db901b92..f6ac79b34 100644
--- a/src/pq_flash_index.cpp
+++ b/src/pq_flash_index.cpp
@@ -507,7 +507,7 @@ namespace diskann {
       _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter];
       num_lbls_in_cur_pt = 0;
       counter++;
-      getline(iss, token, '\t');  // first token contains metadata, not used
+      //getline(iss, token, '\t');  // first token contains metadata, not used
       getline(iss, token, '\t');
       std::istringstream new_iss(token);
       while (getline(new_iss, token, ',')) {

From 5650917f381bb5937ae46fef7476b07f9a28a867 Mon Sep 17 00:00:00 2001
From: Neelam Mahapatro <nmahapatro@microsoft.com>
Date: Thu, 9 Feb 2023 22:20:33 +0530
Subject: [PATCH 02/13] Fix comments

---
 include/disk_utils.h                |   2 +-
 include/index.h                     |   2 +-
 src/disk_utils.cpp                  |   7 +-
 src/index.cpp                       | 158 ++++++++++++----------
 src/pq_flash_index.cpp              | 195 +++++++++++++++-------------
 tests/search_disk_index.cpp         |   1 +
 tests/search_memory_index.cpp       |   1 +
 tests/utils/compute_groundtruth.cpp |   3 -
 tests/utils/stats_label_data.cpp    |  17 +--
 9 files changed, 201 insertions(+), 185 deletions(-)

diff --git a/include/disk_utils.h b/include/disk_utils.h
index dabeeeb0c..6aaa02426 100644
--- a/include/disk_utils.h
+++ b/include/disk_utils.h
@@ -75,7 +75,7 @@ namespace diskann {
                                      const _u64 nshards, unsigned max_degree,
                                      const std::string &output_vamana,
                                      const std::string &medoids_file, bool use_filters = false,
-      const std::string &labels_to_medoids_file = std::string(""));
+                                     const std::string &labels_to_medoids_file = std::string(""));
 
 DISKANN_DLLEXPORT void extract_shard_labels(
       const std::string &in_label_file, const std::string &shard_ids_bin,
diff --git a/include/index.h b/include/index.h
index 9025a22f5..871c2ddd7 100644
--- a/include/index.h
+++ b/include/index.h
@@ -238,7 +238,7 @@ namespace diskann {
     // determines navigating node of the graph by calculating medoid of datafopt
     unsigned calculate_entry_point();
 
-    void parse_label_file(const std::string& map_file);
+    size_t parse_label_file(const std::string &map_file);
 
     std::pair<uint32_t, uint32_t> iterate_to_fixed_point(
         const T *node_coords, const unsigned Lindex,
diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp
index 9b7c3ca24..e6aa02fdb 100644
--- a/src/disk_utils.cpp
+++ b/src/disk_utils.cpp
@@ -323,6 +323,7 @@ namespace diskann {
       }
 
       std::ofstream mapping_writer(labels_to_medoids_file);
+      assert(mapping_writer.is_open());
       for (auto iter : global_label_to_medoids) {
         mapping_writer << iter.first << ", ";
         auto &vec = iter.second;
@@ -534,6 +535,7 @@ namespace diskann {
       diskann::cout << labels_per_point.size() << " is the new number of points"
                     << std::endl;
       std::ofstream label_writer(out_labels_file);
+      assert(label_writer.is_open());
       for (_u32 i = 0; i < labels_per_point.size(); i++) {
         for (_u32 j = 0; j < (labels_per_point[i].size() - 1); j++) {
           label_writer << labels_per_point[i][j] << ",";
@@ -551,6 +553,7 @@ namespace diskann {
       data = (T *) std::realloc((void *) data,
                                 labels_per_point.size() * ndims * sizeof(T));
       std::ofstream dummy_writer(out_metadata_file);
+      assert(dummy_writer.is_open());
       for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) {
         dummy_writer << i->first << "," << i->second << std::endl;
         std::memcpy(data + i->first * ndims, data + i->second * ndims,
@@ -566,7 +569,7 @@ namespace diskann {
       const std::string &in_label_file, const std::string &shard_ids_bin,
       const std::string &shard_label_file) {  // assumes ith row is for ith
                                               // point in labels file
-    std::cout << "Extracting labels for shard" << std::endl;
+    diskann::cout << "Extracting labels for shard" << std::endl;
 
     _u32 *ids = nullptr;
     _u64  num_ids, tmp_dim;
@@ -578,6 +581,8 @@ namespace diskann {
     std::ifstream label_reader(in_label_file);
     std::ofstream label_writer(shard_label_file);
 
+    assert(label_reader.is_open());
+    assert(label_writer.is_open());
     if (label_reader && label_writer) {
       while (std::getline(label_reader, cur_line)) {
         if (shard_counter >= num_ids) {
diff --git a/src/index.cpp b/src/index.cpp
index f0e5e19a5..5ff702317 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -278,9 +278,9 @@ namespace diskann {
         if (_filter_to_medoid_id.size() > 0) {
           std::ofstream medoid_writer(std::string(filename) +
                                       "_labels_to_medoids.txt");
+          assert(medoid_writer.is_open());
           for (auto iter : _filter_to_medoid_id) {
             medoid_writer << iter.first << ", " << iter.second << std::endl;
-            // std::cout << iter.first << ", " << iter.second << std::endl;
           }
           medoid_writer.close();
         }
@@ -288,12 +288,14 @@ namespace diskann {
         if (_use_universal_label) {
           std::ofstream universal_label_writer(std::string(filename) +
                                                "_universal_label.txt");
+          assert(universal_label_writer.is_open());
           universal_label_writer << _universal_label << std::endl;
           universal_label_writer.close();
         }
 
         if (_pts_to_labels.size() > 0) {
           std::ofstream label_writer(std::string(filename) + "_labels.txt");
+          assert(label_writer.is_open());
           for (_u32 i = 0; i < _pts_to_labels.size(); i++) {
             for (_u32 j = 0; j < (_pts_to_labels[i].size() - 1); j++) {
               label_writer << _pts_to_labels[i][j] << ",";
@@ -476,52 +478,12 @@ namespace diskann {
 
     _has_built = true;
 
-    size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0;
+    size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, label_num_pts = 0;
 
     std::string mem_index_file(filename);
     std::string labels_file = mem_index_file + "_labels.txt";
     std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt";
-    if (file_exists(labels_file)) {
-      parse_label_file(labels_file);
-      if (file_exists(labels_to_medoids)) {
-        std::ifstream medoid_stream(labels_to_medoids);
-
-        std::string line, token;
-        unsigned    line_cnt = 0;
-
-        _filter_to_medoid_id.clear();
-
-        while (std::getline(medoid_stream, line)) {
-          std::istringstream iss(line);
-          _u32               cnt = 0;
-          _u32               medoid = 0;
-          label              label;
-          while (std::getline(iss, token, ',')) {
-            token.erase(std::remove(token.begin(), token.end(), '\n'),
-                        token.end());
-            token.erase(std::remove(token.begin(), token.end(), '\r'),
-                        token.end());
-            unsigned token_as_num = std::stoul(token);
-            if (cnt == 0)
-              label = token_as_num;
-            else
-              medoid = token_as_num;
-            cnt++;
-          }
-          _filter_to_medoid_id[label] = medoid;
-          line_cnt++;
-        }
-      }
-
-      std::string universal_label_file(filename);
-      universal_label_file += "_universal_label.txt";
-      if (file_exists(universal_label_file)) {
-        std::ifstream universal_label_reader(universal_label_file);
-        universal_label_reader >> _universal_label;
-        _use_universal_label = true;
-        universal_label_reader.close();
-      }
-    }
+    
 
     if (!_save_as_one_file) {
       // For DLVS Store, we will not support saving the index in multiple files.
@@ -560,6 +522,53 @@ namespace diskann {
       throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__,
                                   __LINE__);
     }
+    if (file_exists(labels_file)) {
+      label_num_pts = parse_label_file(labels_file);
+      assert(label_num_pts == data_file_num_pts);
+      if (file_exists(labels_to_medoids)) {
+        std::ifstream medoid_stream(labels_to_medoids);
+        assert(medoid_stream.is_open());        
+        std::string line, token;
+        unsigned    line_cnt = 0;
+
+        _filter_to_medoid_id.clear();
+        try {
+          while (std::getline(medoid_stream, line)) {
+            std::istringstream iss(line);
+            _u32               cnt = 0;
+            _u32               medoid = 0;
+            label              label;
+            while (std::getline(iss, token, ',')) {
+              token.erase(std::remove(token.begin(), token.end(), '\n'),
+                          token.end());
+              token.erase(std::remove(token.begin(), token.end(), '\r'),
+                          token.end());
+              unsigned token_as_num = std::stoul(token);
+              if (cnt == 0)
+                label = token_as_num;
+              else
+                medoid = token_as_num;
+              cnt++;
+            }
+            _filter_to_medoid_id[label] = medoid;
+            line_cnt++;
+          }
+        } catch (std::system_error &e) {
+          throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__,
+                              __LINE__);
+        }
+      }
+
+      std::string universal_label_file(filename);
+      universal_label_file += "_universal_label.txt";
+      if (file_exists(universal_label_file)) {
+        std::ifstream universal_label_reader(universal_label_file);
+        assert(universal_label_reader.is_open());
+        universal_label_reader >> _universal_label;
+        _use_universal_label = true;
+        universal_label_reader.close();
+      }
+    }
 
     _nd = data_file_num_pts - _num_frozen_pts;
     _empty_slots.clear();
@@ -1645,13 +1654,13 @@ namespace diskann {
   }
 
   template<typename T, typename TagT>
-  void Index<T, TagT>::parse_label_file(const std::string &map_file) {
+  size_t Index<T, TagT>::parse_label_file(const std::string &label_file) {
     // Format of Label txt file: filters with comma separators
-
-    std::ifstream infile(map_file);
+    std::ifstream infile(label_file);
+    assert(infile.is_open());
     std::string   line, token;
     unsigned      line_cnt = 0;
-
+    
     while (std::getline(infile, line)) {
       line_cnt++;
     }
@@ -1660,28 +1669,35 @@ namespace diskann {
     infile.clear();
     infile.seekg(0, std::ios::beg);
     line_cnt = 0;
-    while (std::getline(infile, line)) {
-      std::istringstream iss(line);
-      std::vector<label> lbls(0);
-      getline(iss, token, '\t');
-      std::istringstream new_iss(token);
-      while (getline(new_iss, token, ',')) {
-        token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-        token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-        unsigned token_as_num = std::stoul(token);
-        lbls.push_back(token_as_num);
-        _labels.insert(token_as_num);
-      }
-      if (lbls.size() <= 0) {
-        std::cout << "No label found";
-        exit(-1);
+    try {
+      while (std::getline(infile, line)) {
+        std::istringstream iss(line);
+        std::vector<label> lbls(0);
+        getline(iss, token, '\t');
+        std::istringstream new_iss(token);
+        while (getline(new_iss, token, ',')) {
+          token.erase(std::remove(token.begin(), token.end(), '\n'),
+                      token.end());
+          token.erase(std::remove(token.begin(), token.end(), '\r'),
+                      token.end());
+          unsigned token_as_num = std::stoul(token);
+          lbls.push_back(token_as_num);
+          _labels.insert(token_as_num);
+        }
+        if (lbls.size() <= 0) {
+          diskann::cout << "No label found";
+          exit(-1);
+        }
+        std::sort(lbls.begin(), lbls.end());
+        _pts_to_labels[line_cnt] = lbls;
+        line_cnt++;
       }
-      std::sort(lbls.begin(), lbls.end());
-      _pts_to_labels[line_cnt] = lbls;
-      line_cnt++;
+      diskann::cout << "Identified " << _labels.size() << " distinct label(s)"
+                    << std::endl;
+    } catch (std::system_error &e) {
+      throw FileException(label_file, e, __FUNCSIG__, __FILE__, __LINE__);
     }
-    std::cout << "Identified " << _labels.size() << " distinct label(s)"
-              << std::endl;
+    return line_cnt;
   }
 
   template<typename T, typename TagT>
@@ -1699,9 +1715,9 @@ namespace diskann {
     _labels_file = label_file;
     _filtered_index = true;
     _filter_to_medoid_id.clear();
-    parse_label_file(label_file);  // determines medoid for each label and
+    size_t label_num_pts = parse_label_file(label_file);  // determines medoid for each label and
                                    // identifies the points to label mapping
-
+    assert(label_num_pts != num_points_to_load);
     _u32 counter = 0;
 #pragma omp parallel for schedule(dynamic, 1)
     for (int lbl = 0; lbl < _labels.size(); lbl++) {
@@ -1741,15 +1757,13 @@ namespace diskann {
           _filter_to_medoid_id[x] = best_medoid;
           _medoid_counts[best_medoid]++;
           std::stringstream a;
-          // a << "Medoid of " << x << " is " << best_medoid <<
-          // std::endl; std::cout << a.str();
         }
       }
 #pragma omp critical
       counter++;
       std::stringstream a;
       a << ((100.0 * counter) / _labels.size()) << "\% processed \r";
-      std::cout << a.str() << std::flush;
+      diskann::cout << a.str() << std::flush;
     }
 
     this->build(filename, num_points_to_load, parameters, tags);
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
index f6ac79b34..a86c14869 100644
--- a/src/pq_flash_index.cpp
+++ b/src/pq_flash_index.cpp
@@ -464,7 +464,7 @@ namespace diskann {
       num_pts++;
     }
 
-    std::cout << "Labels file metadata: num_points: " << num_pts
+    diskann::cout << "Labels file metadata: num_points: " << num_pts
               << ", #total_labels: " << num_total_labels << std::endl;
     infile.close();
   }
@@ -484,55 +484,64 @@ namespace diskann {
     return ret_val;
   }
   template<typename T>
-  void PQFlashIndex<T>::parse_label_file(const std::string &map_file) {
+  void PQFlashIndex<T>::parse_label_file(const std::string &label_file) {
     //_filtered_ann = 1;
 
-    std::ifstream infile(map_file);
+    std::ifstream infile(label_file);
+    
     std::string   line, token;
     _u32          line_cnt = 0;
 
     _u32 num_pts_in_label_file;
     _u32 num_total_labels;
-    get_label_file_metadata(map_file, num_pts_in_label_file, num_total_labels);
-
+    get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels);
+    assert(num_pts_in_label_file == this->num_points);
+    
     _pts_to_label_offsets = new _u32[num_pts_in_label_file];
     _pts_to_labels = new _u32[num_pts_in_label_file + num_total_labels];
     _u32 counter = 0;
-
-    while (std::getline(infile, line)) {
-      std::istringstream iss(line);
-      std::vector<_u32>  lbls(0);
-
-      _pts_to_label_offsets[line_cnt] = counter;
-      _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter];
-      num_lbls_in_cur_pt = 0;
-      counter++;
-      //getline(iss, token, '\t');  // first token contains metadata, not used
-      getline(iss, token, '\t');
-      std::istringstream new_iss(token);
-      while (getline(new_iss, token, ',')) {
-        token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-        token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-        if (_labels.find(token) == _labels.end()) {
-          _filter_list.emplace_back(token);
+    infile.clear();
+    infile.seekg(0, std::ios::beg);
+    assert(infile.is_open());
+    try {
+      while (std::getline(infile, line)) {
+        std::istringstream iss(line);
+        std::vector<_u32>  lbls(0);
+
+        _pts_to_label_offsets[line_cnt] = counter;
+        _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter];
+        num_lbls_in_cur_pt = 0;
+        counter++;
+        getline(iss, token, '\t');
+        std::istringstream new_iss(token);
+        while (getline(new_iss, token, ',')) {
+          token.erase(std::remove(token.begin(), token.end(), '\n'),
+                      token.end());
+          token.erase(std::remove(token.begin(), token.end(), '\r'),
+                      token.end());
+          if (_labels.find(token) == _labels.end()) {
+            _filter_list.emplace_back(token);
+          }
+          int32_t filter_num = get_filter_number(token);
+          if (filter_num == -1) {
+            diskann::cout << "Error!! " << std::endl;
+            exit(-1);
+          }
+          _pts_to_labels[counter++] = filter_num;
+          num_lbls_in_cur_pt++;
+          _labels.insert(token);
         }
-        int32_t filter_num = get_filter_number(token);
-        if (filter_num == -1) {
-          std::cout << "Error!! " << std::endl;
+
+        if (num_lbls_in_cur_pt == 0) {
+          diskann::cout << "No label found for point " << line_cnt << std::endl;
           exit(-1);
         }
-        _pts_to_labels[counter++] = filter_num;
-        num_lbls_in_cur_pt++;
-        _labels.insert(token);
+        line_cnt++;
       }
-
-      if (num_lbls_in_cur_pt == 0) {
-        std::cout << "No label found for point " << line_cnt << std::endl;
-        exit(-1);
-      }
-      line_cnt++;
+    }catch (std::system_error &e) {
+      throw FileException(label_file, e, __FUNCSIG__, __FILE__,
+                          __LINE__);
     }
-    // MallocExtension::instance()->ReleaseFreeMemory();
     infile.close();
   }
 
@@ -540,7 +549,7 @@ namespace diskann {
   void PQFlashIndex<T>::set_universal_label(const std::string &label) {
     int32_t temp_filter_num = get_filter_number(label);
     if (temp_filter_num == -1) {
-      std::cout << "Error, could not find universal label. Exitting."
+      diskann::cout << "Error, could not find universal label. Exitting."
                 << std::endl;
       exit(-1);
     } else {
@@ -596,29 +605,70 @@ namespace diskann {
         std ::string(disk_index_file) + "_labels_to_medoids.txt";
     std::string dummy_map_file =
         std ::string(disk_index_file) + "_dummy_map.txt";
+
+    size_t pq_file_dim, pq_file_num_centroids;
+#ifdef EXEC_ENV_OLS
+    get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim,
+                     METADATA_SIZE);
+#else
+    get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim,
+                     METADATA_SIZE);
+#endif
+
+    this->disk_index_file = disk_index_file;
+
+    if (pq_file_num_centroids != 256) {
+      diskann::cout << "Error. Number of PQ centroids is not 256. Exiting."
+                    << std::endl;
+      return -1;
+    }
+
+    this->data_dim = pq_file_dim;
+    // will reset later if we use PQ on disk
+    this->disk_data_dim = this->data_dim;
+    // will change later if we use PQ on disk or if we are using
+    // inner product without PQ
+    this->disk_bytes_per_point = this->data_dim * sizeof(T);
+    this->aligned_dim = ROUND_UP(pq_file_dim, 8);
+
+    size_t npts_u64, nchunks_u64;
+#ifdef EXEC_ENV_OLS
+    diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64,
+                           nchunks_u64);
+#else
+    diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64,
+                           nchunks_u64);
+#endif
+
+    this->num_points = npts_u64;
+    this->n_chunks = nchunks_u64;
+
     if (file_exists(labels_file)) {
       parse_label_file(labels_file);
       if (file_exists(labels_to_medoids)) {
         std::ifstream medoid_stream(labels_to_medoids);
-
+        assert(medoid_stream.is_open());
         std::string line, token;
-        //        unsigned    line_cnt = 0;
-
         _filter_to_medoid_id.clear();
 
-        while (std::getline(medoid_stream, line)) {
-          std::istringstream iss(line);
-          _u32               cnt = 0;
-          _u32               medoid;
-          std::string        label;
-          while (std::getline(iss, token, ',')) {
-            if (cnt == 0)
-              label = token;
-            else
-              medoid = (_u32) stoul(token);
-            cnt++;
+        try {
+          while (std::getline(medoid_stream, line)) {
+            std::istringstream iss(line);
+            _u32               cnt = 0;
+            _u32               medoid;
+            std::string        label;
+            while (std::getline(iss, token, ',')) {
+              if (cnt == 0)
+                label = token;
+              else
+                medoid = (_u32) stoul(token);
+              cnt++;
+            }
+            _filter_to_medoid_id[label] = medoid;
           }
-          _filter_to_medoid_id[label] = medoid;
+        } catch (std::system_error &e) {
+          throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__,
+                              __LINE__);
         }
       }
       std::string univ_label_file =
@@ -634,8 +684,6 @@ namespace diskann {
         std::ifstream dummy_map_stream(dummy_map_file);
 
         std::string line, token;
-        //        unsigned    line_cnt = 0;
-
         while (std::getline(dummy_map_stream, line)) {
           std::istringstream iss(line);
           _u32               cnt = 0;
@@ -658,47 +706,10 @@ namespace diskann {
           _real_to_dummy_map[real_id].emplace_back(dummy_id);
         }
         dummy_map_stream.close();
-        std::cout << "Loaded dummy map" << std::endl;
+        diskann::cout << "Loaded dummy map" << std::endl;
       }
     }
 
-    size_t pq_file_dim, pq_file_num_centroids;
-#ifdef EXEC_ENV_OLS
-    get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim,
-                     METADATA_SIZE);
-#else
-    get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim,
-                     METADATA_SIZE);
-#endif
-
-    this->disk_index_file = disk_index_file;
-
-    if (pq_file_num_centroids != 256) {
-      diskann::cout << "Error. Number of PQ centroids is not 256. Exiting."
-                    << std::endl;
-      return -1;
-    }
-
-    this->data_dim = pq_file_dim;
-    // will reset later if we use PQ on disk
-    this->disk_data_dim = this->data_dim;
-    // will change later if we use PQ on disk or if we are using
-    // inner product without PQ
-    this->disk_bytes_per_point = this->data_dim * sizeof(T);
-    this->aligned_dim = ROUND_UP(pq_file_dim, 8);
-
-    size_t npts_u64, nchunks_u64;
-#ifdef EXEC_ENV_OLS
-    diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64,
-                           nchunks_u64);
-#else
-    diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64,
-                           nchunks_u64);
-#endif
-
-    this->num_points = npts_u64;
-    this->n_chunks = nchunks_u64;
-
 #ifdef EXEC_ENV_OLS
     pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64);
 #else
@@ -736,7 +747,7 @@ namespace diskann {
       disk_bytes_per_point =
           disk_pq_n_chunks *
           sizeof(_u8);  // revising disk_bytes_per_point since DISK PQ is used.
-      std::cout << "Disk index uses PQ data compressed down to "
+      diskann::cout << "Disk index uses PQ data compressed down to "
                 << disk_pq_n_chunks << " bytes per point." << std::endl;
     }
 
diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp
index a765669ef..2825ac1cc 100644
--- a/tests/search_disk_index.cpp
+++ b/tests/search_disk_index.cpp
@@ -425,4 +425,5 @@ int main(int argc, char** argv) {
     diskann::cerr << "Index search failed." << std::endl;
     return -1;
   }
+
 }
\ No newline at end of file
diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp
index c0e3c094a..153662988 100644
--- a/tests/search_memory_index.cpp
+++ b/tests/search_memory_index.cpp
@@ -357,4 +357,5 @@ int main(int argc, char** argv) {
     diskann::cerr << "Index search failed." << std::endl;
     return -1;
   }
+
 }
\ No newline at end of file
diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp
index d57eb70c3..5de1e7351 100644
--- a/tests/utils/compute_groundtruth.cpp
+++ b/tests/utils/compute_groundtruth.cpp
@@ -327,14 +327,11 @@ inline std::vector<size_t> load_filtered_bin_as_float(
   std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl;
   reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t),
                std::ios::beg);
-  //    data = new T[nptsuint64_t * ndimsuint64_t];
   T *data_T = new T[nptsuint64_t * ndimsuint64_t];
   reader.read((char *) data_T, sizeof(T) * nptsuint64_t * ndimsuint64_t);
   std::cout << "Finished reading part of the bin file." << std::endl;
   reader.close();
-  //  data =  (nptsuint64_t*ndimsuint64_t, ALIGNMENT);
   data = aligned_malloc<float>(nptsuint64_t * ndimsuint64_t, ALIGNMENT);
-  //#pragma omp parallel for schedule(dynamic, 32768)
   for (int64_t i = 0; i < (int64_t) nptsuint64_t; i++) {
     if (std::find(pts_to_labels[start_id + i].begin(),
                   pts_to_labels[start_id + i].end(),
diff --git a/tests/utils/stats_label_data.cpp b/tests/utils/stats_label_data.cpp
index 3287179c3..b983b02d3 100644
--- a/tests/utils/stats_label_data.cpp
+++ b/tests/utils/stats_label_data.cpp
@@ -1,8 +1,5 @@
-//
-// Created by Neel Karia on 10/01/2021.
-//
-// Generate base file and groundtruth file for points belonging to a particular
-// label
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
 
 #include <iostream>
 #include <fstream>
@@ -33,7 +30,6 @@ void stats_analysis(const std::string labels_file, std::string univeral_label,
                     _u32 density = 10) {
   std::string   token, line;
   std::ifstream labels_stream(labels_file);
-  // int           cnt[51] = {0};
   std::unordered_map<std::string, _u32> label_counts;
   std::string                           label_with_max_points;
   _u32                                  max_points = 0;
@@ -68,20 +64,11 @@ void stats_analysis(const std::string labels_file, std::string univeral_label,
             << (float) dense_pts / (float) labels_per_point.size() << std::endl;
   std::sort(labels_per_point.begin(), labels_per_point.end());
 
-  /*  for (_u32 p = 90; p < 100; p += 1) {
-      std::cout << labels_per_point[(p * labels_per_point.size()) / 100.0]
-                << std::endl;
-    }
-    std::cout << labels_per_point[labels_per_point.size() - 1] << std::endl;
-  */
-
   std::vector<std::pair<std::string, _u32>> label_count_vec;
 
   for (auto it = label_counts.begin(); it != label_counts.end(); it++) {
     auto& lbl = *it;
     label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second));
-    //    std::cout << "Label " << lbl.first << " has " << lbl.second
-    //              << " data points" << std::endl;
     if (lbl.second > max_points) {
       max_points = lbl.second;
       label_with_max_points = lbl.first;

From e450b7c26b1ee7a32df340b5710608e4e8f9df90 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Thu, 25 Aug 2022 18:01:54 +0530
Subject: [PATCH 03/13] Add initial implmentation of stitched_vamana

---
 tests/CMakeLists.txt           |   3 +
 tests/build_stitched_index.cpp | 596 +++++++++++++++++++++++++++++++++
 2 files changed, 599 insertions(+)
 create mode 100644 tests/build_stitched_index.cpp

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index c34be5690..6aa5532ef 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -6,6 +6,9 @@ set(CMAKE_CXX_STANDARD 14)
 add_executable(build_memory_index build_memory_index.cpp)
 target_link_libraries(build_memory_index ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options)
 
+add_executable(build_stitched_index build_stitched_index.cpp)
+target_link_libraries(build_stitched_index ${PROJECT_NAME} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options)
+
 add_executable(search_memory_index search_memory_index.cpp)
 target_link_libraries(search_memory_index ${PROJECT_NAME} ${DISKANN_ASYNC_LIB} ${DISKANN_TOOLS_TCMALLOC_LINK_OPTIONS} Boost::program_options)
 
diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
new file mode 100644
index 000000000..64cc244a5
--- /dev/null
+++ b/tests/build_stitched_index.cpp
@@ -0,0 +1,596 @@
+// authors: Siddharth Gollapudi, Varun Sivashankar
+// emails: t-gollapudis@microsoft.com, t-varunsi@microsoft.com
+
+#include <boost/program_options.hpp>
+#include <bits/types/struct_iovec.h>
+#include <fcntl.h>
+#include <omp.h>
+#include <sys/mman.h>
+#include <sys/uio.h>
+#include <unistd.h>
+#include <chrono>
+#include <cstring>
+#include <tuple>
+#include <random>
+
+#include "index.h"
+#include "utils.h"
+
+
+namespace po = boost::program_options;
+
+// oh god
+#define   LIKELY(condition) __builtin_expect(static_cast<bool>(condition), 1)
+#define UNLIKELY(condition) __builtin_expect(static_cast<bool>(condition), 0)
+
+// custom types (for readability)
+typedef std::string label;
+typedef tsl::robin_set<label> label_set;
+typedef std::string path;
+
+// structs for returning multiple items from a function
+typedef std::tuple<std::vector<label_set>, tsl::robin_map<label, _u32>, label_set> parse_label_file_return_values;
+typedef std::tuple<std::vector<std::vector<_u32>>, _u64> load_label_index_return_values;
+typedef std::tuple<std::vector<std::vector<_u32>>, _u64> stitch_indices_return_values;
+
+
+/*
+ * Inline function to generate a random integer in a range.
+ */
+inline size_t random(size_t range_from, size_t range_to) {
+  std::random_device                    rand_dev;
+  std::mt19937                          generator(rand_dev());
+  std::uniform_int_distribution<size_t> distr(range_from, range_to);
+  return distr(generator);
+}
+
+/*
+ * Inline function to handle command line parsing.
+ *
+ * Arguments are merely the inputs from the command line.
+ */
+inline size_t handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, 
+                          path &final_index_path_prefix, path &label_data_path, std::string &universal_label, 
+                          unsigned &num_threads, unsigned &R, unsigned &L, unsigned &stitched_R, float &alpha) {
+  po::options_description desc{"Arguments"};
+  try {
+    desc.add_options()("help,h", "Print information on arguments");
+    desc.add_options()("data_type",
+                       po::value<std::string>(&data_type)->required(),
+                       "data type <int8/uint8/float>");
+    desc.add_options()("data_path",
+                       po::value<path>(&input_data_path)->required(),
+                       "Input data file in bin format");
+    desc.add_options()("index_path_prefix",
+                       po::value<path>(&final_index_path_prefix)->required(),
+                       "Path prefix for saving index file components");
+    desc.add_options()("max_degree,R",
+                       po::value<uint32_t>(&R)->default_value(64),
+                       "Maximum graph degree");
+    desc.add_options()(
+        "Lbuild,L", po::value<uint32_t>(&L)->default_value(100),
+        "Build complexity, higher value results in better graphs");
+    desc.add_options()(
+        "stitched_R", po::value<uint32_t>(&L)->default_value(100),
+        "Degree to prune final graph down to");
+    desc.add_options()(
+        "alpha", po::value<float>(&alpha)->default_value(1.2f),
+        "alpha controls density and diameter of graph, set 1 for sparse graph, "
+        "1.2 or 1.4 for denser graphs with lower diameter");
+    desc.add_options()(
+        "num_threads,T",
+        po::value<uint32_t>(&num_threads)->default_value(omp_get_num_procs()),
+        "Number of threads used for building index (defaults to "
+        "omp_get_num_procs())");
+    desc.add_options()(
+        "label_file",
+        po::value<path>(&label_data_path)->default_value(""),
+        "Input label file in txt format if present");
+    desc.add_options()(
+        "universal_label",
+        po::value<label>(&universal_label)->default_value(""),
+        "Universal label, if using it, only in conjunction with labels_file");
+
+    po::variables_map vm;
+    po::store(po::parse_command_line(argc, argv, desc), vm);
+    if (vm.count("help")) {
+      std::cout << desc;
+      return 0;
+    }
+    po::notify(vm);
+  } catch (const std::exception& ex) {
+    std::cerr << ex.what() << '\n';
+    return -1;
+  }
+  
+  return 1;
+}
+
+/*
+ * Parses the label datafile, which has the following line format:
+ *            point_id -> \t -> comma-separated labels
+ *
+ * Returns three objects in via std::tuple:
+ * 1. the label universe as a set
+ */
+parse_label_file_return_values parse_label_file (path label_data_path, std::string universal_label) {
+  std::ifstream label_data_stream(label_data_path);
+  std::string   line, token;
+  unsigned      line_cnt = 0;
+
+  // allows us to reserve space for the points_to_labels vector
+  while (std::getline(label_data_stream, line))
+    line_cnt++;
+  label_data_stream.clear();
+  label_data_stream.seekg(0, std::ios::beg);
+
+  // values to return
+  std::vector<label_set> point_ids_to_labels(line_cnt);
+  tsl::robin_map<label, _u32> labels_to_number_of_points;
+  label_set all_labels;
+  
+  std::vector<_u32> points_with_universal_label;
+  while (std::getline(label_data_stream, line)) {
+    std::istringstream       current_point_and_labels(line);
+    label_set current_labels;
+
+    // get point id
+    getline(current_point_and_labels, token, '\t');
+    _u32 point_id = (_u32) std::stoul(token);
+
+    // parse comma separated labels
+    getline(current_point_and_labels, token, '\t');
+    std::istringstream current_labels_comma_separated(token);
+    bool current_universal_label_check = false;
+    while (getline(current_labels_comma_separated, token, ',')) {
+      token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
+      token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
+
+      // if token is empty, there's no labels for the point
+      if (token == universal_label) {
+        points_with_universal_label.push_back(point_id);
+        current_universal_label_check = true;
+      } else {
+        all_labels.insert(token);
+        current_labels.insert(token);
+        labels_to_number_of_points[token]++;
+      }
+    }
+
+    if (current_labels.size() <= 0 && !current_universal_label_check) {
+      std::cerr << "Error: " << point_id << " has no labels." << std::endl;
+      exit(-1);
+    } 
+    point_ids_to_labels[point_id] = current_labels;
+  }
+
+  // for every point with universal label, set its label set to all labels
+  // also, increment the count for number of points a label has
+  for (const auto &point_id : points_with_universal_label) {
+    point_ids_to_labels[point_id] = all_labels;
+    for (const auto &label : all_labels)
+      labels_to_number_of_points[label]++;
+  }
+
+  std::cout << "Identified " << all_labels.size() << " distinct label(s) \n"
+            << std::endl;
+
+  return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels);
+}
+
+
+/*
+ * For each label, generates a file containing all vectors that have said label.
+ * Also copies data from original bin file to new dimension-aligned file.
+ * Utilizes platform-specific functions mmap and writev.
+ *
+ * Each data file is saved under the following format:
+ *    input_data_path + "_" + label
+ */
+template<typename T>
+tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(path input_data_path, 
+                                                                              tsl::robin_map<label, _u32> labels_to_number_of_points,
+                                                                              std::vector<label_set> point_ids_to_labels, label_set all_labels) {
+  void *input_memblock;
+  int input_data_fd;
+  const size_t METADATA = 2 * sizeof(_u32);
+  struct stat input_data_st;
+
+  auto file_writing_timer = std::chrono::high_resolution_clock::now();
+  // UNLIKELY moves the branch instr. body elsewhere, opening up icache space
+  // for code that actually executes
+  input_data_fd = open(input_data_path.c_str(), O_RDONLY, (mode_t) 0444);
+  if (UNLIKELY(input_data_fd == -1))
+    throw;
+  if (UNLIKELY(fstat(input_data_fd, &input_data_st))) {
+    close(input_data_fd);
+    throw;
+  }
+
+  // mmap could be faster than just doing a regular sequential read?
+  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_PRIVATE, input_data_fd, 0);
+  if (UNLIKELY(input_memblock == MAP_FAILED)) {
+    close(input_data_fd);
+    throw;
+  }
+  const char *input_start = static_cast<char const *>(input_memblock);
+  _u32 number_of_points, dimension;
+  std::memcpy(&number_of_points, input_start, sizeof(_u32));
+  std::memcpy(&dimension, input_start + sizeof(_u32), sizeof(_u32));
+
+  // iovec necessary for using writev
+  tsl::robin_map<label, iovec> label_to_vectors_map;
+
+  // mapping from id in label file to original file
+  tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id;
+  for (_u32 point_id = 0; point_id < number_of_points; point_id++) {
+    for (const auto &label : point_ids_to_labels[point_id]) {
+      // first add metadata before adding vectors
+      if (!label_to_vectors_map.count(label)) {
+        iovec curr_iovec;
+        curr_iovec.iov_base = malloc(METADATA + labels_to_number_of_points[label] * dimension * sizeof(T)); 
+        std::memcpy((char *) curr_iovec.iov_base, &labels_to_number_of_points[label], sizeof(_u32));
+        curr_iovec.iov_len = sizeof(_u32);
+        std::memcpy((char *) curr_iovec.iov_base + curr_iovec.iov_len, &dimension, sizeof(_u32));
+        curr_iovec.iov_len += sizeof(_u32);
+        label_to_vectors_map[label] = curr_iovec;
+        label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);
+      }
+      char *current_iovec_buffer = (char *) label_to_vectors_map[label].iov_base;
+      size_t *current_iovec_buffer_size = &label_to_vectors_map[label].iov_len;
+      std::memcpy(current_iovec_buffer + *current_iovec_buffer_size, input_start + METADATA + dimension * point_id, sizeof(T) * dimension);
+      *current_iovec_buffer_size += sizeof(T) * dimension;
+      label_id_to_orig_id[label].push_back(point_id);
+    }
+    
+  }
+
+  if (UNLIKELY(munmap(input_memblock, input_data_st.st_size) == -1)) {
+    close(input_data_fd);
+    throw;
+  }
+  close(input_data_fd);
+
+  // write each label iovec to resp. file
+  for (const auto &label : all_labels) {
+    int label_input_data_fd;
+    std::string curr_label_input_data_path(input_data_path + "_" + label);
+
+    label_input_data_fd = open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY, (mode_t) 0644);
+    if (UNLIKELY(label_input_data_fd == -1))
+      throw;
+
+    // using writev leads to one write syscall per label file
+    int return_value = writev(label_input_data_fd, &label_to_vectors_map[label], 1);
+    if (UNLIKELY(return_value == -1)) {
+      close(label_input_data_fd);
+      throw;
+    }
+
+    free(label_to_vectors_map[label].iov_base);
+    close(label_input_data_fd);
+  }
+
+  std::chrono::duration<double> file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer;
+  std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time "
+    << file_writing_time.count() << "\n" << std::endl;
+
+  return label_id_to_orig_id;
+}
+
+
+/*
+ * Using passed in parameters and files generated from step 3, 
+ * builds a vanilla diskANN index for each label.
+ *
+ * Each index is saved under the following path:
+ *  final_index_path_prefix + "_" + label
+ */
+template <typename T>
+void generate_label_indices(path input_data_path, path final_index_path_prefix,
+                            label_set all_labels, unsigned R, unsigned L, float alpha, 
+                            unsigned num_threads) {
+  diskann::Parameters label_index_build_parameters;
+  label_index_build_parameters.Set<unsigned>("R", R);
+  label_index_build_parameters.Set<unsigned>("L", L);
+  label_index_build_parameters.Set<unsigned>("C", 750);  // maximum candidate set size du
+  label_index_build_parameters.Set<bool>("saturate_graph", 0);
+  label_index_build_parameters.Set<float>("alpha", alpha);
+  label_index_build_parameters.Set<unsigned>("num_threads", num_threads);
+
+  // for each label, build an index on resp. points
+  double total_indexing_time = 0.0;
+  auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr);
+  auto std_cout_buffer = std::cout.rdbuf(nullptr);
+  for (const auto &label : all_labels) {
+    path curr_label_input_data_path(input_data_path + "_" + label);
+    path curr_label_index_path(final_index_path_prefix + "_" + label);
+
+    size_t number_of_label_points, dimension;
+    diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension);
+    diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
+
+    auto index_build_timer = std::chrono::high_resolution_clock::now();
+    index.build(curr_label_input_data_path.c_str(), number_of_label_points, label_index_build_parameters);
+    std::chrono::duration<double> current_indexing_time =
+      std::chrono::high_resolution_clock::now() - index_build_timer;
+  
+    total_indexing_time += current_indexing_time.count();
+    index.save(curr_label_index_path.c_str());
+  }
+  diskann::cout.rdbuf(diskann_cout_buffer);
+  std::cout.rdbuf(std_cout_buffer);
+
+  std::cout << "generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl;
+}
+
+
+/*
+ * Manually loads a graph index in from a given file.
+ * 
+ * Returns both the graph index and the size of the file in bytes.
+ */
+load_label_index_return_values load_label_index(path label_index_path, _u32 label_number_of_points) {
+  std::ifstream label_index_stream;
+  label_index_stream.exceptions(std::ios::badbit | std::ios::failbit);
+  label_index_stream.open(label_index_path, std::ios::binary);
+
+  _u64 index_file_size, index_num_frozen_points;
+  _u32 index_max_observed_degree, index_entry_point;
+  const size_t INDEX_METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
+  label_index_stream.read((char *) &index_file_size, sizeof(_u64));
+  label_index_stream.read((char *) &index_max_observed_degree, sizeof(_u32));
+  label_index_stream.read((char *) &index_entry_point, sizeof(_u32));
+  label_index_stream.read((char *) &index_num_frozen_points, sizeof(_u64));
+  size_t bytes_read = INDEX_METADATA;
+
+  std::vector<std::vector<_u32>> label_index(label_number_of_points);
+  _u32 nodes_read = 0;
+  while (bytes_read != index_file_size) {
+    _u32 current_node_num_neighbors;
+    label_index_stream.read((char *) &current_node_num_neighbors, sizeof(_u32));
+    nodes_read++;
+    
+    std::vector<_u32> current_node_neighbors(current_node_num_neighbors);
+    label_index_stream.read((char *) current_node_neighbors.data(), current_node_num_neighbors * sizeof(_u32));
+    label_index[nodes_read - 1].swap(current_node_neighbors);
+    bytes_read += sizeof(_u32) * (current_node_num_neighbors + 1);
+  }
+  
+  return std::make_tuple(label_index, index_file_size);
+}
+
+
+void save_full_index(path final_index_path_prefix, path input_data_path,
+                     _u64 final_index_size, std::vector<std::vector<_u32>> stitched_graph,
+                     tsl::robin_map<label, _u32> entry_points,
+                     label universal_label, path label_data_path) {
+
+  std::cout << "test 1" << std::endl;
+  auto saving_index_timer = std::chrono::high_resolution_clock::now();
+  std::ifstream original_label_data_stream;
+  original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
+  original_label_data_stream.open(label_data_path, std::ios::binary);
+  std::ofstream new_label_data_stream;
+  new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
+  new_label_data_stream.open(final_index_path_prefix + "_labels.txt", std::ios::binary);
+  new_label_data_stream << original_label_data_stream.rdbuf();
+  original_label_data_stream.close();
+  new_label_data_stream.close();
+
+  std::cout << "test 1" << std::endl;
+  std::ifstream original_input_data_stream;
+  original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
+  original_input_data_stream.open(input_data_path, std::ios::binary);
+  std::ofstream new_input_data_stream;
+  new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
+  new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary);
+  new_input_data_stream << original_input_data_stream.rdbuf();
+  original_input_data_stream.close();
+  new_label_data_stream.close();
+
+  std::cout << "test 1" << std::endl;
+  std::ofstream labels_to_medoids_writer;
+  labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit);
+  labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt");
+  for (auto iter : entry_points) {
+    std::cout << iter.first << ", " << iter.second << std::endl;;
+    labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl;
+  }
+  labels_to_medoids_writer.close();
+
+  std::cout << "test 1" << std::endl;
+  if (universal_label != "") {
+    std::ofstream universal_label_writer;
+    universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit);
+    universal_label_writer.open(final_index_path_prefix + "_universal_label.txt");
+    universal_label_writer << universal_label << std::endl;
+    universal_label_writer.close();
+  }
+
+  std::cout << "test 1" << std::endl;
+  _u64 index_num_frozen_points = 0, index_num_edges = 0;
+  _u32 index_max_observed_degree = 0, index_entry_point = 0;
+  for (auto& point_neighbors : stitched_graph) {
+    index_max_observed_degree = std::max(index_max_observed_degree, (_u32) point_neighbors.size());
+  }
+
+  void *memblock;
+  const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
+  struct stat sb;
+  int stitched_graph_fd;
+
+  stitched_graph_fd = open(final_index_path_prefix.c_str(), O_RDWR | O_CREAT, (mode_t) 0644);
+  if (UNLIKELY(stitched_graph_fd == -1))
+    throw;
+  if (UNLIKELY(fstat(stitched_graph_fd, &sb))) {
+    close(stitched_graph_fd);
+    throw;
+  }
+
+  if (UNLIKELY(ftruncate(stitched_graph_fd, final_index_size) == -1)) {
+    close(stitched_graph_fd);
+    throw;
+  }
+
+  // mmap could be faster than just doing a regular sequential write?
+  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_PRIVATE, stitched_graph_fd, 0);
+  if (UNLIKELY(memblock == MAP_FAILED)) {
+    close(stitched_graph_fd);
+    throw;
+  }
+  
+  std::cout << "test 1" << std::endl;
+  char *begin = static_cast<char *>(memblock);
+  std::memcpy(begin, &final_index_size, sizeof(_u64));
+  std::memcpy(begin + sizeof(_u64), &index_max_observed_degree, sizeof(_u32));
+  std::memcpy(begin + sizeof(_u64) + sizeof(_u32), &index_entry_point, sizeof(_u32));
+  std::memcpy(begin + sizeof(_u64) + 2 * sizeof(_u32), &index_num_frozen_points, sizeof(_u64));
+  size_t bytes_written = METADATA;
+  for (_u32 node_point = 0; node_point < stitched_graph.size(); node_point++) {
+    _u32 current_node_num_neighbors = stitched_graph[node_point].size();
+    std::memcpy(begin + bytes_written, &current_node_num_neighbors, sizeof(_u32));
+    bytes_written += sizeof(_u32);
+    
+    for (const auto &current_node_neighbor : stitched_graph[node_point]) {
+      std::memcpy(begin + bytes_written, &current_node_neighbor, sizeof(_u32));
+      bytes_written += sizeof(_u32);
+    }
+    index_num_edges += current_node_num_neighbors;
+  }
+
+  if (UNLIKELY(bytes_written != final_index_size)) {
+    std::cerr << "Error: written bytes does not match allocated space" << std::endl;
+    throw;
+  }
+
+  if (UNLIKELY(msync(memblock, final_index_size, MS_SYNC) == -1)) {
+    std::cerr << "Error: could not sync memory-mapped writes to disk" << std::endl;
+    close(stitched_graph_fd);
+    throw;
+  }
+
+  if (UNLIKELY(munmap(memblock, final_index_size) == -1)) {
+    std::cerr << "Error: could not release the memory-mapped file" << std::endl;
+    close(stitched_graph_fd);
+    throw;
+  }
+
+  close(stitched_graph_fd);
+
+  std::chrono::duration<double> saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer;
+  std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl;
+  std::cout << "Average degree: " << ((float) index_num_edges) / ((float) (stitched_graph.size())) << std::endl;
+  std::cout << "Max degree: " << index_max_observed_degree << std::endl;
+}
+
+
+/*
+ * Unions the per-label graph indices together via the following policy:
+ *  - any two nodes can only have at most one edge between them -
+ * 
+ * Returns the "stitched" graph and its expected file size.
+ */
+template <typename T>
+stitch_indices_return_values stitch_label_indices(path final_index_path_prefix, _u32 total_number_of_points, label_set all_labels, 
+                          tsl::robin_map<label, _u32> labels_to_number_of_points, 
+                          tsl::robin_map<label, _u32> &label_entry_points,
+                          tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map) {
+  size_t final_index_size = 0;
+  std::vector<std::vector<_u32>> stitched_graph(total_number_of_points);
+  
+  
+  auto stitching_index_timer = std::chrono::high_resolution_clock::now();
+  for (const auto &label : all_labels) {
+    path curr_label_index_path(final_index_path_prefix + "_" + label);
+    std::vector<std::vector<_u32>> curr_label_index;
+    _u64 curr_label_index_size;
+    _u32 curr_label_entry_point;
+    std::tie(curr_label_index, curr_label_index_size) = load_label_index(curr_label_index_path, labels_to_number_of_points[label]);
+    final_index_size += curr_label_index_size;
+    curr_label_entry_point = random(0, curr_label_index.size());
+    label_entry_points[label] = label_id_to_orig_id_map[label][curr_label_entry_point];
+  
+    for (_u32 node_point = 0; node_point < curr_label_index.size(); node_point++) {
+      _u32 original_point_id = label_id_to_orig_id_map[label][node_point];
+      for (auto &node_neighbor : curr_label_index[node_point]) {
+        _u32 original_neighbor_id = label_id_to_orig_id_map[label][node_neighbor];
+        std::vector<_u32> curr_point_neighbors = stitched_graph[original_point_id];
+        if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == curr_point_neighbors.end())
+          stitched_graph[original_point_id].push_back(original_neighbor_id);
+        else
+          final_index_size -= sizeof(T);
+      }
+    }
+  }
+
+  size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
+  final_index_size -= all_labels.size() * METADATA;
+  final_index_size += METADATA;
+
+  std::chrono::duration<double> stitching_index_time = std::chrono::high_resolution_clock::now() - stitching_index_timer;
+  std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl;
+
+  return std::make_tuple(stitched_graph, final_index_size);
+}
+
+
+int main(int argc, char **argv) {
+  // 1. handle cmdline inputs
+  std::string data_type;
+  path input_data_path, final_index_path_prefix, label_data_path;
+  label universal_label;
+  unsigned num_threads, R, L, stitched_R;
+  float alpha;
+
+  handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, num_threads, R, L, stitched_R, alpha);
+
+  // 2. parse label file and create necessary data structures
+  std::vector<label_set> point_ids_to_labels;
+  tsl::robin_map<label, _u32> labels_to_number_of_points;
+  label_set all_labels;
+  
+  std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = parse_label_file(label_data_path, universal_label);
+
+  // 3. for each label, make a separate data file
+  tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map;
+  _u32 total_number_of_points = point_ids_to_labels.size();
+  
+  if (data_type == "uint8")
+    label_id_to_orig_id_map = generate_label_specific_vector_files<uint8_t>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+  else if (data_type == "int8")
+    label_id_to_orig_id_map = generate_label_specific_vector_files<int8_t>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+  else if (data_type == "float")
+    label_id_to_orig_id_map = generate_label_specific_vector_files<float>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+  else
+    throw;
+
+  // 4. for each created data file, create a vanilla diskANN index
+  if (data_type == "uint8")
+    generate_label_indices<uint8_t>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+  else if (data_type == "int8")
+    generate_label_indices<int8_t>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+  else if (data_type == "float")
+    generate_label_indices<float>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+  else
+   throw;
+
+  // 5. "stitch" the indices together
+  std::vector<std::vector<_u32>> stitched_graph;
+  tsl::robin_map<label, _u32> label_entry_points;
+  _u64 stitched_graph_size;
+
+  if (data_type == "uint8")
+    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<uint8_t>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+  else if (data_type == "int8")
+    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<int8_t>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+  else if (data_type == "float")
+    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<float>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+  else
+   throw;
+
+  // 5a. save the stitched graph to disk
+  save_full_index(final_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, universal_label, label_data_path);
+
+  // TODO: 6. run a prune on the stitched index, and save to disk
+}

From 8f6059a16d16e308579f93b0cb98161b5325d48c Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Mon, 29 Aug 2022 19:18:12 +0530
Subject: [PATCH 04/13] Add bugfixes for major disk IO and graph stitching,
 plus cleanup for intermediate file artifacts

---
 include/index.h                |   2 +
 src/index.cpp                  |  78 ++++++++++--
 tests/build_stitched_index.cpp | 215 ++++++++++++++++++++++++---------
 3 files changed, 231 insertions(+), 64 deletions(-)

diff --git a/include/index.h b/include/index.h
index 7ea309e34..33d6bc240 100644
--- a/include/index.h
+++ b/include/index.h
@@ -192,6 +192,8 @@ namespace diskann {
     DISKANN_DLLEXPORT consolidation_report
     consolidate_deletes(const Parameters &parameters);
 
+		DISKANN_DLLEXPORT void prune_all_nbrs(const Parameters &parameters);
+
     DISKANN_DLLEXPORT bool is_index_saved();
 
     // repositions frozen points to the end of _data - if they have been moved
diff --git a/src/index.cpp b/src/index.cpp
index 35eab43c3..957cb6643 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -483,6 +483,14 @@ namespace diskann {
     std::string labels_file = mem_index_file + "_labels.txt";
     std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt";
     if (file_exists(labels_file)) {
+      std::string universal_label_file(filename);
+      universal_label_file += "_universal_label.txt";
+      if (file_exists(universal_label_file)) {
+        std::ifstream universal_label_reader(universal_label_file);
+        universal_label_reader >> _universal_label;
+        _use_universal_label = true;
+        universal_label_reader.close();
+      }
       parse_label_file(labels_file);
       if (file_exists(labels_to_medoids)) {
         std::ifstream medoid_stream(labels_to_medoids);
@@ -514,14 +522,7 @@ namespace diskann {
         }
       }
 
-      std::string universal_label_file(filename);
-      universal_label_file += "_universal_label.txt";
-      if (file_exists(universal_label_file)) {
-        std::ifstream universal_label_reader(universal_label_file);
-        universal_label_reader >> _universal_label;
-        _use_universal_label = true;
-        universal_label_reader.close();
-      }
+      
     }
 
     if (!_save_as_one_file) {
@@ -688,13 +689,13 @@ namespace diskann {
       }
     }
 #else
-
     size_t   bytes_read = vamana_metadata_size;
     size_t   cc = 0;
     unsigned nodes_read = 0;
     while (bytes_read != expected_file_size) {
       unsigned k;
       in.read((char *) &k, sizeof(unsigned));
+
       if (k == 0) {
         diskann::cerr << "ERROR: Point found with no out-neighbors, point#"
                       << nodes_read << std::endl;
@@ -1378,6 +1379,65 @@ namespace diskann {
     }
   }
 
+  template<typename T, typename TagT>
+  void Index<T, TagT>::prune_all_nbrs(const Parameters &parameters) {
+    const unsigned range = parameters.Get<unsigned>("R");
+    const unsigned maxc = parameters.Get<unsigned>("C");
+    const float alpha = parameters.Get<unsigned>("alpha");
+    _filtered_index = true;
+
+    diskann::Timer timer;
+#pragma omp parallel for
+    for (_s64 node = 0; node < (_s64)(_max_points + _num_frozen_pts); node++) {
+      if ((size_t) node < _nd || (size_t) node == _max_points) {
+        if (_final_graph[node].size() > range) {
+          tsl::robin_set<unsigned> dummy_visited(0);
+          std::vector<Neighbor>    dummy_pool(0);
+          std::vector<unsigned>    new_out_neighbors;
+
+					ScratchStoreManager<InMemQueryScratch<T>> manager(_query_scratch);
+      		auto scratch = manager.scratch_space();
+
+          for (auto cur_nbr : _final_graph[node]) {
+            if (dummy_visited.find(cur_nbr) == dummy_visited.end() &&
+                cur_nbr != node) {
+              float dist =
+                  _distance->compare(_data + _aligned_dim * (size_t) node,
+                                     _data + _aligned_dim * (size_t) cur_nbr,
+                                     (unsigned) _aligned_dim);
+              dummy_pool.emplace_back(Neighbor(cur_nbr, dist, true));
+              dummy_visited.insert(cur_nbr);
+            }
+          }
+          
+          prune_neighbors((_u32) node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch);
+          _final_graph[node].clear();
+          for (auto id : new_out_neighbors)
+            _final_graph[node].emplace_back(id);
+        }
+      }
+    }
+
+    diskann::cout << "Prune time : " << timer.elapsed() / 1000 << "ms"
+                  << std::endl;
+    size_t max = 0, min = 1 << 30, total = 0, cnt = 0;
+    for (size_t i = 0; i < (_nd + _num_frozen_pts); i++) {
+      std::vector<unsigned> pool = _final_graph[i];
+      max = (std::max)(max, pool.size());
+      min = (std::min)(min, pool.size());
+      total += pool.size();
+      if (pool.size() < 2)
+        cnt++;
+    }
+    if (min > max)
+      min = max;
+    if (_nd > 0) {
+      diskann::cout << "Index built with degree: max:" << max << "  avg:"
+                    << (float) total / (float) (_nd + _num_frozen_pts)
+                    << "  min:" << min << "  count(deg<2):" << cnt << std::endl;
+    }
+  }
+
   template<typename T, typename TagT>
   void Index<T, TagT>::set_start_point(T *data) {
     std::unique_lock<std::shared_timed_mutex> ul(_update_lock);
diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 64cc244a5..34f00b605 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -9,19 +9,23 @@
 #include <sys/uio.h>
 #include <unistd.h>
 #include <chrono>
+#include <cstdio>
 #include <cstring>
 #include <tuple>
 #include <random>
 
 #include "index.h"
+#include "parameters.h"
 #include "utils.h"
 
 
 namespace po = boost::program_options;
 
-// oh god
+// macros (posix reasons)
 #define   LIKELY(condition) __builtin_expect(static_cast<bool>(condition), 1)
 #define UNLIKELY(condition) __builtin_expect(static_cast<bool>(condition), 0)
+#define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"
+#define PBWIDTH 60
 
 // custom types (for readability)
 typedef std::string label;
@@ -34,6 +38,18 @@ typedef std::tuple<std::vector<std::vector<_u32>>, _u64> load_label_index_return
 typedef std::tuple<std::vector<std::vector<_u32>>, _u64> stitch_indices_return_values;
 
 
+
+/*
+ * Inline function to display progress bar.
+ */
+inline void print_progress(double percentage) {
+    int val = (int) (percentage * 100);
+    int lpad = (int) (percentage * PBWIDTH);
+    int rpad = PBWIDTH - lpad;
+    printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, "");
+    fflush(stdout);
+}
+
 /*
  * Inline function to generate a random integer in a range.
  */
@@ -45,11 +61,11 @@ inline size_t random(size_t range_from, size_t range_to) {
 }
 
 /*
- * Inline function to handle command line parsing.
+ * function to handle command line parsing.
  *
  * Arguments are merely the inputs from the command line.
  */
-inline size_t handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, 
+size_t handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, 
                           path &final_index_path_prefix, path &label_data_path, std::string &universal_label, 
                           unsigned &num_threads, unsigned &R, unsigned &L, unsigned &stitched_R, float &alpha) {
   po::options_description desc{"Arguments"};
@@ -71,7 +87,7 @@ inline size_t handle_args(int argc, char **argv, std::string &data_type, path &i
         "Lbuild,L", po::value<uint32_t>(&L)->default_value(100),
         "Build complexity, higher value results in better graphs");
     desc.add_options()(
-        "stitched_R", po::value<uint32_t>(&L)->default_value(100),
+        "stitched_R", po::value<uint32_t>(&stitched_R)->default_value(100),
         "Degree to prune final graph down to");
     desc.add_options()(
         "alpha", po::value<float>(&alpha)->default_value(1.2f),
@@ -102,10 +118,10 @@ inline size_t handle_args(int argc, char **argv, std::string &data_type, path &i
     std::cerr << ex.what() << '\n';
     return -1;
   }
-  
   return 1;
 }
 
+
 /*
  * Parses the label datafile, which has the following line format:
  *            point_id -> \t -> comma-separated labels
@@ -172,7 +188,7 @@ parse_label_file_return_values parse_label_file (path label_data_path, std::stri
       labels_to_number_of_points[label]++;
   }
 
-  std::cout << "Identified " << all_labels.size() << " distinct label(s) \n"
+  std::cout << "Identified " << all_labels.size() << " distinct label(s) for " << point_ids_to_labels.size() << " points\n"
             << std::endl;
 
   return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels);
@@ -193,10 +209,10 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
                                                                               std::vector<label_set> point_ids_to_labels, label_set all_labels) {
   void *input_memblock;
   int input_data_fd;
-  const size_t METADATA = 2 * sizeof(_u32);
   struct stat input_data_st;
-
+  const size_t METADATA = 2 * sizeof(_u32);
   auto file_writing_timer = std::chrono::high_resolution_clock::now();
+
   // UNLIKELY moves the branch instr. body elsewhere, opening up icache space
   // for code that actually executes
   input_data_fd = open(input_data_path.c_str(), O_RDONLY, (mode_t) 0444);
@@ -208,7 +224,7 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
   }
 
   // mmap could be faster than just doing a regular sequential read?
-  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_PRIVATE, input_data_fd, 0);
+  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_SHARED, input_data_fd, 0);
   if (UNLIKELY(input_memblock == MAP_FAILED)) {
     close(input_data_fd);
     throw;
@@ -217,29 +233,35 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
   _u32 number_of_points, dimension;
   std::memcpy(&number_of_points, input_start, sizeof(_u32));
   std::memcpy(&dimension, input_start + sizeof(_u32), sizeof(_u32));
+  const _u32 VECTOR_SIZE = dimension * sizeof(T);
+  if (number_of_points != point_ids_to_labels.size()) {
+    std::cerr << "Error: number of points in labels file and data file differ." << std::endl;
+    throw;
+  }
 
-  // iovec necessary for using writev
-  tsl::robin_map<label, iovec> label_to_vectors_map;
+  // iovecs necessary for using writev
+  tsl::robin_map<label, iovec> label_to_iovec_map;
 
   // mapping from id in label file to original file
   tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id;
   for (_u32 point_id = 0; point_id < number_of_points; point_id++) {
+    const char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id);
     for (const auto &label : point_ids_to_labels[point_id]) {
       // first add metadata before adding vectors
-      if (!label_to_vectors_map.count(label)) {
+      if (!label_to_iovec_map.count(label)) {
         iovec curr_iovec;
-        curr_iovec.iov_base = malloc(METADATA + labels_to_number_of_points[label] * dimension * sizeof(T)); 
+        curr_iovec.iov_base = malloc(METADATA + labels_to_number_of_points[label] * VECTOR_SIZE); 
         std::memcpy((char *) curr_iovec.iov_base, &labels_to_number_of_points[label], sizeof(_u32));
         curr_iovec.iov_len = sizeof(_u32);
         std::memcpy((char *) curr_iovec.iov_base + curr_iovec.iov_len, &dimension, sizeof(_u32));
         curr_iovec.iov_len += sizeof(_u32);
-        label_to_vectors_map[label] = curr_iovec;
+        label_to_iovec_map[label] = curr_iovec;
         label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);
       }
-      char *current_iovec_buffer = (char *) label_to_vectors_map[label].iov_base;
-      size_t *current_iovec_buffer_size = &label_to_vectors_map[label].iov_len;
-      std::memcpy(current_iovec_buffer + *current_iovec_buffer_size, input_start + METADATA + dimension * point_id, sizeof(T) * dimension);
-      *current_iovec_buffer_size += sizeof(T) * dimension;
+      char *current_iovec_buffer = (char *) label_to_iovec_map[label].iov_base;
+      size_t *current_iovec_buffer_size = &label_to_iovec_map[label].iov_len;
+      std::memcpy(current_iovec_buffer + *current_iovec_buffer_size, curr_point, VECTOR_SIZE);
+      *current_iovec_buffer_size += VECTOR_SIZE;
       label_id_to_orig_id[label].push_back(point_id);
     }
     
@@ -254,20 +276,20 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
   // write each label iovec to resp. file
   for (const auto &label : all_labels) {
     int label_input_data_fd;
-    std::string curr_label_input_data_path(input_data_path + "_" + label);
+    path curr_label_input_data_path(input_data_path + "_" + label);
 
     label_input_data_fd = open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY, (mode_t) 0644);
     if (UNLIKELY(label_input_data_fd == -1))
       throw;
 
     // using writev leads to one write syscall per label file
-    int return_value = writev(label_input_data_fd, &label_to_vectors_map[label], 1);
+    int return_value = writev(label_input_data_fd, &label_to_iovec_map[label], 1);
     if (UNLIKELY(return_value == -1)) {
       close(label_input_data_fd);
       throw;
     }
 
-    free(label_to_vectors_map[label].iov_base);
+    free(label_to_iovec_map[label].iov_base);
     close(label_input_data_fd);
   }
 
@@ -293,21 +315,23 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
   diskann::Parameters label_index_build_parameters;
   label_index_build_parameters.Set<unsigned>("R", R);
   label_index_build_parameters.Set<unsigned>("L", L);
-  label_index_build_parameters.Set<unsigned>("C", 750);  // maximum candidate set size du
+  label_index_build_parameters.Set<unsigned>("C", 750);  // maximum candidate set size during prune
   label_index_build_parameters.Set<bool>("saturate_graph", 0);
   label_index_build_parameters.Set<float>("alpha", alpha);
   label_index_build_parameters.Set<unsigned>("num_threads", num_threads);
 
+  std::cout << "Generating indices per label..." << std::endl;
   // for each label, build an index on resp. points
-  double total_indexing_time = 0.0;
-  auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr);
-  auto std_cout_buffer = std::cout.rdbuf(nullptr);
+  double total_indexing_time = 0.0, indexing_percentage = 0.0;
+  std::cout.setstate(std::ios_base::failbit);
+  diskann::cout.setstate(std::ios_base::failbit);
   for (const auto &label : all_labels) {
     path curr_label_input_data_path(input_data_path + "_" + label);
     path curr_label_index_path(final_index_path_prefix + "_" + label);
-
+    
     size_t number_of_label_points, dimension;
     diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension);
+    std::cout << number_of_label_points << " " << dimension << std::endl;
     diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
 
     auto index_build_timer = std::chrono::high_resolution_clock::now();
@@ -316,15 +340,21 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
       std::chrono::high_resolution_clock::now() - index_build_timer;
   
     total_indexing_time += current_indexing_time.count();
+    indexing_percentage += (1 / (double) all_labels.size());
+    print_progress(indexing_percentage);
+
     index.save(curr_label_index_path.c_str());
   }
-  diskann::cout.rdbuf(diskann_cout_buffer);
-  std::cout.rdbuf(std_cout_buffer);
+  std::cout.clear();
+  diskann::cout.clear();
 
-  std::cout << "generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl;
+  std::cout << "\nDone. Generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl;
 }
 
 
+
+
+
 /*
  * Manually loads a graph index in from a given file.
  * 
@@ -356,17 +386,24 @@ load_label_index_return_values load_label_index(path label_index_path, _u32 labe
     label_index[nodes_read - 1].swap(current_node_neighbors);
     bytes_read += sizeof(_u32) * (current_node_num_neighbors + 1);
   }
-  
+
   return std::make_tuple(label_index, index_file_size);
 }
 
-
+/*
+ * Custom index save to write the in-memory index to disk.
+ * Also writes required files for diskANN API - 
+ *  1. labels_to_medoids
+ *  2. universal_label
+ *  3. data (redundant for static indices)
+ *  4. labels (redundant for static indices)
+ */
 void save_full_index(path final_index_path_prefix, path input_data_path,
                      _u64 final_index_size, std::vector<std::vector<_u32>> stitched_graph,
                      tsl::robin_map<label, _u32> entry_points,
                      label universal_label, path label_data_path) {
 
-  std::cout << "test 1" << std::endl;
+  // aux. file 1
   auto saving_index_timer = std::chrono::high_resolution_clock::now();
   std::ifstream original_label_data_stream;
   original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
@@ -378,7 +415,7 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   original_label_data_stream.close();
   new_label_data_stream.close();
 
-  std::cout << "test 1" << std::endl;
+  // aux. file 2
   std::ifstream original_input_data_stream;
   original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
   original_input_data_stream.open(input_data_path, std::ios::binary);
@@ -387,19 +424,17 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary);
   new_input_data_stream << original_input_data_stream.rdbuf();
   original_input_data_stream.close();
-  new_label_data_stream.close();
+  new_input_data_stream.close();
 
-  std::cout << "test 1" << std::endl;
+  // aux. file 3
   std::ofstream labels_to_medoids_writer;
   labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit);
   labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt");
-  for (auto iter : entry_points) {
-    std::cout << iter.first << ", " << iter.second << std::endl;;
+  for (auto iter : entry_points)
     labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl;
-  }
   labels_to_medoids_writer.close();
 
-  std::cout << "test 1" << std::endl;
+  // aux. file 4
   if (universal_label != "") {
     std::ofstream universal_label_writer;
     universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit);
@@ -408,7 +443,7 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
     universal_label_writer.close();
   }
 
-  std::cout << "test 1" << std::endl;
+  // main index
   _u64 index_num_frozen_points = 0, index_num_edges = 0;
   _u32 index_max_observed_degree = 0, index_entry_point = 0;
   for (auto& point_neighbors : stitched_graph) {
@@ -434,13 +469,12 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   }
 
   // mmap could be faster than just doing a regular sequential write?
-  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_PRIVATE, stitched_graph_fd, 0);
+  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_SHARED, stitched_graph_fd, 0);
   if (UNLIKELY(memblock == MAP_FAILED)) {
     close(stitched_graph_fd);
     throw;
   }
   
-  std::cout << "test 1" << std::endl;
   char *begin = static_cast<char *>(memblock);
   std::memcpy(begin, &final_index_size, sizeof(_u64));
   std::memcpy(begin + sizeof(_u64), &index_max_observed_degree, sizeof(_u32));
@@ -449,10 +483,10 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   size_t bytes_written = METADATA;
   for (_u32 node_point = 0; node_point < stitched_graph.size(); node_point++) {
     _u32 current_node_num_neighbors = stitched_graph[node_point].size();
+    std::vector<_u32> current_node_neighbors = stitched_graph[node_point];
     std::memcpy(begin + bytes_written, &current_node_num_neighbors, sizeof(_u32));
     bytes_written += sizeof(_u32);
-    
-    for (const auto &current_node_neighbor : stitched_graph[node_point]) {
+    for (const auto &current_node_neighbor : current_node_neighbors) {
       std::memcpy(begin + bytes_written, &current_node_neighbor, sizeof(_u32));
       bytes_written += sizeof(_u32);
     }
@@ -480,8 +514,8 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
 
   std::chrono::duration<double> saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer;
   std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl;
-  std::cout << "Average degree: " << ((float) index_num_edges) / ((float) (stitched_graph.size())) << std::endl;
-  std::cout << "Max degree: " << index_max_observed_degree << std::endl;
+  std::cout << "Stitched graph average degree: " << ((float) index_num_edges) / ((float) (stitched_graph.size())) << std::endl;
+  std::cout << "Stitched graph max degree: " << index_max_observed_degree << std::endl << std::endl;
 }
 
 
@@ -499,15 +533,14 @@ stitch_indices_return_values stitch_label_indices(path final_index_path_prefix,
   size_t final_index_size = 0;
   std::vector<std::vector<_u32>> stitched_graph(total_number_of_points);
   
-  
   auto stitching_index_timer = std::chrono::high_resolution_clock::now();
   for (const auto &label : all_labels) {
     path curr_label_index_path(final_index_path_prefix + "_" + label);
     std::vector<std::vector<_u32>> curr_label_index;
     _u64 curr_label_index_size;
     _u32 curr_label_entry_point;
+
     std::tie(curr_label_index, curr_label_index_size) = load_label_index(curr_label_index_path, labels_to_number_of_points[label]);
-    final_index_size += curr_label_index_size;
     curr_label_entry_point = random(0, curr_label_index.size());
     label_entry_points[label] = label_id_to_orig_id_map[label][curr_label_entry_point];
   
@@ -516,17 +549,16 @@ stitch_indices_return_values stitch_label_indices(path final_index_path_prefix,
       for (auto &node_neighbor : curr_label_index[node_point]) {
         _u32 original_neighbor_id = label_id_to_orig_id_map[label][node_neighbor];
         std::vector<_u32> curr_point_neighbors = stitched_graph[original_point_id];
-        if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == curr_point_neighbors.end())
+        if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == curr_point_neighbors.end()) {
           stitched_graph[original_point_id].push_back(original_neighbor_id);
-        else
-          final_index_size -= sizeof(T);
+          final_index_size += sizeof(_u32);
+        }
       }
     }
   }
 
-  size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
-  final_index_size -= all_labels.size() * METADATA;
-  final_index_size += METADATA;
+  const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
+  final_index_size += (total_number_of_points * sizeof(_u32) + METADATA);
 
   std::chrono::duration<double> stitching_index_time = std::chrono::high_resolution_clock::now() - stitching_index_timer;
   std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl;
@@ -535,6 +567,62 @@ stitch_indices_return_values stitch_label_indices(path final_index_path_prefix,
 }
 
 
+template <typename T>
+void prune_and_save(path final_index_path_prefix, path input_data_path, std::vector<std::vector<_u32>> stitched_graph, unsigned stitched_R,
+                    tsl::robin_map<label, _u32> label_entry_points, label universal_label, path label_data_path, unsigned num_threads) {
+    size_t dimension, number_of_label_points;
+    auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr);
+    auto std_cout_buffer = std::cout.rdbuf(nullptr);
+    auto pruning_index_timer = std::chrono::high_resolution_clock::now();
+
+    diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension);
+    diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
+
+    // not searching this index, set search_l to 0
+    index.load(final_index_path_prefix.c_str(), 0, 1);
+
+    diskann::Parameters paras;
+    paras.Set<unsigned>("R", stitched_R);
+    paras.Set<unsigned>(
+      "C", 750);  // maximum candidate set size during pruning procedure
+    paras.Set<float>("alpha", 1.2);
+    paras.Set<bool>("saturate_graph", 1);
+      std::cout << "parsing labels" << std::endl;
+
+    index.prune_all_nbrs(paras);
+    index.save((final_index_path_prefix + "_pruned").c_str());
+    
+    diskann::cout.rdbuf(diskann_cout_buffer);
+    std::cout.rdbuf(std_cout_buffer);
+    std::chrono::duration<double> pruning_index_time = std::chrono::high_resolution_clock::now() - pruning_index_timer;
+    std::cout << "pruning performed in " << pruning_index_time.count() << " seconds\n" << std::endl;
+
+
+}
+
+/*
+ * Delete all temporary artifacts.
+ * In the process of creating the stitched index, some temporary artifacts are created:
+ * 1. the separate bin files for each labels' points
+ * 2. the separate diskANN indices built for each label
+ * 3. the '.data' file created while generating the indices
+ */
+void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) {
+  for (const auto &label : all_labels) {
+    path curr_label_input_data_path(input_data_path + "_" + label);
+    path curr_label_index_path(final_index_path_prefix + "_" + label);
+    path curr_label_index_path_data(curr_label_index_path + ".data");
+
+    if (std::remove(curr_label_index_path.c_str()) != 0)
+      throw;
+    if (std::remove(curr_label_input_data_path.c_str()) != 0)
+      throw;
+    if (std::remove(curr_label_index_path_data.c_str()) != 0)
+      throw;
+  }
+}
+
+
 int main(int argc, char **argv) {
   // 1. handle cmdline inputs
   std::string data_type;
@@ -543,6 +631,7 @@ int main(int argc, char **argv) {
   unsigned num_threads, R, L, stitched_R;
   float alpha;
 
+  auto index_timer = std::chrono::high_resolution_clock::now();
   handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, num_threads, R, L, stitched_R, alpha);
 
   // 2. parse label file and create necessary data structures
@@ -592,5 +681,21 @@ int main(int argc, char **argv) {
   // 5a. save the stitched graph to disk
   save_full_index(final_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, universal_label, label_data_path);
 
-  // TODO: 6. run a prune on the stitched index, and save to disk
+  // 6. run a prune on the stitched index, and save to disk
+  if (data_type == "uint8")
+    prune_and_save<uint8_t>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
+                    label_entry_points, universal_label, label_data_path, num_threads);
+  else if (data_type == "int8")
+    prune_and_save<int8_t>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
+                    label_entry_points, universal_label, label_data_path, num_threads);
+  else if (data_type == "float")
+    prune_and_save<float>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
+                    label_entry_points, universal_label, label_data_path, num_threads);
+  else
+   throw;
+  
+  std::chrono::duration<double> index_time = std::chrono::high_resolution_clock::now() - index_timer;
+  std::cout << "pruned/stitched graph generated in " << index_time.count() << " seconds" << std::endl;
+
+  clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels);
 }

From f8227148979031f964f12431fad914fb5ffa33ab Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Thu, 2 Feb 2023 23:37:10 +0530
Subject: [PATCH 05/13] Change all filter comparisons in stitched_vamana from
 string to integer for speedup

---
 tests/build_stitched_index.cpp | 37 +++++++++++++++++-----------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 34f00b605..1980f88ac 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -11,8 +11,9 @@
 #include <chrono>
 #include <cstdio>
 #include <cstring>
-#include <tuple>
 #include <random>
+#include <string>
+#include <tuple>
 
 #include "index.h"
 #include "parameters.h"
@@ -28,7 +29,7 @@ namespace po = boost::program_options;
 #define PBWIDTH 60
 
 // custom types (for readability)
-typedef std::string label;
+typedef unsigned label;
 typedef tsl::robin_set<label> label_set;
 typedef std::string path;
 
@@ -66,7 +67,7 @@ inline size_t random(size_t range_from, size_t range_to) {
  * Arguments are merely the inputs from the command line.
  */
 size_t handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, 
-                          path &final_index_path_prefix, path &label_data_path, std::string &universal_label, 
+                          path &final_index_path_prefix, path &label_data_path, label &universal_label, 
                           unsigned &num_threads, unsigned &R, unsigned &L, unsigned &stitched_R, float &alpha) {
   po::options_description desc{"Arguments"};
   try {
@@ -104,7 +105,7 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
         "Input label file in txt format if present");
     desc.add_options()(
         "universal_label",
-        po::value<label>(&universal_label)->default_value(""),
+        po::value<label>(&universal_label)->default_value(65537),
         "Universal label, if using it, only in conjunction with labels_file");
 
     po::variables_map vm;
@@ -129,7 +130,7 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
  * Returns three objects in via std::tuple:
  * 1. the label universe as a set
  */
-parse_label_file_return_values parse_label_file (path label_data_path, std::string universal_label) {
+parse_label_file_return_values parse_label_file (path label_data_path, label universal_label) {
   std::ifstream label_data_stream(label_data_path);
   std::string   line, token;
   unsigned      line_cnt = 0;
@@ -163,13 +164,14 @@ parse_label_file_return_values parse_label_file (path label_data_path, std::stri
       token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
 
       // if token is empty, there's no labels for the point
-      if (token == universal_label) {
+			unsigned token_as_num = std::stoul(token);
+      if (token_as_num == universal_label) {
         points_with_universal_label.push_back(point_id);
         current_universal_label_check = true;
       } else {
-        all_labels.insert(token);
-        current_labels.insert(token);
-        labels_to_number_of_points[token]++;
+        all_labels.insert(token_as_num);
+        current_labels.insert(token_as_num);
+        labels_to_number_of_points[token_as_num]++;
       }
     }
 
@@ -276,7 +278,7 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
   // write each label iovec to resp. file
   for (const auto &label : all_labels) {
     int label_input_data_fd;
-    path curr_label_input_data_path(input_data_path + "_" + label);
+    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
 
     label_input_data_fd = open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY, (mode_t) 0644);
     if (UNLIKELY(label_input_data_fd == -1))
@@ -326,12 +328,11 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
   std::cout.setstate(std::ios_base::failbit);
   diskann::cout.setstate(std::ios_base::failbit);
   for (const auto &label : all_labels) {
-    path curr_label_input_data_path(input_data_path + "_" + label);
-    path curr_label_index_path(final_index_path_prefix + "_" + label);
+    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
+    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
     
     size_t number_of_label_points, dimension;
     diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension);
-    std::cout << number_of_label_points << " " << dimension << std::endl;
     diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
 
     auto index_build_timer = std::chrono::high_resolution_clock::now();
@@ -435,7 +436,7 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   labels_to_medoids_writer.close();
 
   // aux. file 4
-  if (universal_label != "") {
+  if (universal_label != 65537) {
     std::ofstream universal_label_writer;
     universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit);
     universal_label_writer.open(final_index_path_prefix + "_universal_label.txt");
@@ -535,7 +536,7 @@ stitch_indices_return_values stitch_label_indices(path final_index_path_prefix,
   
   auto stitching_index_timer = std::chrono::high_resolution_clock::now();
   for (const auto &label : all_labels) {
-    path curr_label_index_path(final_index_path_prefix + "_" + label);
+    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
     std::vector<std::vector<_u32>> curr_label_index;
     _u64 curr_label_index_size;
     _u32 curr_label_entry_point;
@@ -609,8 +610,8 @@ void prune_and_save(path final_index_path_prefix, path input_data_path, std::vec
  */
 void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) {
   for (const auto &label : all_labels) {
-    path curr_label_input_data_path(input_data_path + "_" + label);
-    path curr_label_index_path(final_index_path_prefix + "_" + label);
+    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
+    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
     path curr_label_index_path_data(curr_label_index_path + ".data");
 
     if (std::remove(curr_label_index_path.c_str()) != 0)
@@ -627,7 +628,7 @@ int main(int argc, char **argv) {
   // 1. handle cmdline inputs
   std::string data_type;
   path input_data_path, final_index_path_prefix, label_data_path;
-  label universal_label;
+	label universal_label;
   unsigned num_threads, R, L, stitched_R;
   float alpha;
 

From 15ac703b8c11bdb0f4c66f846498d36a36a64a04 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Tue, 31 Jan 2023 14:05:55 +0530
Subject: [PATCH 06/13] Add fixes for file generation

---
 tests/build_stitched_index.cpp | 579 ++++++++++++++++++++-------------
 1 file changed, 344 insertions(+), 235 deletions(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 1980f88ac..496636d13 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -1,5 +1,5 @@
-// authors: Siddharth Gollapudi, Varun Sivashankar
-// emails: t-gollapudis@microsoft.com, t-varunsi@microsoft.com
+// author: Siddharth Gollapudi
+// email: t-gollapudis@microsoft.com
 
 #include <boost/program_options.hpp>
 #include <bits/types/struct_iovec.h>
@@ -19,36 +19,37 @@
 #include "parameters.h"
 #include "utils.h"
 
-
 namespace po = boost::program_options;
 
-// macros (posix reasons)
-#define   LIKELY(condition) __builtin_expect(static_cast<bool>(condition), 1)
+// macros
+#define LIKELY(condition) __builtin_expect(static_cast<bool>(condition), 1)
 #define UNLIKELY(condition) __builtin_expect(static_cast<bool>(condition), 0)
 #define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"
 #define PBWIDTH 60
 
 // custom types (for readability)
-typedef unsigned label;
+typedef unsigned              label;
 typedef tsl::robin_set<label> label_set;
-typedef std::string path;
+typedef std::string           path;
 
 // structs for returning multiple items from a function
-typedef std::tuple<std::vector<label_set>, tsl::robin_map<label, _u32>, label_set> parse_label_file_return_values;
-typedef std::tuple<std::vector<std::vector<_u32>>, _u64> load_label_index_return_values;
-typedef std::tuple<std::vector<std::vector<_u32>>, _u64> stitch_indices_return_values;
-
-
+typedef std::tuple<std::vector<label_set>, tsl::robin_map<label, _u32>,
+                   label_set>
+    parse_label_file_return_values;
+typedef std::tuple<std::vector<std::vector<_u32>>, _u64>
+    load_label_index_return_values;
+typedef std::tuple<std::vector<std::vector<_u32>>, _u64>
+    stitch_indices_return_values;
 
 /*
  * Inline function to display progress bar.
  */
 inline void print_progress(double percentage) {
-    int val = (int) (percentage * 100);
-    int lpad = (int) (percentage * PBWIDTH);
-    int rpad = PBWIDTH - lpad;
-    printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, "");
-    fflush(stdout);
+  int val = (int) (percentage * 100);
+  int lpad = (int) (percentage * PBWIDTH);
+  int rpad = PBWIDTH - lpad;
+  printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, "");
+  fflush(stdout);
 }
 
 /*
@@ -66,9 +67,11 @@ inline size_t random(size_t range_from, size_t range_to) {
  *
  * Arguments are merely the inputs from the command line.
  */
-size_t handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, 
-                          path &final_index_path_prefix, path &label_data_path, label &universal_label, 
-                          unsigned &num_threads, unsigned &R, unsigned &L, unsigned &stitched_R, float &alpha) {
+size_t handle_args(int argc, char **argv, std::string &data_type,
+                   path &input_data_path, path &final_index_path_prefix,
+                   path &label_data_path, label &universal_label,
+                   unsigned &num_threads, unsigned &R, unsigned &L,
+                   unsigned &stitched_R, float &alpha) {
   po::options_description desc{"Arguments"};
   try {
     desc.add_options()("help,h", "Print information on arguments");
@@ -87,9 +90,9 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
     desc.add_options()(
         "Lbuild,L", po::value<uint32_t>(&L)->default_value(100),
         "Build complexity, higher value results in better graphs");
-    desc.add_options()(
-        "stitched_R", po::value<uint32_t>(&stitched_R)->default_value(100),
-        "Degree to prune final graph down to");
+    desc.add_options()("stitched_R",
+                       po::value<uint32_t>(&stitched_R)->default_value(100),
+                       "Degree to prune final graph down to");
     desc.add_options()(
         "alpha", po::value<float>(&alpha)->default_value(1.2f),
         "alpha controls density and diameter of graph, set 1 for sparse graph, "
@@ -99,10 +102,9 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
         po::value<uint32_t>(&num_threads)->default_value(omp_get_num_procs()),
         "Number of threads used for building index (defaults to "
         "omp_get_num_procs())");
-    desc.add_options()(
-        "label_file",
-        po::value<path>(&label_data_path)->default_value(""),
-        "Input label file in txt format if present");
+    desc.add_options()("label_file",
+                       po::value<path>(&label_data_path)->default_value(""),
+                       "Input label file in txt format if present");
     desc.add_options()(
         "universal_label",
         po::value<label>(&universal_label)->default_value(65537),
@@ -112,17 +114,16 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
     po::store(po::parse_command_line(argc, argv, desc), vm);
     if (vm.count("help")) {
       std::cout << desc;
-      return 0;
+      exit(0);
     }
     po::notify(vm);
-  } catch (const std::exception& ex) {
+  } catch (const std::exception &ex) {
     std::cerr << ex.what() << '\n';
-    return -1;
+    throw;
   }
   return 1;
 }
 
-
 /*
  * Parses the label datafile, which has the following line format:
  *            point_id -> \t -> comma-separated labels
@@ -130,7 +131,8 @@ size_t handle_args(int argc, char **argv, std::string &data_type, path &input_da
  * Returns three objects in via std::tuple:
  * 1. the label universe as a set
  */
-parse_label_file_return_values parse_label_file (path label_data_path, label universal_label) {
+parse_label_file_return_values parse_label_file(path  label_data_path,
+                                                label universal_label) {
   std::ifstream label_data_stream(label_data_path);
   std::string   line, token;
   unsigned      line_cnt = 0;
@@ -142,14 +144,14 @@ parse_label_file_return_values parse_label_file (path label_data_path, label uni
   label_data_stream.seekg(0, std::ios::beg);
 
   // values to return
-  std::vector<label_set> point_ids_to_labels(line_cnt);
+  std::vector<label_set>      point_ids_to_labels(line_cnt);
   tsl::robin_map<label, _u32> labels_to_number_of_points;
-  label_set all_labels;
-  
+  label_set                   all_labels;
+
   std::vector<_u32> points_with_universal_label;
   while (std::getline(label_data_stream, line)) {
-    std::istringstream       current_point_and_labels(line);
-    label_set current_labels;
+    std::istringstream current_point_and_labels(line);
+    label_set          current_labels;
 
     // get point id
     getline(current_point_and_labels, token, '\t');
@@ -158,13 +160,13 @@ parse_label_file_return_values parse_label_file (path label_data_path, label uni
     // parse comma separated labels
     getline(current_point_and_labels, token, '\t');
     std::istringstream current_labels_comma_separated(token);
-    bool current_universal_label_check = false;
+    bool               current_universal_label_check = false;
     while (getline(current_labels_comma_separated, token, ',')) {
       token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
       token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
 
       // if token is empty, there's no labels for the point
-			unsigned token_as_num = std::stoul(token);
+      unsigned token_as_num = std::stoul(token);
       if (token_as_num == universal_label) {
         points_with_universal_label.push_back(point_id);
         current_universal_label_check = true;
@@ -178,7 +180,7 @@ parse_label_file_return_values parse_label_file (path label_data_path, label uni
     if (current_labels.size() <= 0 && !current_universal_label_check) {
       std::cerr << "Error: " << point_id << " has no labels." << std::endl;
       exit(-1);
-    } 
+    }
     point_ids_to_labels[point_id] = current_labels;
   }
 
@@ -190,13 +192,14 @@ parse_label_file_return_values parse_label_file (path label_data_path, label uni
       labels_to_number_of_points[label]++;
   }
 
-  std::cout << "Identified " << all_labels.size() << " distinct label(s) for " << point_ids_to_labels.size() << " points\n"
+  std::cout << "Identified " << all_labels.size() << " distinct label(s) for "
+            << point_ids_to_labels.size() << " points\n"
             << std::endl;
 
-  return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels);
+  return std::make_tuple(point_ids_to_labels, labels_to_number_of_points,
+                         all_labels);
 }
 
-
 /*
  * For each label, generates a file containing all vectors that have said label.
  * Also copies data from original bin file to new dimension-aligned file.
@@ -206,17 +209,19 @@ parse_label_file_return_values parse_label_file (path label_data_path, label uni
  *    input_data_path + "_" + label
  */
 template<typename T>
-tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(path input_data_path, 
-                                                                              tsl::robin_map<label, _u32> labels_to_number_of_points,
-                                                                              std::vector<label_set> point_ids_to_labels, label_set all_labels) {
-  void *input_memblock;
-  int input_data_fd;
-  struct stat input_data_st;
+tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
+    path                        input_data_path,
+    tsl::robin_map<label, _u32> labels_to_number_of_points,
+    std::vector<label_set> point_ids_to_labels, label_set all_labels) {
+  void *       input_memblock;
+  int          input_data_fd;
+  struct stat  input_data_st;
   const size_t METADATA = 2 * sizeof(_u32);
-  auto file_writing_timer = std::chrono::high_resolution_clock::now();
+  auto         file_writing_timer = std::chrono::high_resolution_clock::now();
 
   // UNLIKELY moves the branch instr. body elsewhere, opening up icache space
-  // for code that actually executes
+  // for code that actually executes (but this is probably irrelevant unless
+  // there is a massive number of labels)
   input_data_fd = open(input_data_path.c_str(), O_RDONLY, (mode_t) 0444);
   if (UNLIKELY(input_data_fd == -1))
     throw;
@@ -225,99 +230,132 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(pa
     throw;
   }
 
-  // mmap could be faster than just doing a regular sequential read?
-  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_SHARED, input_data_fd, 0);
+  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_SHARED,
+                        input_data_fd, 0);
   if (UNLIKELY(input_memblock == MAP_FAILED)) {
     close(input_data_fd);
     throw;
   }
-  const char *input_start = static_cast<char const *>(input_memblock);
-  _u32 number_of_points, dimension;
+
+  char *input_start = static_cast<char *>(input_memblock);
+  _u32  number_of_points, dimension;
   std::memcpy(&number_of_points, input_start, sizeof(_u32));
   std::memcpy(&dimension, input_start + sizeof(_u32), sizeof(_u32));
   const _u32 VECTOR_SIZE = dimension * sizeof(T);
   if (number_of_points != point_ids_to_labels.size()) {
-    std::cerr << "Error: number of points in labels file and data file differ." << std::endl;
+    std::cerr << "Error: number of points in labels file and data file differ."
+              << std::endl;
     throw;
   }
 
   // iovecs necessary for using writev
-  tsl::robin_map<label, iovec> label_to_iovec_map;
-
-  // mapping from id in label file to original file
+  tsl::robin_map<label, iovec *>           label_to_iovec_map;
+  tsl::robin_map<label, _u32>              label_to_curr_iovec;
   tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id;
+
+  // setup iovec list for each label
+  for (const auto &label : all_labels) {
+    iovec *label_iovecs =
+        (iovec *) malloc(labels_to_number_of_points[label] * sizeof(iovec));
+    label_to_iovec_map[label] = label_iovecs;
+    label_to_curr_iovec[label] = 0;
+    label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);
+  }
+
+  // each point added to corresponding per-label iovec list
   for (_u32 point_id = 0; point_id < number_of_points; point_id++) {
-    const char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id);
+    char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id);
+    iovec curr_iovec;
+
+    curr_iovec.iov_base = curr_point;
+    curr_iovec.iov_len = VECTOR_SIZE;
     for (const auto &label : point_ids_to_labels[point_id]) {
-      // first add metadata before adding vectors
-      if (!label_to_iovec_map.count(label)) {
-        iovec curr_iovec;
-        curr_iovec.iov_base = malloc(METADATA + labels_to_number_of_points[label] * VECTOR_SIZE); 
-        std::memcpy((char *) curr_iovec.iov_base, &labels_to_number_of_points[label], sizeof(_u32));
-        curr_iovec.iov_len = sizeof(_u32);
-        std::memcpy((char *) curr_iovec.iov_base + curr_iovec.iov_len, &dimension, sizeof(_u32));
-        curr_iovec.iov_len += sizeof(_u32);
-        label_to_iovec_map[label] = curr_iovec;
-        label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);
-      }
-      char *current_iovec_buffer = (char *) label_to_iovec_map[label].iov_base;
-      size_t *current_iovec_buffer_size = &label_to_iovec_map[label].iov_len;
-      std::memcpy(current_iovec_buffer + *current_iovec_buffer_size, curr_point, VECTOR_SIZE);
-      *current_iovec_buffer_size += VECTOR_SIZE;
+      *(label_to_iovec_map[label] + label_to_curr_iovec[label]) = curr_iovec;
+      label_to_curr_iovec[label]++;
       label_id_to_orig_id[label].push_back(point_id);
     }
-    
-  }
-
-  if (UNLIKELY(munmap(input_memblock, input_data_st.st_size) == -1)) {
-    close(input_data_fd);
-    throw;
   }
-  close(input_data_fd);
 
   // write each label iovec to resp. file
   for (const auto &label : all_labels) {
-    int label_input_data_fd;
-    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
-
-    label_input_data_fd = open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY, (mode_t) 0644);
+    int  label_input_data_fd;
+    path curr_label_input_data_path(input_data_path + "_" +
+                                    std::to_string(label));
+    _u32 curr_num_pts = labels_to_number_of_points[label];
+
+    label_input_data_fd =
+        open(curr_label_input_data_path.c_str(),
+             O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t) 0644);
     if (UNLIKELY(label_input_data_fd == -1))
       throw;
 
-    // using writev leads to one write syscall per label file
-    int return_value = writev(label_input_data_fd, &label_to_iovec_map[label], 1);
+    // write metadata
+    _u32 metadata[2] = {curr_num_pts, dimension};
+    int  return_value = write(label_input_data_fd, metadata, sizeof(_u32) * 2);
     if (UNLIKELY(return_value == -1)) {
       close(label_input_data_fd);
       throw;
     }
 
-    free(label_to_iovec_map[label].iov_base);
+    // limits on number of iovec structs per writev means we need to perform
+    // multiple writevs
+    size_t i = 0;
+    while (curr_num_pts > IOV_MAX) {
+      return_value =
+          writev(label_input_data_fd,
+                 (label_to_iovec_map[label] + (IOV_MAX * i)), IOV_MAX);
+      if (UNLIKELY(return_value == -1)) {
+        close(label_input_data_fd);
+        throw;
+      }
+      curr_num_pts -= IOV_MAX;
+      i += 1;
+    }
+    return_value =
+        writev(label_input_data_fd, (label_to_iovec_map[label] + (IOV_MAX * i)),
+               curr_num_pts);
+    if (UNLIKELY(return_value == -1)) {
+      close(label_input_data_fd);
+      throw;
+    }
+
+    free(label_to_iovec_map[label]);
     close(label_input_data_fd);
   }
 
-  std::chrono::duration<double> file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer;
-  std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time "
-    << file_writing_time.count() << "\n" << std::endl;
+  // un-mmap (or munmap)
+  if (UNLIKELY(munmap(input_memblock, input_data_st.st_size) == -1)) {
+    close(input_data_fd);
+    throw;
+  }
+  close(input_data_fd);
+
+  std::chrono::duration<double> file_writing_time =
+      std::chrono::high_resolution_clock::now() - file_writing_timer;
+  std::cout << "generated " << all_labels.size()
+            << " label-specific vector files for index building in time "
+            << file_writing_time.count() << "\n"
+            << std::endl;
 
   return label_id_to_orig_id;
 }
 
-
 /*
- * Using passed in parameters and files generated from step 3, 
+ * Using passed in parameters and files generated from step 3,
  * builds a vanilla diskANN index for each label.
  *
  * Each index is saved under the following path:
  *  final_index_path_prefix + "_" + label
  */
-template <typename T>
+template<typename T>
 void generate_label_indices(path input_data_path, path final_index_path_prefix,
-                            label_set all_labels, unsigned R, unsigned L, float alpha, 
-                            unsigned num_threads) {
+                            label_set all_labels, unsigned R, unsigned L,
+                            float alpha, unsigned num_threads) {
   diskann::Parameters label_index_build_parameters;
   label_index_build_parameters.Set<unsigned>("R", R);
   label_index_build_parameters.Set<unsigned>("L", L);
-  label_index_build_parameters.Set<unsigned>("C", 750);  // maximum candidate set size during prune
+  label_index_build_parameters.Set<unsigned>(
+      "C", 750);  // maximum candidate set size during prune
   label_index_build_parameters.Set<bool>("saturate_graph", 0);
   label_index_build_parameters.Set<float>("alpha", alpha);
   label_index_build_parameters.Set<unsigned>("num_threads", num_threads);
@@ -328,18 +366,23 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
   std::cout.setstate(std::ios_base::failbit);
   diskann::cout.setstate(std::ios_base::failbit);
   for (const auto &label : all_labels) {
-    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
-    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
-    
+    path curr_label_input_data_path(input_data_path + "_" +
+                                    std::to_string(label));
+    path curr_label_index_path(final_index_path_prefix + "_" +
+                               std::to_string(label));
+
     size_t number_of_label_points, dimension;
-    diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension);
-    diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
+    diskann::get_bin_metadata(curr_label_input_data_path,
+                              number_of_label_points, dimension);
+    diskann::Index<T> index(diskann::Metric::L2, dimension,
+                            number_of_label_points, false, false);
 
     auto index_build_timer = std::chrono::high_resolution_clock::now();
-    index.build(curr_label_input_data_path.c_str(), number_of_label_points, label_index_build_parameters);
+    index.build(curr_label_input_data_path.c_str(), number_of_label_points,
+                label_index_build_parameters);
     std::chrono::duration<double> current_indexing_time =
-      std::chrono::high_resolution_clock::now() - index_build_timer;
-  
+        std::chrono::high_resolution_clock::now() - index_build_timer;
+
     total_indexing_time += current_indexing_time.count();
     indexing_percentage += (1 / (double) all_labels.size());
     print_progress(indexing_percentage);
@@ -349,25 +392,24 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
   std::cout.clear();
   diskann::cout.clear();
 
-  std::cout << "\nDone. Generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl;
+  std::cout << "\nDone. Generated per-label indices in " << total_indexing_time
+            << " seconds\n"
+            << std::endl;
 }
 
-
-
-
-
 /*
  * Manually loads a graph index in from a given file.
- * 
+ *
  * Returns both the graph index and the size of the file in bytes.
  */
-load_label_index_return_values load_label_index(path label_index_path, _u32 label_number_of_points) {
+load_label_index_return_values load_label_index(path label_index_path,
+                                                _u32 label_number_of_points) {
   std::ifstream label_index_stream;
   label_index_stream.exceptions(std::ios::badbit | std::ios::failbit);
   label_index_stream.open(label_index_path, std::ios::binary);
 
-  _u64 index_file_size, index_num_frozen_points;
-  _u32 index_max_observed_degree, index_entry_point;
+  _u64         index_file_size, index_num_frozen_points;
+  _u32         index_max_observed_degree, index_entry_point;
   const size_t INDEX_METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
   label_index_stream.read((char *) &index_file_size, sizeof(_u64));
   label_index_stream.read((char *) &index_max_observed_degree, sizeof(_u32));
@@ -376,14 +418,15 @@ load_label_index_return_values load_label_index(path label_index_path, _u32 labe
   size_t bytes_read = INDEX_METADATA;
 
   std::vector<std::vector<_u32>> label_index(label_number_of_points);
-  _u32 nodes_read = 0;
+  _u32                           nodes_read = 0;
   while (bytes_read != index_file_size) {
     _u32 current_node_num_neighbors;
     label_index_stream.read((char *) &current_node_num_neighbors, sizeof(_u32));
     nodes_read++;
-    
+
     std::vector<_u32> current_node_neighbors(current_node_num_neighbors);
-    label_index_stream.read((char *) current_node_neighbors.data(), current_node_num_neighbors * sizeof(_u32));
+    label_index_stream.read((char *) current_node_neighbors.data(),
+                            current_node_num_neighbors * sizeof(_u32));
     label_index[nodes_read - 1].swap(current_node_neighbors);
     bytes_read += sizeof(_u32) * (current_node_num_neighbors + 1);
   }
@@ -393,25 +436,26 @@ load_label_index_return_values load_label_index(path label_index_path, _u32 labe
 
 /*
  * Custom index save to write the in-memory index to disk.
- * Also writes required files for diskANN API - 
+ * Also writes required files for diskANN API -
  *  1. labels_to_medoids
  *  2. universal_label
  *  3. data (redundant for static indices)
  *  4. labels (redundant for static indices)
  */
 void save_full_index(path final_index_path_prefix, path input_data_path,
-                     _u64 final_index_size, std::vector<std::vector<_u32>> stitched_graph,
-                     tsl::robin_map<label, _u32> entry_points,
+                     _u64                           final_index_size,
+                     std::vector<std::vector<_u32>> stitched_graph,
+                     tsl::robin_map<label, _u32>    entry_points,
                      label universal_label, path label_data_path) {
-
   // aux. file 1
-  auto saving_index_timer = std::chrono::high_resolution_clock::now();
+  auto          saving_index_timer = std::chrono::high_resolution_clock::now();
   std::ifstream original_label_data_stream;
   original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
   original_label_data_stream.open(label_data_path, std::ios::binary);
   std::ofstream new_label_data_stream;
   new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
-  new_label_data_stream.open(final_index_path_prefix + "_labels.txt", std::ios::binary);
+  new_label_data_stream.open(final_index_path_prefix + "_labels.txt",
+                             std::ios::binary);
   new_label_data_stream << original_label_data_stream.rdbuf();
   original_label_data_stream.close();
   new_label_data_stream.close();
@@ -422,7 +466,8 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   original_input_data_stream.open(input_data_path, std::ios::binary);
   std::ofstream new_input_data_stream;
   new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit);
-  new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary);
+  new_input_data_stream.open(final_index_path_prefix + ".data",
+                             std::ios::binary);
   new_input_data_stream << original_input_data_stream.rdbuf();
   original_input_data_stream.close();
   new_input_data_stream.close();
@@ -430,7 +475,8 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   // aux. file 3
   std::ofstream labels_to_medoids_writer;
   labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit);
-  labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt");
+  labels_to_medoids_writer.open(final_index_path_prefix +
+                                "_labels_to_medoids.txt");
   for (auto iter : entry_points)
     labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl;
   labels_to_medoids_writer.close();
@@ -439,7 +485,8 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   if (universal_label != 65537) {
     std::ofstream universal_label_writer;
     universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit);
-    universal_label_writer.open(final_index_path_prefix + "_universal_label.txt");
+    universal_label_writer.open(final_index_path_prefix +
+                                "_universal_label.txt");
     universal_label_writer << universal_label << std::endl;
     universal_label_writer.close();
   }
@@ -447,16 +494,18 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   // main index
   _u64 index_num_frozen_points = 0, index_num_edges = 0;
   _u32 index_max_observed_degree = 0, index_entry_point = 0;
-  for (auto& point_neighbors : stitched_graph) {
-    index_max_observed_degree = std::max(index_max_observed_degree, (_u32) point_neighbors.size());
+  for (auto &point_neighbors : stitched_graph) {
+    index_max_observed_degree =
+        std::max(index_max_observed_degree, (_u32) point_neighbors.size());
   }
 
-  void *memblock;
+  void *       memblock;
   const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
-  struct stat sb;
-  int stitched_graph_fd;
+  struct stat  sb;
+  int          stitched_graph_fd;
 
-  stitched_graph_fd = open(final_index_path_prefix.c_str(), O_RDWR | O_CREAT, (mode_t) 0644);
+  stitched_graph_fd =
+      open(final_index_path_prefix.c_str(), O_RDWR | O_CREAT, (mode_t) 0644);
   if (UNLIKELY(stitched_graph_fd == -1))
     throw;
   if (UNLIKELY(fstat(stitched_graph_fd, &sb))) {
@@ -470,22 +519,26 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   }
 
   // mmap could be faster than just doing a regular sequential write?
-  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_SHARED, stitched_graph_fd, 0);
+  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_SHARED,
+                  stitched_graph_fd, 0);
   if (UNLIKELY(memblock == MAP_FAILED)) {
     close(stitched_graph_fd);
     throw;
   }
-  
+
   char *begin = static_cast<char *>(memblock);
   std::memcpy(begin, &final_index_size, sizeof(_u64));
   std::memcpy(begin + sizeof(_u64), &index_max_observed_degree, sizeof(_u32));
-  std::memcpy(begin + sizeof(_u64) + sizeof(_u32), &index_entry_point, sizeof(_u32));
-  std::memcpy(begin + sizeof(_u64) + 2 * sizeof(_u32), &index_num_frozen_points, sizeof(_u64));
+  std::memcpy(begin + sizeof(_u64) + sizeof(_u32), &index_entry_point,
+              sizeof(_u32));
+  std::memcpy(begin + sizeof(_u64) + 2 * sizeof(_u32), &index_num_frozen_points,
+              sizeof(_u64));
   size_t bytes_written = METADATA;
   for (_u32 node_point = 0; node_point < stitched_graph.size(); node_point++) {
     _u32 current_node_num_neighbors = stitched_graph[node_point].size();
     std::vector<_u32> current_node_neighbors = stitched_graph[node_point];
-    std::memcpy(begin + bytes_written, &current_node_num_neighbors, sizeof(_u32));
+    std::memcpy(begin + bytes_written, &current_node_num_neighbors,
+                sizeof(_u32));
     bytes_written += sizeof(_u32);
     for (const auto &current_node_neighbor : current_node_neighbors) {
       std::memcpy(begin + bytes_written, &current_node_neighbor, sizeof(_u32));
@@ -495,12 +548,14 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   }
 
   if (UNLIKELY(bytes_written != final_index_size)) {
-    std::cerr << "Error: written bytes does not match allocated space" << std::endl;
+    std::cerr << "Error: written bytes does not match allocated space"
+              << std::endl;
     throw;
   }
 
   if (UNLIKELY(msync(memblock, final_index_size, MS_SYNC) == -1)) {
-    std::cerr << "Error: could not sync memory-mapped writes to disk" << std::endl;
+    std::cerr << "Error: could not sync memory-mapped writes to disk"
+              << std::endl;
     close(stitched_graph_fd);
     throw;
   }
@@ -513,44 +568,58 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
 
   close(stitched_graph_fd);
 
-  std::chrono::duration<double> saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer;
-  std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl;
-  std::cout << "Stitched graph average degree: " << ((float) index_num_edges) / ((float) (stitched_graph.size())) << std::endl;
-  std::cout << "Stitched graph max degree: " << index_max_observed_degree << std::endl << std::endl;
+  std::chrono::duration<double> saving_index_time =
+      std::chrono::high_resolution_clock::now() - saving_index_timer;
+  std::cout << "Stitched graph written in " << saving_index_time.count()
+            << " seconds" << std::endl;
+  std::cout << "Stitched graph average degree: "
+            << ((float) index_num_edges) / ((float) (stitched_graph.size()))
+            << std::endl;
+  std::cout << "Stitched graph max degree: " << index_max_observed_degree
+            << std::endl
+            << std::endl;
 }
 
-
 /*
  * Unions the per-label graph indices together via the following policy:
  *  - any two nodes can only have at most one edge between them -
- * 
+ *
  * Returns the "stitched" graph and its expected file size.
  */
-template <typename T>
-stitch_indices_return_values stitch_label_indices(path final_index_path_prefix, _u32 total_number_of_points, label_set all_labels, 
-                          tsl::robin_map<label, _u32> labels_to_number_of_points, 
-                          tsl::robin_map<label, _u32> &label_entry_points,
-                          tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map) {
-  size_t final_index_size = 0;
+template<typename T>
+stitch_indices_return_values stitch_label_indices(
+    path final_index_path_prefix, _u32 total_number_of_points,
+    label_set                                all_labels,
+    tsl::robin_map<label, _u32>              labels_to_number_of_points,
+    tsl::robin_map<label, _u32> &            label_entry_points,
+    tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map) {
+  size_t                         final_index_size = 0;
   std::vector<std::vector<_u32>> stitched_graph(total_number_of_points);
-  
+
   auto stitching_index_timer = std::chrono::high_resolution_clock::now();
   for (const auto &label : all_labels) {
-    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
+    path curr_label_index_path(final_index_path_prefix + "_" +
+                               std::to_string(label));
     std::vector<std::vector<_u32>> curr_label_index;
-    _u64 curr_label_index_size;
-    _u32 curr_label_entry_point;
+    _u64                           curr_label_index_size;
+    _u32                           curr_label_entry_point;
 
-    std::tie(curr_label_index, curr_label_index_size) = load_label_index(curr_label_index_path, labels_to_number_of_points[label]);
+    std::tie(curr_label_index, curr_label_index_size) = load_label_index(
+        curr_label_index_path, labels_to_number_of_points[label]);
     curr_label_entry_point = random(0, curr_label_index.size());
-    label_entry_points[label] = label_id_to_orig_id_map[label][curr_label_entry_point];
-  
-    for (_u32 node_point = 0; node_point < curr_label_index.size(); node_point++) {
+    label_entry_points[label] =
+        label_id_to_orig_id_map[label][curr_label_entry_point];
+
+    for (_u32 node_point = 0; node_point < curr_label_index.size();
+         node_point++) {
       _u32 original_point_id = label_id_to_orig_id_map[label][node_point];
       for (auto &node_neighbor : curr_label_index[node_point]) {
-        _u32 original_neighbor_id = label_id_to_orig_id_map[label][node_neighbor];
-        std::vector<_u32> curr_point_neighbors = stitched_graph[original_point_id];
-        if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == curr_point_neighbors.end()) {
+        _u32 original_neighbor_id =
+            label_id_to_orig_id_map[label][node_neighbor];
+        std::vector<_u32> curr_point_neighbors =
+            stitched_graph[original_point_id];
+        if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(),
+                      original_neighbor_id) == curr_point_neighbors.end()) {
           stitched_graph[original_point_id].push_back(original_neighbor_id);
           final_index_size += sizeof(_u32);
         }
@@ -561,57 +630,68 @@ stitch_indices_return_values stitch_label_indices(path final_index_path_prefix,
   const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
   final_index_size += (total_number_of_points * sizeof(_u32) + METADATA);
 
-  std::chrono::duration<double> stitching_index_time = std::chrono::high_resolution_clock::now() - stitching_index_timer;
-  std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl;
+  std::chrono::duration<double> stitching_index_time =
+      std::chrono::high_resolution_clock::now() - stitching_index_timer;
+  std::cout << "stitched graph generated in memory in "
+            << stitching_index_time.count() << " seconds" << std::endl;
 
   return std::make_tuple(stitched_graph, final_index_size);
 }
 
-
-template <typename T>
-void prune_and_save(path final_index_path_prefix, path input_data_path, std::vector<std::vector<_u32>> stitched_graph, unsigned stitched_R,
-                    tsl::robin_map<label, _u32> label_entry_points, label universal_label, path label_data_path, unsigned num_threads) {
-    size_t dimension, number_of_label_points;
-    auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr);
-    auto std_cout_buffer = std::cout.rdbuf(nullptr);
-    auto pruning_index_timer = std::chrono::high_resolution_clock::now();
-
-    diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension);
-    diskann::Index<T> index(diskann::Metric::L2, dimension, number_of_label_points, false, false);
-
-    // not searching this index, set search_l to 0
-    index.load(final_index_path_prefix.c_str(), 0, 1);
-
-    diskann::Parameters paras;
-    paras.Set<unsigned>("R", stitched_R);
-    paras.Set<unsigned>(
+template<typename T>
+void prune_and_save(path final_index_path_prefix, path input_data_path,
+                    std::vector<std::vector<_u32>> stitched_graph,
+                    unsigned                       stitched_R,
+                    tsl::robin_map<label, _u32>    label_entry_points,
+                    label universal_label, path label_data_path,
+                    unsigned num_threads) {
+  size_t dimension, number_of_label_points;
+  auto   diskann_cout_buffer = diskann::cout.rdbuf(nullptr);
+  auto   std_cout_buffer = std::cout.rdbuf(nullptr);
+  auto   pruning_index_timer = std::chrono::high_resolution_clock::now();
+
+  diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension);
+  diskann::Index<T> index(diskann::Metric::L2, dimension,
+                          number_of_label_points, false, false);
+
+  // not searching this index, set search_l to 0
+  index.load(final_index_path_prefix.c_str(), 0, 1);
+
+  diskann::Parameters paras;
+  paras.Set<unsigned>("R", stitched_R);
+  paras.Set<unsigned>(
       "C", 750);  // maximum candidate set size during pruning procedure
-    paras.Set<float>("alpha", 1.2);
-    paras.Set<bool>("saturate_graph", 1);
-      std::cout << "parsing labels" << std::endl;
-
-    index.prune_all_nbrs(paras);
-    index.save((final_index_path_prefix + "_pruned").c_str());
-    
-    diskann::cout.rdbuf(diskann_cout_buffer);
-    std::cout.rdbuf(std_cout_buffer);
-    std::chrono::duration<double> pruning_index_time = std::chrono::high_resolution_clock::now() - pruning_index_timer;
-    std::cout << "pruning performed in " << pruning_index_time.count() << " seconds\n" << std::endl;
-
-
+  paras.Set<float>("alpha", 1.2);
+  paras.Set<bool>("saturate_graph", 1);
+  std::cout << "parsing labels" << std::endl;
+
+  index.prune_all_nbrs(paras);
+  index.save((final_index_path_prefix + "_pruned").c_str());
+
+  diskann::cout.rdbuf(diskann_cout_buffer);
+  std::cout.rdbuf(std_cout_buffer);
+  std::chrono::duration<double> pruning_index_time =
+      std::chrono::high_resolution_clock::now() - pruning_index_timer;
+  std::cout << "pruning performed in " << pruning_index_time.count()
+            << " seconds\n"
+            << std::endl;
 }
 
 /*
  * Delete all temporary artifacts.
- * In the process of creating the stitched index, some temporary artifacts are created:
+ * In the process of creating the stitched index, some temporary artifacts are
+ * created:
  * 1. the separate bin files for each labels' points
  * 2. the separate diskANN indices built for each label
  * 3. the '.data' file created while generating the indices
  */
-void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) {
+void clean_up_artifacts(path input_data_path, path final_index_path_prefix,
+                        label_set all_labels) {
   for (const auto &label : all_labels) {
-    path curr_label_input_data_path(input_data_path + "_" + std::to_string(label));
-    path curr_label_index_path(final_index_path_prefix + "_" + std::to_string(label));
+    path curr_label_input_data_path(input_data_path + "_" +
+                                    std::to_string(label));
+    path curr_label_index_path(final_index_path_prefix + "_" +
+                               std::to_string(label));
     path curr_label_index_path_data(curr_label_index_path + ".data");
 
     if (std::remove(curr_label_index_path.c_str()) != 0)
@@ -623,80 +703,109 @@ void clean_up_artifacts(path input_data_path, path final_index_path_prefix, labe
   }
 }
 
-
 int main(int argc, char **argv) {
   // 1. handle cmdline inputs
   std::string data_type;
-  path input_data_path, final_index_path_prefix, label_data_path;
-	label universal_label;
-  unsigned num_threads, R, L, stitched_R;
-  float alpha;
+  path        input_data_path, final_index_path_prefix, label_data_path;
+  label       universal_label;
+  unsigned    num_threads, R, L, stitched_R;
+  float       alpha;
 
   auto index_timer = std::chrono::high_resolution_clock::now();
-  handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, num_threads, R, L, stitched_R, alpha);
+  handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix,
+              label_data_path, universal_label, num_threads, R, L, stitched_R,
+              alpha);
 
   // 2. parse label file and create necessary data structures
-  std::vector<label_set> point_ids_to_labels;
+  std::vector<label_set>      point_ids_to_labels;
   tsl::robin_map<label, _u32> labels_to_number_of_points;
-  label_set all_labels;
-  
-  std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = parse_label_file(label_data_path, universal_label);
+  label_set                   all_labels;
+
+  std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) =
+      parse_label_file(label_data_path, universal_label);
 
   // 3. for each label, make a separate data file
   tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map;
   _u32 total_number_of_points = point_ids_to_labels.size();
-  
+
   if (data_type == "uint8")
-    label_id_to_orig_id_map = generate_label_specific_vector_files<uint8_t>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+    label_id_to_orig_id_map = generate_label_specific_vector_files<uint8_t>(
+        input_data_path, labels_to_number_of_points, point_ids_to_labels,
+        all_labels);
   else if (data_type == "int8")
-    label_id_to_orig_id_map = generate_label_specific_vector_files<int8_t>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+    label_id_to_orig_id_map = generate_label_specific_vector_files<int8_t>(
+        input_data_path, labels_to_number_of_points, point_ids_to_labels,
+        all_labels);
   else if (data_type == "float")
-    label_id_to_orig_id_map = generate_label_specific_vector_files<float>(input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels);
+    label_id_to_orig_id_map = generate_label_specific_vector_files<float>(
+        input_data_path, labels_to_number_of_points, point_ids_to_labels,
+        all_labels);
   else
     throw;
 
   // 4. for each created data file, create a vanilla diskANN index
   if (data_type == "uint8")
-    generate_label_indices<uint8_t>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+    generate_label_indices<uint8_t>(input_data_path, final_index_path_prefix,
+                                    all_labels, R, L, alpha, num_threads);
   else if (data_type == "int8")
-    generate_label_indices<int8_t>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+    generate_label_indices<int8_t>(input_data_path, final_index_path_prefix,
+                                   all_labels, R, L, alpha, num_threads);
   else if (data_type == "float")
-    generate_label_indices<float>(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, num_threads);
+    generate_label_indices<float>(input_data_path, final_index_path_prefix,
+                                  all_labels, R, L, alpha, num_threads);
   else
-   throw;
+    throw;
 
   // 5. "stitch" the indices together
   std::vector<std::vector<_u32>> stitched_graph;
-  tsl::robin_map<label, _u32> label_entry_points;
-  _u64 stitched_graph_size;
+  tsl::robin_map<label, _u32>    label_entry_points;
+  _u64                           stitched_graph_size;
 
   if (data_type == "uint8")
-    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<uint8_t>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+    std::tie(stitched_graph, stitched_graph_size) =
+        stitch_label_indices<uint8_t>(
+            final_index_path_prefix, total_number_of_points, all_labels,
+            labels_to_number_of_points, label_entry_points,
+            label_id_to_orig_id_map);
   else if (data_type == "int8")
-    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<int8_t>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+    std::tie(stitched_graph, stitched_graph_size) =
+        stitch_label_indices<int8_t>(
+            final_index_path_prefix, total_number_of_points, all_labels,
+            labels_to_number_of_points, label_entry_points,
+            label_id_to_orig_id_map);
   else if (data_type == "float")
-    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<float>(final_index_path_prefix, total_number_of_points, all_labels, labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map);
+    std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices<float>(
+        final_index_path_prefix, total_number_of_points, all_labels,
+        labels_to_number_of_points, label_entry_points,
+        label_id_to_orig_id_map);
   else
-   throw;
+    throw;
 
   // 5a. save the stitched graph to disk
-  save_full_index(final_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, universal_label, label_data_path);
+  save_full_index(final_index_path_prefix, input_data_path, stitched_graph_size,
+                  stitched_graph, label_entry_points, universal_label,
+                  label_data_path);
 
   // 6. run a prune on the stitched index, and save to disk
   if (data_type == "uint8")
-    prune_and_save<uint8_t>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
-                    label_entry_points, universal_label, label_data_path, num_threads);
+    prune_and_save<uint8_t>(final_index_path_prefix, input_data_path,
+                            stitched_graph, stitched_R, label_entry_points,
+                            universal_label, label_data_path, num_threads);
   else if (data_type == "int8")
-    prune_and_save<int8_t>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
-                    label_entry_points, universal_label, label_data_path, num_threads);
+    prune_and_save<int8_t>(final_index_path_prefix, input_data_path,
+                           stitched_graph, stitched_R, label_entry_points,
+                           universal_label, label_data_path, num_threads);
   else if (data_type == "float")
-    prune_and_save<float>(final_index_path_prefix, input_data_path, stitched_graph, stitched_R,
-                    label_entry_points, universal_label, label_data_path, num_threads);
+    prune_and_save<float>(final_index_path_prefix, input_data_path,
+                          stitched_graph, stitched_R, label_entry_points,
+                          universal_label, label_data_path, num_threads);
   else
-   throw;
-  
-  std::chrono::duration<double> index_time = std::chrono::high_resolution_clock::now() - index_timer;
-  std::cout << "pruned/stitched graph generated in " << index_time.count() << " seconds" << std::endl;
+    throw;
+
+  std::chrono::duration<double> index_time =
+      std::chrono::high_resolution_clock::now() - index_timer;
+  std::cout << "pruned/stitched graph generated in " << index_time.count()
+            << " seconds" << std::endl;
 
   clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels);
 }

From 586ccb21cbf115a320dca01920cb9189ecb93b78 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Fri, 3 Feb 2023 01:31:38 +0530
Subject: [PATCH 07/13] Add small fix to ensure compat. with scratch

---
 tests/build_stitched_index.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 496636d13..8ab389c11 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -655,7 +655,7 @@ void prune_and_save(path final_index_path_prefix, path input_data_path,
                           number_of_label_points, false, false);
 
   // not searching this index, set search_l to 0
-  index.load(final_index_path_prefix.c_str(), 0, 1);
+  index.load(final_index_path_prefix.c_str(), num_threads, 1);
 
   diskann::Parameters paras;
   paras.Set<unsigned>("R", stitched_R);

From 2e9cbb623d50b602170688f7c15434d344a0b8d4 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Fri, 3 Feb 2023 17:00:18 +0530
Subject: [PATCH 08/13] Update the label parsing to match new format

---
 tests/build_stitched_index.cpp | 35 +++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 8ab389c11..ad3dec0c6 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -2,12 +2,6 @@
 // email: t-gollapudis@microsoft.com
 
 #include <boost/program_options.hpp>
-#include <bits/types/struct_iovec.h>
-#include <fcntl.h>
-#include <omp.h>
-#include <sys/mman.h>
-#include <sys/uio.h>
-#include <unistd.h>
 #include <chrono>
 #include <cstdio>
 #include <cstring>
@@ -15,6 +9,12 @@
 #include <string>
 #include <tuple>
 
+#include <fcntl.h>
+#include <omp.h>
+#include <sys/mman.h>
+#include <sys/uio.h>
+#include <unistd.h>
+
 #include "index.h"
 #include "parameters.h"
 #include "utils.h"
@@ -125,11 +125,13 @@ size_t handle_args(int argc, char **argv, std::string &data_type,
 }
 
 /*
- * Parses the label datafile, which has the following line format:
- *            point_id -> \t -> comma-separated labels
+ * Parses the label datafile, which has comma-separated labels on
+ * each line. Line i corresponds to point id i.
  *
- * Returns three objects in via std::tuple:
- * 1. the label universe as a set
+ * Returns three objects via std::tuple:
+ * 1. map: key is point id, value is vector of labels said point has
+ * 2. map: key is label, value is number of points with the label
+ * 3. the label universe as a set
  */
 parse_label_file_return_values parse_label_file(path  label_data_path,
                                                 label universal_label) {
@@ -149,17 +151,15 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
   label_set                   all_labels;
 
   std::vector<_u32> points_with_universal_label;
+	line_cnt = 0;
   while (std::getline(label_data_stream, line)) {
-    std::istringstream current_point_and_labels(line);
+    std::istringstream current_labels_comma_separated(line);
     label_set          current_labels;
 
     // get point id
-    getline(current_point_and_labels, token, '\t');
-    _u32 point_id = (_u32) std::stoul(token);
+    _u32 point_id = line_cnt;
 
     // parse comma separated labels
-    getline(current_point_and_labels, token, '\t');
-    std::istringstream current_labels_comma_separated(token);
     bool               current_universal_label_check = false;
     while (getline(current_labels_comma_separated, token, ',')) {
       token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
@@ -182,6 +182,7 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
       exit(-1);
     }
     point_ids_to_labels[point_id] = current_labels;
+		line_cnt++;
   }
 
   // for every point with universal label, set its label set to all labels
@@ -257,6 +258,10 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
   for (const auto &label : all_labels) {
     iovec *label_iovecs =
         (iovec *) malloc(labels_to_number_of_points[label] * sizeof(iovec));
+		if (UNLIKELY(label_iovecs == nullptr)) {
+			close(input_data_fd);
+			throw;
+		}
     label_to_iovec_map[label] = label_iovecs;
     label_to_curr_iovec[label] = 0;
     label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);

From 91aa0019dfd9ae23cbdcf431f2c515bf5b27edd2 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Sun, 5 Feb 2023 17:07:53 +0530
Subject: [PATCH 09/13] Add Windows compatibility for stitched_vamana

---
 tests/build_stitched_index.cpp | 247 +++++++++++++++++++--------------
 1 file changed, 143 insertions(+), 104 deletions(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index ad3dec0c6..6c64c6dc5 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -9,21 +9,19 @@
 #include <string>
 #include <tuple>
 
-#include <fcntl.h>
 #include <omp.h>
-#include <sys/mman.h>
+#ifndef _WINDOWS
 #include <sys/uio.h>
-#include <unistd.h>
+#endif
 
 #include "index.h"
+#include "memory_mapper.h"
 #include "parameters.h"
 #include "utils.h"
 
 namespace po = boost::program_options;
 
 // macros
-#define LIKELY(condition) __builtin_expect(static_cast<bool>(condition), 1)
-#define UNLIKELY(condition) __builtin_expect(static_cast<bool>(condition), 0)
 #define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"
 #define PBWIDTH 60
 
@@ -107,7 +105,7 @@ size_t handle_args(int argc, char **argv, std::string &data_type,
                        "Input label file in txt format if present");
     desc.add_options()(
         "universal_label",
-        po::value<label>(&universal_label)->default_value(65537),
+        po::value<label>(&universal_label)->default_value(UINT32_MAX),
         "Universal label, if using it, only in conjunction with labels_file");
 
     po::variables_map vm;
@@ -151,7 +149,7 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
   label_set                   all_labels;
 
   std::vector<_u32> points_with_universal_label;
-	line_cnt = 0;
+  line_cnt = 0;
   while (std::getline(label_data_stream, line)) {
     std::istringstream current_labels_comma_separated(line);
     label_set          current_labels;
@@ -160,7 +158,7 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
     _u32 point_id = line_cnt;
 
     // parse comma separated labels
-    bool               current_universal_label_check = false;
+    bool current_universal_label_check = false;
     while (getline(current_labels_comma_separated, token, ',')) {
       token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
       token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
@@ -182,7 +180,7 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
       exit(-1);
     }
     point_ids_to_labels[point_id] = current_labels;
-		line_cnt++;
+    line_cnt++;
   }
 
   // for every point with universal label, set its label set to all labels
@@ -204,7 +202,9 @@ parse_label_file_return_values parse_label_file(path  label_data_path,
 /*
  * For each label, generates a file containing all vectors that have said label.
  * Also copies data from original bin file to new dimension-aligned file.
- * Utilizes platform-specific functions mmap and writev.
+ *
+ * Utilizes POSIX functions mmap and writev in order to minimize memory
+ * overhead, so we include an STL version as well.
  *
  * Each data file is saved under the following format:
  *    input_data_path + "_" + label
@@ -214,42 +214,21 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
     path                        input_data_path,
     tsl::robin_map<label, _u32> labels_to_number_of_points,
     std::vector<label_set> point_ids_to_labels, label_set all_labels) {
-  void *       input_memblock;
-  int          input_data_fd;
-  struct stat  input_data_st;
-  const size_t METADATA = 2 * sizeof(_u32);
-  auto         file_writing_timer = std::chrono::high_resolution_clock::now();
+  auto file_writing_timer = std::chrono::high_resolution_clock::now();
+  diskann::MemoryMapper input_data(input_data_path);
+  char *                input_start = input_data.getBuf();
 
-  // UNLIKELY moves the branch instr. body elsewhere, opening up icache space
-  // for code that actually executes (but this is probably irrelevant unless
-  // there is a massive number of labels)
-  input_data_fd = open(input_data_path.c_str(), O_RDONLY, (mode_t) 0444);
-  if (UNLIKELY(input_data_fd == -1))
-    throw;
-  if (UNLIKELY(fstat(input_data_fd, &input_data_st))) {
-    close(input_data_fd);
-    throw;
-  }
-
-  input_memblock = mmap(nullptr, input_data_st.st_size, PROT_READ, MAP_SHARED,
-                        input_data_fd, 0);
-  if (UNLIKELY(input_memblock == MAP_FAILED)) {
-    close(input_data_fd);
-    throw;
-  }
-
-  char *input_start = static_cast<char *>(input_memblock);
-  _u32  number_of_points, dimension;
+  _u32 number_of_points, dimension;
   std::memcpy(&number_of_points, input_start, sizeof(_u32));
   std::memcpy(&dimension, input_start + sizeof(_u32), sizeof(_u32));
-  const _u32 VECTOR_SIZE = dimension * sizeof(T);
+  const _u32   VECTOR_SIZE = dimension * sizeof(T);
+  const size_t METADATA = 2 * sizeof(_u32);
   if (number_of_points != point_ids_to_labels.size()) {
     std::cerr << "Error: number of points in labels file and data file differ."
               << std::endl;
     throw;
   }
 
-  // iovecs necessary for using writev
   tsl::robin_map<label, iovec *>           label_to_iovec_map;
   tsl::robin_map<label, _u32>              label_to_curr_iovec;
   tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id;
@@ -258,10 +237,9 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
   for (const auto &label : all_labels) {
     iovec *label_iovecs =
         (iovec *) malloc(labels_to_number_of_points[label] * sizeof(iovec));
-		if (UNLIKELY(label_iovecs == nullptr)) {
-			close(input_data_fd);
-			throw;
-		}
+    if (label_iovecs == nullptr) {
+      throw;
+    }
     label_to_iovec_map[label] = label_iovecs;
     label_to_curr_iovec[label] = 0;
     label_id_to_orig_id[label].reserve(labels_to_number_of_points[label]);
@@ -291,14 +269,13 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
     label_input_data_fd =
         open(curr_label_input_data_path.c_str(),
              O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t) 0644);
-    if (UNLIKELY(label_input_data_fd == -1))
+    if (label_input_data_fd == -1)
       throw;
 
     // write metadata
     _u32 metadata[2] = {curr_num_pts, dimension};
     int  return_value = write(label_input_data_fd, metadata, sizeof(_u32) * 2);
-    if (UNLIKELY(return_value == -1)) {
-      close(label_input_data_fd);
+    if (return_value == -1) {
       throw;
     }
 
@@ -309,7 +286,7 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
       return_value =
           writev(label_input_data_fd,
                  (label_to_iovec_map[label] + (IOV_MAX * i)), IOV_MAX);
-      if (UNLIKELY(return_value == -1)) {
+      if (return_value == -1) {
         close(label_input_data_fd);
         throw;
       }
@@ -319,7 +296,7 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
     return_value =
         writev(label_input_data_fd, (label_to_iovec_map[label] + (IOV_MAX * i)),
                curr_num_pts);
-    if (UNLIKELY(return_value == -1)) {
+    if (return_value == -1) {
       close(label_input_data_fd);
       throw;
     }
@@ -328,12 +305,82 @@ tsl::robin_map<label, std::vector<_u32>> generate_label_specific_vector_files(
     close(label_input_data_fd);
   }
 
-  // un-mmap (or munmap)
-  if (UNLIKELY(munmap(input_memblock, input_data_st.st_size) == -1)) {
-    close(input_data_fd);
+  std::chrono::duration<double> file_writing_time =
+      std::chrono::high_resolution_clock::now() - file_writing_timer;
+  std::cout << "generated " << all_labels.size()
+            << " label-specific vector files for index building in time "
+            << file_writing_time.count() << "\n"
+            << std::endl;
+
+  return label_id_to_orig_id;
+}
+
+// for use on systems without writev (i.e. Windows)
+template<typename T>
+tsl::robin_map<label, std::vector<_u32>>
+generate_label_specific_vector_files_compat(
+    path                        input_data_path,
+    tsl::robin_map<label, _u32> labels_to_number_of_points,
+    std::vector<label_set> point_ids_to_labels, label_set all_labels) {
+  auto          file_writing_timer = std::chrono::high_resolution_clock::now();
+  std::ifstream input_data_stream(input_data_path);
+
+  _u32 number_of_points, dimension;
+  input_data_stream.read((char *) &number_of_points, sizeof(_u32));
+  input_data_stream.read((char *) &dimension, sizeof(_u32));
+  const _u32 VECTOR_SIZE = dimension * sizeof(T);
+  if (number_of_points != point_ids_to_labels.size()) {
+    std::cerr << "Error: number of points in labels file and data file differ."
+              << std::endl;
     throw;
   }
-  close(input_data_fd);
+
+  tsl::robin_map<label, char *>            labels_to_vectors;
+  tsl::robin_map<label, _u32>              labels_to_curr_vector;
+  tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id;
+
+  for (const auto &label : all_labels) {
+    _u32  number_of_label_pts = labels_to_number_of_points[label];
+    char *vectors = (char *) malloc(number_of_label_pts * VECTOR_SIZE);
+    if (vectors == nullptr) {
+      throw;
+    }
+    labels_to_vectors[label] = vectors;
+    labels_to_curr_vector[label] = 0;
+    label_id_to_orig_id[label].reserve(number_of_label_pts);
+  }
+
+  for (_u32 point_id = 0; point_id < number_of_points; point_id++) {
+    char *curr_vector = (char *) malloc(VECTOR_SIZE);
+    input_data_stream.read(curr_vector, VECTOR_SIZE);
+    for (const auto &label : point_ids_to_labels[point_id]) {
+      char *curr_label_vector_ptr =
+          labels_to_vectors[label] +
+          (labels_to_curr_vector[label] * VECTOR_SIZE);
+      memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE);
+      labels_to_curr_vector[label]++;
+      label_id_to_orig_id[label].push_back(point_id);
+    }
+    free(curr_vector);
+  }
+
+  for (const auto &label : all_labels) {
+    path curr_label_input_data_path(input_data_path + "_" +
+                                    std::to_string(label));
+    _u32 number_of_label_pts = labels_to_number_of_points[label];
+
+    std::ofstream label_file_stream;
+    label_file_stream.exceptions(std::ios::badbit | std::ios::failbit);
+    label_file_stream.open(curr_label_input_data_path);
+
+    label_file_stream.write((char *) &number_of_label_pts, sizeof(_u32));
+    label_file_stream.write((char *) &dimension, sizeof(_u32));
+    label_file_stream.write((char *) labels_to_vectors[label],
+                            number_of_label_pts * VECTOR_SIZE);
+
+    free(labels_to_vectors[label]);
+    label_file_stream.close();
+  }
 
   std::chrono::duration<double> file_writing_time =
       std::chrono::high_resolution_clock::now() - file_writing_timer;
@@ -486,8 +533,8 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
     labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl;
   labels_to_medoids_writer.close();
 
-  // aux. file 4
-  if (universal_label != 65537) {
+  // aux. file 4 (only if we're using a universal label)
+  if (universal_label != UINT32_MAX) {
     std::ofstream universal_label_writer;
     universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit);
     universal_label_writer.open(final_index_path_prefix +
@@ -497,81 +544,46 @@ void save_full_index(path final_index_path_prefix, path input_data_path,
   }
 
   // main index
-  _u64 index_num_frozen_points = 0, index_num_edges = 0;
-  _u32 index_max_observed_degree = 0, index_entry_point = 0;
+  _u64         index_num_frozen_points = 0, index_num_edges = 0;
+  _u32         index_max_observed_degree = 0, index_entry_point = 0;
+  const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
   for (auto &point_neighbors : stitched_graph) {
     index_max_observed_degree =
         std::max(index_max_observed_degree, (_u32) point_neighbors.size());
   }
 
-  void *       memblock;
-  const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32);
-  struct stat  sb;
-  int          stitched_graph_fd;
-
-  stitched_graph_fd =
-      open(final_index_path_prefix.c_str(), O_RDWR | O_CREAT, (mode_t) 0644);
-  if (UNLIKELY(stitched_graph_fd == -1))
-    throw;
-  if (UNLIKELY(fstat(stitched_graph_fd, &sb))) {
-    close(stitched_graph_fd);
-    throw;
-  }
-
-  if (UNLIKELY(ftruncate(stitched_graph_fd, final_index_size) == -1)) {
-    close(stitched_graph_fd);
-    throw;
-  }
+  std::ofstream stitched_graph_writer;
+  stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit);
+  stitched_graph_writer.open(final_index_path_prefix);
 
-  // mmap could be faster than just doing a regular sequential write?
-  memblock = mmap(nullptr, final_index_size, PROT_READ | PROT_WRITE, MAP_SHARED,
-                  stitched_graph_fd, 0);
-  if (UNLIKELY(memblock == MAP_FAILED)) {
-    close(stitched_graph_fd);
-    throw;
-  }
+  stitched_graph_writer.write((char *) &final_index_size, sizeof(_u64));
+  stitched_graph_writer.write((char *) &index_max_observed_degree,
+                              sizeof(_u32));
+  stitched_graph_writer.write((char *) &index_entry_point, sizeof(_u32));
+  stitched_graph_writer.write((char *) &index_num_frozen_points, sizeof(_u64));
 
-  char *begin = static_cast<char *>(memblock);
-  std::memcpy(begin, &final_index_size, sizeof(_u64));
-  std::memcpy(begin + sizeof(_u64), &index_max_observed_degree, sizeof(_u32));
-  std::memcpy(begin + sizeof(_u64) + sizeof(_u32), &index_entry_point,
-              sizeof(_u32));
-  std::memcpy(begin + sizeof(_u64) + 2 * sizeof(_u32), &index_num_frozen_points,
-              sizeof(_u64));
   size_t bytes_written = METADATA;
   for (_u32 node_point = 0; node_point < stitched_graph.size(); node_point++) {
     _u32 current_node_num_neighbors = stitched_graph[node_point].size();
     std::vector<_u32> current_node_neighbors = stitched_graph[node_point];
-    std::memcpy(begin + bytes_written, &current_node_num_neighbors,
-                sizeof(_u32));
+    stitched_graph_writer.write((char *) &current_node_num_neighbors,
+                                sizeof(_u32));
     bytes_written += sizeof(_u32);
     for (const auto &current_node_neighbor : current_node_neighbors) {
-      std::memcpy(begin + bytes_written, &current_node_neighbor, sizeof(_u32));
+      stitched_graph_writer.write((char *) &current_node_neighbor,
+                                  sizeof(_u32));
       bytes_written += sizeof(_u32);
     }
     index_num_edges += current_node_num_neighbors;
   }
 
-  if (UNLIKELY(bytes_written != final_index_size)) {
+  if (bytes_written != final_index_size) {
     std::cerr << "Error: written bytes does not match allocated space"
               << std::endl;
     throw;
   }
 
-  if (UNLIKELY(msync(memblock, final_index_size, MS_SYNC) == -1)) {
-    std::cerr << "Error: could not sync memory-mapped writes to disk"
-              << std::endl;
-    close(stitched_graph_fd);
-    throw;
-  }
-
-  if (UNLIKELY(munmap(memblock, final_index_size) == -1)) {
-    std::cerr << "Error: could not release the memory-mapped file" << std::endl;
-    close(stitched_graph_fd);
-    throw;
-  }
-
-  close(stitched_graph_fd);
+  stitched_graph_writer.close();
 
   std::chrono::duration<double> saving_index_time =
       std::chrono::high_resolution_clock::now() - saving_index_timer;
@@ -643,6 +655,13 @@ stitch_indices_return_values stitch_label_indices(
   return std::make_tuple(stitched_graph, final_index_size);
 }
 
+/*
+ * Applies the prune_neighbors function from src/index.cpp to
+ * every node in the stitched graph.
+ *
+ * This is an optional step, hence the saving of both the full
+ * and pruned graph.
+ */
 template<typename T>
 void prune_and_save(path final_index_path_prefix, path input_data_path,
                     std::vector<std::vector<_u32>> stitched_graph,
@@ -733,6 +752,7 @@ int main(int argc, char **argv) {
   tsl::robin_map<label, std::vector<_u32>> label_id_to_orig_id_map;
   _u32 total_number_of_points = point_ids_to_labels.size();
 
+#ifndef _WINDOWS
   if (data_type == "uint8")
     label_id_to_orig_id_map = generate_label_specific_vector_files<uint8_t>(
         input_data_path, labels_to_number_of_points, point_ids_to_labels,
@@ -747,6 +767,25 @@ int main(int argc, char **argv) {
         all_labels);
   else
     throw;
+#else
+  if (data_type == "uint8")
+    label_id_to_orig_id_map =
+        generate_label_specific_vector_files_compat<uint8_t>(
+            input_data_path, labels_to_number_of_points, point_ids_to_labels,
+            all_labels);
+  else if (data_type == "int8")
+    label_id_to_orig_id_map =
+        generate_label_specific_vector_files_compat<int8_t>(
+            input_data_path, labels_to_number_of_points, point_ids_to_labels,
+            all_labels);
+  else if (data_type == "float")
+    label_id_to_orig_id_map =
+        generate_label_specific_vector_files_compat<float>(
+            input_data_path, labels_to_number_of_points, point_ids_to_labels,
+            all_labels);
+  else
+    throw;
+#endif
 
   // 4. for each created data file, create a vanilla diskANN index
   if (data_type == "uint8")

From 4d19df769a105e0a7c76c067d693f5bb1239cd6d Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Sun, 5 Feb 2023 18:27:28 +0530
Subject: [PATCH 10/13] Fix bug with new Neighbor definition

---
 src/index.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/index.cpp b/src/index.cpp
index 957cb6643..990a18cd4 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -1405,7 +1405,7 @@ namespace diskann {
                   _distance->compare(_data + _aligned_dim * (size_t) node,
                                      _data + _aligned_dim * (size_t) cur_nbr,
                                      (unsigned) _aligned_dim);
-              dummy_pool.emplace_back(Neighbor(cur_nbr, dist, true));
+              dummy_pool.emplace_back(Neighbor(cur_nbr, dist));
               dummy_visited.insert(cur_nbr);
             }
           }

From e9680639d7640856b0c1ed41529da0244f369c03 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Mon, 6 Feb 2023 16:35:27 +0530
Subject: [PATCH 11/13] Add license info to build_stitched_index and apply
 clang-format to index.h

---
 include/index.h                |  2 +-
 src/index.cpp                  | 17 ++++++++---------
 tests/build_stitched_index.cpp |  8 +++++---
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/include/index.h b/include/index.h
index 33d6bc240..ecdf2cac8 100644
--- a/include/index.h
+++ b/include/index.h
@@ -192,7 +192,7 @@ namespace diskann {
     DISKANN_DLLEXPORT consolidation_report
     consolidate_deletes(const Parameters &parameters);
 
-		DISKANN_DLLEXPORT void prune_all_nbrs(const Parameters &parameters);
+    DISKANN_DLLEXPORT void prune_all_nbrs(const Parameters &parameters);
 
     DISKANN_DLLEXPORT bool is_index_saved();
 
diff --git a/src/index.cpp b/src/index.cpp
index 990a18cd4..46017275f 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -483,14 +483,6 @@ namespace diskann {
     std::string labels_file = mem_index_file + "_labels.txt";
     std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt";
     if (file_exists(labels_file)) {
-      std::string universal_label_file(filename);
-      universal_label_file += "_universal_label.txt";
-      if (file_exists(universal_label_file)) {
-        std::ifstream universal_label_reader(universal_label_file);
-        universal_label_reader >> _universal_label;
-        _use_universal_label = true;
-        universal_label_reader.close();
-      }
       parse_label_file(labels_file);
       if (file_exists(labels_to_medoids)) {
         std::ifstream medoid_stream(labels_to_medoids);
@@ -522,7 +514,14 @@ namespace diskann {
         }
       }
 
-      
+			std::string universal_label_file(filename);
+      universal_label_file += "_universal_label.txt";
+      if (file_exists(universal_label_file)) {
+        std::ifstream universal_label_reader(universal_label_file);
+        universal_label_reader >> _universal_label;
+        _use_universal_label = true;
+        universal_label_reader.close();
+      }
     }
 
     if (!_save_as_one_file) {
diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 6c64c6dc5..19816d31b 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -1,5 +1,5 @@
-// author: Siddharth Gollapudi
-// email: t-gollapudis@microsoft.com
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
 
 #include <boost/program_options.hpp>
 #include <chrono>
@@ -106,7 +106,9 @@ size_t handle_args(int argc, char **argv, std::string &data_type,
     desc.add_options()(
         "universal_label",
         po::value<label>(&universal_label)->default_value(UINT32_MAX),
-        "Universal label, if using it, only in conjunction with labels_file");
+        "If a point comes with the specified universal label (and only the "
+        "univ. "
+        "label), then the point is considered to have every possible label");
 
     po::variables_map vm;
     po::store(po::parse_command_line(argc, argv, desc), vm);

From fd6c0d3998365ba0db6151eaa6c1179720125664 Mon Sep 17 00:00:00 2001
From: Siddharth Gollapudi <t-gollapudis@microsoft.com>
Date: Thu, 9 Feb 2023 03:27:45 +0530
Subject: [PATCH 12/13] Fix issue with --FilteredLbuild and stitched_vamana

---
 tests/build_stitched_index.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp
index 19816d31b..a714a51f5 100644
--- a/tests/build_stitched_index.cpp
+++ b/tests/build_stitched_index.cpp
@@ -408,8 +408,8 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix,
   diskann::Parameters label_index_build_parameters;
   label_index_build_parameters.Set<unsigned>("R", R);
   label_index_build_parameters.Set<unsigned>("L", L);
-  label_index_build_parameters.Set<unsigned>(
-      "C", 750);  // maximum candidate set size during prune
+  label_index_build_parameters.Set<unsigned>("C", 750);
+	label_index_build_parameters.Set<unsigned>("Lf", 0);
   label_index_build_parameters.Set<bool>("saturate_graph", 0);
   label_index_build_parameters.Set<float>("alpha", alpha);
   label_index_build_parameters.Set<unsigned>("num_threads", num_threads);

From 2f91e090cbc712005a12c34fd0289b2a1425d9cb Mon Sep 17 00:00:00 2001
From: Neelam Mahapatro <nmahapatro@microsoft.com>
Date: Thu, 9 Feb 2023 22:20:33 +0530
Subject: [PATCH 13/13] Fix comments

---
 include/disk_utils.h                |   2 +-
 include/index.h                     |   2 +-
 src/disk_utils.cpp                  |   7 +-
 src/index.cpp                       | 158 +++++++++++-----------
 src/pq_flash_index.cpp              | 194 +++++++++++++++-------------
 tests/search_disk_index.cpp         |   1 +
 tests/search_memory_index.cpp       |   1 +
 tests/utils/compute_groundtruth.cpp |   3 -
 tests/utils/stats_label_data.cpp    |  17 +--
 9 files changed, 201 insertions(+), 184 deletions(-)

diff --git a/include/disk_utils.h b/include/disk_utils.h
index dabeeeb0c..6aaa02426 100644
--- a/include/disk_utils.h
+++ b/include/disk_utils.h
@@ -75,7 +75,7 @@ namespace diskann {
                                      const _u64 nshards, unsigned max_degree,
                                      const std::string &output_vamana,
                                      const std::string &medoids_file, bool use_filters = false,
-      const std::string &labels_to_medoids_file = std::string(""));
+                                     const std::string &labels_to_medoids_file = std::string(""));
 
 DISKANN_DLLEXPORT void extract_shard_labels(
       const std::string &in_label_file, const std::string &shard_ids_bin,
diff --git a/include/index.h b/include/index.h
index ecdf2cac8..a866591bf 100644
--- a/include/index.h
+++ b/include/index.h
@@ -240,7 +240,7 @@ namespace diskann {
     // determines navigating node of the graph by calculating medoid of datafopt
     unsigned calculate_entry_point();
 
-    void parse_label_file(const std::string &map_file);
+    size_t parse_label_file(const std::string &map_file);
 
     std::pair<uint32_t, uint32_t> iterate_to_fixed_point(
         const T *node_coords, const unsigned Lindex,
diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp
index 9b7c3ca24..e6aa02fdb 100644
--- a/src/disk_utils.cpp
+++ b/src/disk_utils.cpp
@@ -323,6 +323,7 @@ namespace diskann {
       }
 
       std::ofstream mapping_writer(labels_to_medoids_file);
+      assert(mapping_writer.is_open());
       for (auto iter : global_label_to_medoids) {
         mapping_writer << iter.first << ", ";
         auto &vec = iter.second;
@@ -534,6 +535,7 @@ namespace diskann {
       diskann::cout << labels_per_point.size() << " is the new number of points"
                     << std::endl;
       std::ofstream label_writer(out_labels_file);
+      assert(label_writer.is_open());
       for (_u32 i = 0; i < labels_per_point.size(); i++) {
         for (_u32 j = 0; j < (labels_per_point[i].size() - 1); j++) {
           label_writer << labels_per_point[i][j] << ",";
@@ -551,6 +553,7 @@ namespace diskann {
       data = (T *) std::realloc((void *) data,
                                 labels_per_point.size() * ndims * sizeof(T));
       std::ofstream dummy_writer(out_metadata_file);
+      assert(dummy_writer.is_open());
       for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) {
         dummy_writer << i->first << "," << i->second << std::endl;
         std::memcpy(data + i->first * ndims, data + i->second * ndims,
@@ -566,7 +569,7 @@ namespace diskann {
       const std::string &in_label_file, const std::string &shard_ids_bin,
       const std::string &shard_label_file) {  // assumes ith row is for ith
                                               // point in labels file
-    std::cout << "Extracting labels for shard" << std::endl;
+    diskann::cout << "Extracting labels for shard" << std::endl;
 
     _u32 *ids = nullptr;
     _u64  num_ids, tmp_dim;
@@ -578,6 +581,8 @@ namespace diskann {
     std::ifstream label_reader(in_label_file);
     std::ofstream label_writer(shard_label_file);
 
+    assert(label_reader.is_open());
+    assert(label_writer.is_open());
     if (label_reader && label_writer) {
       while (std::getline(label_reader, cur_line)) {
         if (shard_counter >= num_ids) {
diff --git a/src/index.cpp b/src/index.cpp
index 46017275f..344ca86d6 100644
--- a/src/index.cpp
+++ b/src/index.cpp
@@ -279,9 +279,9 @@ namespace diskann {
         if (_filter_to_medoid_id.size() > 0) {
           std::ofstream medoid_writer(std::string(filename) +
                                       "_labels_to_medoids.txt");
+          assert(medoid_writer.is_open());
           for (auto iter : _filter_to_medoid_id) {
             medoid_writer << iter.first << ", " << iter.second << std::endl;
-            // std::cout << iter.first << ", " << iter.second << std::endl;
           }
           medoid_writer.close();
         }
@@ -289,12 +289,14 @@ namespace diskann {
         if (_use_universal_label) {
           std::ofstream universal_label_writer(std::string(filename) +
                                                "_universal_label.txt");
+          assert(universal_label_writer.is_open());
           universal_label_writer << _universal_label << std::endl;
           universal_label_writer.close();
         }
 
         if (_pts_to_labels.size() > 0) {
           std::ofstream label_writer(std::string(filename) + "_labels.txt");
+          assert(label_writer.is_open());
           for (_u32 i = 0; i < _pts_to_labels.size(); i++) {
             for (_u32 j = 0; j < (_pts_to_labels[i].size() - 1); j++) {
               label_writer << _pts_to_labels[i][j] << ",";
@@ -477,52 +479,12 @@ namespace diskann {
 
     _has_built = true;
 
-    size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0;
+    size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, label_num_pts = 0;
 
     std::string mem_index_file(filename);
     std::string labels_file = mem_index_file + "_labels.txt";
     std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt";
-    if (file_exists(labels_file)) {
-      parse_label_file(labels_file);
-      if (file_exists(labels_to_medoids)) {
-        std::ifstream medoid_stream(labels_to_medoids);
-
-        std::string line, token;
-        unsigned    line_cnt = 0;
-
-        _filter_to_medoid_id.clear();
-
-        while (std::getline(medoid_stream, line)) {
-          std::istringstream iss(line);
-          _u32               cnt = 0;
-          _u32               medoid = 0;
-          label              label;
-          while (std::getline(iss, token, ',')) {
-            token.erase(std::remove(token.begin(), token.end(), '\n'),
-                        token.end());
-            token.erase(std::remove(token.begin(), token.end(), '\r'),
-                        token.end());
-            unsigned token_as_num = std::stoul(token);
-            if (cnt == 0)
-              label = token_as_num;
-            else
-              medoid = token_as_num;
-            cnt++;
-          }
-          _filter_to_medoid_id[label] = medoid;
-          line_cnt++;
-        }
-      }
-
-			std::string universal_label_file(filename);
-      universal_label_file += "_universal_label.txt";
-      if (file_exists(universal_label_file)) {
-        std::ifstream universal_label_reader(universal_label_file);
-        universal_label_reader >> _universal_label;
-        _use_universal_label = true;
-        universal_label_reader.close();
-      }
-    }
+    
 
     if (!_save_as_one_file) {
       // For DLVS Store, we will not support saving the index in multiple files.
@@ -561,6 +523,53 @@ namespace diskann {
       throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__,
                                   __LINE__);
     }
+    if (file_exists(labels_file)) {
+      label_num_pts = parse_label_file(labels_file);
+      assert(label_num_pts == data_file_num_pts);
+      if (file_exists(labels_to_medoids)) {
+        std::ifstream medoid_stream(labels_to_medoids);
+        assert(medoid_stream.is_open());        
+        std::string line, token;
+        unsigned    line_cnt = 0;
+
+        _filter_to_medoid_id.clear();
+        try {
+          while (std::getline(medoid_stream, line)) {
+            std::istringstream iss(line);
+            _u32               cnt = 0;
+            _u32               medoid = 0;
+            label              label;
+            while (std::getline(iss, token, ',')) {
+              token.erase(std::remove(token.begin(), token.end(), '\n'),
+                          token.end());
+              token.erase(std::remove(token.begin(), token.end(), '\r'),
+                          token.end());
+              unsigned token_as_num = std::stoul(token);
+              if (cnt == 0)
+                label = token_as_num;
+              else
+                medoid = token_as_num;
+              cnt++;
+            }
+            _filter_to_medoid_id[label] = medoid;
+            line_cnt++;
+          }
+        } catch (std::system_error &e) {
+          throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__,
+                              __LINE__);
+        }
+      }
+
+      std::string universal_label_file(filename);
+      universal_label_file += "_universal_label.txt";
+      if (file_exists(universal_label_file)) {
+        std::ifstream universal_label_reader(universal_label_file);
+        assert(universal_label_reader.is_open());
+        universal_label_reader >> _universal_label;
+        _use_universal_label = true;
+        universal_label_reader.close();
+      }
+    }
 
     _nd = data_file_num_pts - _num_frozen_pts;
     _empty_slots.clear();
@@ -1707,13 +1716,13 @@ namespace diskann {
   }
 
   template<typename T, typename TagT>
-  void Index<T, TagT>::parse_label_file(const std::string &map_file) {
+  size_t Index<T, TagT>::parse_label_file(const std::string &label_file) {
     // Format of Label txt file: filters with comma separators
-
-    std::ifstream infile(map_file);
+    std::ifstream infile(label_file);
+    assert(infile.is_open());
     std::string   line, token;
     unsigned      line_cnt = 0;
-
+    
     while (std::getline(infile, line)) {
       line_cnt++;
     }
@@ -1722,28 +1731,35 @@ namespace diskann {
     infile.clear();
     infile.seekg(0, std::ios::beg);
     line_cnt = 0;
-    while (std::getline(infile, line)) {
-      std::istringstream iss(line);
-      std::vector<label> lbls(0);
-      getline(iss, token, '\t');
-      std::istringstream new_iss(token);
-      while (getline(new_iss, token, ',')) {
-        token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-        token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-        unsigned token_as_num = std::stoul(token);
-        lbls.push_back(token_as_num);
-        _labels.insert(token_as_num);
-      }
-      if (lbls.size() <= 0) {
-        std::cout << "No label found";
-        exit(-1);
+    try {
+      while (std::getline(infile, line)) {
+        std::istringstream iss(line);
+        std::vector<label> lbls(0);
+        getline(iss, token, '\t');
+        std::istringstream new_iss(token);
+        while (getline(new_iss, token, ',')) {
+          token.erase(std::remove(token.begin(), token.end(), '\n'),
+                      token.end());
+          token.erase(std::remove(token.begin(), token.end(), '\r'),
+                      token.end());
+          unsigned token_as_num = std::stoul(token);
+          lbls.push_back(token_as_num);
+          _labels.insert(token_as_num);
+        }
+        if (lbls.size() <= 0) {
+          diskann::cout << "No label found";
+          exit(-1);
+        }
+        std::sort(lbls.begin(), lbls.end());
+        _pts_to_labels[line_cnt] = lbls;
+        line_cnt++;
       }
-      std::sort(lbls.begin(), lbls.end());
-      _pts_to_labels[line_cnt] = lbls;
-      line_cnt++;
+      diskann::cout << "Identified " << _labels.size() << " distinct label(s)"
+                    << std::endl;
+    } catch (std::system_error &e) {
+      throw FileException(label_file, e, __FUNCSIG__, __FILE__, __LINE__);
     }
-    std::cout << "Identified " << _labels.size() << " distinct label(s)"
-              << std::endl;
+    return line_cnt;
   }
 
   template<typename T, typename TagT>
@@ -1761,9 +1777,9 @@ namespace diskann {
     _labels_file = label_file;
     _filtered_index = true;
     _filter_to_medoid_id.clear();
-    parse_label_file(label_file);  // determines medoid for each label and
+    size_t label_num_pts = parse_label_file(label_file);  // determines medoid for each label and
                                    // identifies the points to label mapping
-
+    assert(label_num_pts != num_points_to_load);
     _u32 counter = 0;
 #pragma omp parallel for schedule(dynamic, 1)
     for (int lbl = 0; lbl < _labels.size(); lbl++) {
@@ -1803,15 +1819,13 @@ namespace diskann {
           _filter_to_medoid_id[x] = best_medoid;
           _medoid_counts[best_medoid]++;
           std::stringstream a;
-          // a << "Medoid of " << x << " is " << best_medoid <<
-          // std::endl; std::cout << a.str();
         }
       }
 #pragma omp critical
       counter++;
       std::stringstream a;
       a << ((100.0 * counter) / _labels.size()) << "\% processed \r";
-      std::cout << a.str() << std::flush;
+      diskann::cout << a.str() << std::flush;
     }
 
     this->build(filename, num_points_to_load, parameters, tags);
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
index fe326a10e..a86c14869 100644
--- a/src/pq_flash_index.cpp
+++ b/src/pq_flash_index.cpp
@@ -464,7 +464,7 @@ namespace diskann {
       num_pts++;
     }
 
-    std::cout << "Labels file metadata: num_points: " << num_pts
+    diskann::cout << "Labels file metadata: num_points: " << num_pts
               << ", #total_labels: " << num_total_labels << std::endl;
     infile.close();
   }
@@ -484,54 +484,64 @@ namespace diskann {
     return ret_val;
   }
   template<typename T>
-  void PQFlashIndex<T>::parse_label_file(const std::string &map_file) {
+  void PQFlashIndex<T>::parse_label_file(const std::string &label_file) {
     //_filtered_ann = 1;
 
-    std::ifstream infile(map_file);
+    std::ifstream infile(label_file);
+    
     std::string   line, token;
     _u32          line_cnt = 0;
 
     _u32 num_pts_in_label_file;
     _u32 num_total_labels;
-    get_label_file_metadata(map_file, num_pts_in_label_file, num_total_labels);
-
+    get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels);
+    assert(num_pts_in_label_file == this->num_points);
+    
     _pts_to_label_offsets = new _u32[num_pts_in_label_file];
     _pts_to_labels = new _u32[num_pts_in_label_file + num_total_labels];
     _u32 counter = 0;
-
-    while (std::getline(infile, line)) {
-      std::istringstream iss(line);
-      std::vector<_u32>  lbls(0);
-
-      _pts_to_label_offsets[line_cnt] = counter;
-      _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter];
-      num_lbls_in_cur_pt = 0;
-      counter++;
-      getline(iss, token, '\t');
-      std::istringstream new_iss(token);
-      while (getline(new_iss, token, ',')) {
-        token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-        token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-        if (_labels.find(token) == _labels.end()) {
-          _filter_list.emplace_back(token);
+    infile.clear();
+    infile.seekg(0, std::ios::beg);
+    assert(infile.is_open());
+    try {
+      while (std::getline(infile, line)) {
+        std::istringstream iss(line);
+        std::vector<_u32>  lbls(0);
+
+        _pts_to_label_offsets[line_cnt] = counter;
+        _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter];
+        num_lbls_in_cur_pt = 0;
+        counter++;
+        getline(iss, token, '\t');
+        std::istringstream new_iss(token);
+        while (getline(new_iss, token, ',')) {
+          token.erase(std::remove(token.begin(), token.end(), '\n'),
+                      token.end());
+          token.erase(std::remove(token.begin(), token.end(), '\r'),
+                      token.end());
+          if (_labels.find(token) == _labels.end()) {
+            _filter_list.emplace_back(token);
+          }
+          int32_t filter_num = get_filter_number(token);
+          if (filter_num == -1) {
+            diskann::cout << "Error!! " << std::endl;
+            exit(-1);
+          }
+          _pts_to_labels[counter++] = filter_num;
+          num_lbls_in_cur_pt++;
+          _labels.insert(token);
         }
-        int32_t filter_num = get_filter_number(token);
-        if (filter_num == -1) {
-          std::cout << "Error!! " << std::endl;
+
+        if (num_lbls_in_cur_pt == 0) {
+          diskann::cout << "No label found for point " << line_cnt << std::endl;
           exit(-1);
         }
-        _pts_to_labels[counter++] = filter_num;
-        num_lbls_in_cur_pt++;
-        _labels.insert(token);
+        line_cnt++;
       }
-
-      if (num_lbls_in_cur_pt == 0) {
-        std::cout << "No label found for point " << line_cnt << std::endl;
-        exit(-1);
-      }
-      line_cnt++;
+    }catch (std::system_error &e) {
+      throw FileException(label_file, e, __FUNCSIG__, __FILE__,
+                          __LINE__);
     }
-    // MallocExtension::instance()->ReleaseFreeMemory();
     infile.close();
   }
 
@@ -539,7 +549,7 @@ namespace diskann {
   void PQFlashIndex<T>::set_universal_label(const std::string &label) {
     int32_t temp_filter_num = get_filter_number(label);
     if (temp_filter_num == -1) {
-      std::cout << "Error, could not find universal label. Exitting."
+      diskann::cout << "Error, could not find universal label. Exitting."
                 << std::endl;
       exit(-1);
     } else {
@@ -595,29 +605,70 @@ namespace diskann {
         std ::string(disk_index_file) + "_labels_to_medoids.txt";
     std::string dummy_map_file =
         std ::string(disk_index_file) + "_dummy_map.txt";
+
+    size_t pq_file_dim, pq_file_num_centroids;
+#ifdef EXEC_ENV_OLS
+    get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim,
+                     METADATA_SIZE);
+#else
+    get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim,
+                     METADATA_SIZE);
+#endif
+
+    this->disk_index_file = disk_index_file;
+
+    if (pq_file_num_centroids != 256) {
+      diskann::cout << "Error. Number of PQ centroids is not 256. Exiting."
+                    << std::endl;
+      return -1;
+    }
+
+    this->data_dim = pq_file_dim;
+    // will reset later if we use PQ on disk
+    this->disk_data_dim = this->data_dim;
+    // will change later if we use PQ on disk or if we are using
+    // inner product without PQ
+    this->disk_bytes_per_point = this->data_dim * sizeof(T);
+    this->aligned_dim = ROUND_UP(pq_file_dim, 8);
+
+    size_t npts_u64, nchunks_u64;
+#ifdef EXEC_ENV_OLS
+    diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64,
+                           nchunks_u64);
+#else
+    diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64,
+                           nchunks_u64);
+#endif
+
+    this->num_points = npts_u64;
+    this->n_chunks = nchunks_u64;
+
     if (file_exists(labels_file)) {
       parse_label_file(labels_file);
       if (file_exists(labels_to_medoids)) {
         std::ifstream medoid_stream(labels_to_medoids);
-
+        assert(medoid_stream.is_open());
         std::string line, token;
-        //        unsigned    line_cnt = 0;
-
         _filter_to_medoid_id.clear();
 
-        while (std::getline(medoid_stream, line)) {
-          std::istringstream iss(line);
-          _u32               cnt = 0;
-          _u32               medoid;
-          std::string        label;
-          while (std::getline(iss, token, ',')) {
-            if (cnt == 0)
-              label = token;
-            else
-              medoid = (_u32) stoul(token);
-            cnt++;
+        try {
+          while (std::getline(medoid_stream, line)) {
+            std::istringstream iss(line);
+            _u32               cnt = 0;
+            _u32               medoid;
+            std::string        label;
+            while (std::getline(iss, token, ',')) {
+              if (cnt == 0)
+                label = token;
+              else
+                medoid = (_u32) stoul(token);
+              cnt++;
+            }
+            _filter_to_medoid_id[label] = medoid;
           }
-          _filter_to_medoid_id[label] = medoid;
+        } catch (std::system_error &e) {
+          throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__,
+                              __LINE__);
         }
       }
       std::string univ_label_file =
@@ -633,8 +684,6 @@ namespace diskann {
         std::ifstream dummy_map_stream(dummy_map_file);
 
         std::string line, token;
-        //        unsigned    line_cnt = 0;
-
         while (std::getline(dummy_map_stream, line)) {
           std::istringstream iss(line);
           _u32               cnt = 0;
@@ -657,47 +706,10 @@ namespace diskann {
           _real_to_dummy_map[real_id].emplace_back(dummy_id);
         }
         dummy_map_stream.close();
-        std::cout << "Loaded dummy map" << std::endl;
+        diskann::cout << "Loaded dummy map" << std::endl;
       }
     }
 
-    size_t pq_file_dim, pq_file_num_centroids;
-#ifdef EXEC_ENV_OLS
-    get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim,
-                     METADATA_SIZE);
-#else
-    get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim,
-                     METADATA_SIZE);
-#endif
-
-    this->disk_index_file = disk_index_file;
-
-    if (pq_file_num_centroids != 256) {
-      diskann::cout << "Error. Number of PQ centroids is not 256. Exiting."
-                    << std::endl;
-      return -1;
-    }
-
-    this->data_dim = pq_file_dim;
-    // will reset later if we use PQ on disk
-    this->disk_data_dim = this->data_dim;
-    // will change later if we use PQ on disk or if we are using
-    // inner product without PQ
-    this->disk_bytes_per_point = this->data_dim * sizeof(T);
-    this->aligned_dim = ROUND_UP(pq_file_dim, 8);
-
-    size_t npts_u64, nchunks_u64;
-#ifdef EXEC_ENV_OLS
-    diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64,
-                           nchunks_u64);
-#else
-    diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64,
-                           nchunks_u64);
-#endif
-
-    this->num_points = npts_u64;
-    this->n_chunks = nchunks_u64;
-
 #ifdef EXEC_ENV_OLS
     pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64);
 #else
@@ -735,7 +747,7 @@ namespace diskann {
       disk_bytes_per_point =
           disk_pq_n_chunks *
           sizeof(_u8);  // revising disk_bytes_per_point since DISK PQ is used.
-      std::cout << "Disk index uses PQ data compressed down to "
+      diskann::cout << "Disk index uses PQ data compressed down to "
                 << disk_pq_n_chunks << " bytes per point." << std::endl;
     }
 
diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp
index a765669ef..2825ac1cc 100644
--- a/tests/search_disk_index.cpp
+++ b/tests/search_disk_index.cpp
@@ -425,4 +425,5 @@ int main(int argc, char** argv) {
     diskann::cerr << "Index search failed." << std::endl;
     return -1;
   }
+
 }
\ No newline at end of file
diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp
index c0e3c094a..153662988 100644
--- a/tests/search_memory_index.cpp
+++ b/tests/search_memory_index.cpp
@@ -357,4 +357,5 @@ int main(int argc, char** argv) {
     diskann::cerr << "Index search failed." << std::endl;
     return -1;
   }
+
 }
\ No newline at end of file
diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp
index d57eb70c3..5de1e7351 100644
--- a/tests/utils/compute_groundtruth.cpp
+++ b/tests/utils/compute_groundtruth.cpp
@@ -327,14 +327,11 @@ inline std::vector<size_t> load_filtered_bin_as_float(
   std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl;
   reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t),
                std::ios::beg);
-  //    data = new T[nptsuint64_t * ndimsuint64_t];
   T *data_T = new T[nptsuint64_t * ndimsuint64_t];
   reader.read((char *) data_T, sizeof(T) * nptsuint64_t * ndimsuint64_t);
   std::cout << "Finished reading part of the bin file." << std::endl;
   reader.close();
-  //  data =  (nptsuint64_t*ndimsuint64_t, ALIGNMENT);
   data = aligned_malloc<float>(nptsuint64_t * ndimsuint64_t, ALIGNMENT);
-  //#pragma omp parallel for schedule(dynamic, 32768)
   for (int64_t i = 0; i < (int64_t) nptsuint64_t; i++) {
     if (std::find(pts_to_labels[start_id + i].begin(),
                   pts_to_labels[start_id + i].end(),
diff --git a/tests/utils/stats_label_data.cpp b/tests/utils/stats_label_data.cpp
index 3287179c3..b983b02d3 100644
--- a/tests/utils/stats_label_data.cpp
+++ b/tests/utils/stats_label_data.cpp
@@ -1,8 +1,5 @@
-//
-// Created by Neel Karia on 10/01/2021.
-//
-// Generate base file and groundtruth file for points belonging to a particular
-// label
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
 
 #include <iostream>
 #include <fstream>
@@ -33,7 +30,6 @@ void stats_analysis(const std::string labels_file, std::string univeral_label,
                     _u32 density = 10) {
   std::string   token, line;
   std::ifstream labels_stream(labels_file);
-  // int           cnt[51] = {0};
   std::unordered_map<std::string, _u32> label_counts;
   std::string                           label_with_max_points;
   _u32                                  max_points = 0;
@@ -68,20 +64,11 @@ void stats_analysis(const std::string labels_file, std::string univeral_label,
             << (float) dense_pts / (float) labels_per_point.size() << std::endl;
   std::sort(labels_per_point.begin(), labels_per_point.end());
 
-  /*  for (_u32 p = 90; p < 100; p += 1) {
-      std::cout << labels_per_point[(p * labels_per_point.size()) / 100.0]
-                << std::endl;
-    }
-    std::cout << labels_per_point[labels_per_point.size() - 1] << std::endl;
-  */
-
   std::vector<std::pair<std::string, _u32>> label_count_vec;
 
   for (auto it = label_counts.begin(); it != label_counts.end(); it++) {
     auto& lbl = *it;
     label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second));
-    //    std::cout << "Label " << lbl.first << " has " << lbl.second
-    //              << " data points" << std::endl;
     if (lbl.second > max_points) {
       max_points = lbl.second;
       label_with_max_points = lbl.first;