read file in one time (#460)

* read whole label file to memory, use string find instead stringstream

* format doc

Author: Sanhaoji2

AnyBuildLogs/latest.txt

@@ -0,0 +1 @@
+20231019-111207-d314f8bf

include/pq_flash_index.h

@@ -118,7 +118,7 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
     DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, LabelT label_id);
     std::unordered_map<std::string, LabelT> load_label_map(std::basic_istream<char> &infile);
     DISKANN_DLLEXPORT void parse_label_file(std::basic_istream<char> &infile, size_t &num_pts_labels);
-    DISKANN_DLLEXPORT void get_label_file_metadata(std::basic_istream<char> &infile, uint32_t &num_pts,
+    DISKANN_DLLEXPORT void get_label_file_metadata(const std::string &fileContent, uint32_t &num_pts,
                                                    uint32_t &num_total_labels);
     DISKANN_DLLEXPORT void generate_random_labels(std::vector<LabelT> &labels, const uint32_t num_labels,
                                                   const uint32_t nthreads);

src/pq_flash_index.cpp

@@ -610,28 +610,47 @@ void PQFlashIndex<T, LabelT>::reset_stream_for_reading(std::basic_istream<char>
 }
 
 template <typename T, typename LabelT>
-void PQFlashIndex<T, LabelT>::get_label_file_metadata(std::basic_istream<char> &infile, uint32_t &num_pts,
+void PQFlashIndex<T, LabelT>::get_label_file_metadata(const std::string &fileContent, uint32_t &num_pts,
                                                       uint32_t &num_total_labels)
 {
-    std::string line, token;
     num_pts = 0;
     num_total_labels = 0;
 
-    while (std::getline(infile, line))
+    size_t file_size = fileContent.length();
+
+    std::string label_str;
+    size_t cur_pos = 0;
+    size_t next_pos = 0;
+    while (cur_pos < file_size && cur_pos != std::string::npos)
     {
-        std::istringstream iss(line);
-        while (getline(iss, token, ','))
+        next_pos = fileContent.find('\n', cur_pos);
+        if (next_pos == std::string::npos)
+        {
+            break;
+        }
+
+        size_t lbl_pos = cur_pos;
+        size_t next_lbl_pos = 0;
+        while (lbl_pos < next_pos && lbl_pos != std::string::npos)
         {
-            token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-            token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
+            next_lbl_pos = fileContent.find(',', lbl_pos);
+            if (next_lbl_pos == std::string::npos) // the last label
+            {
+                next_lbl_pos = next_pos;
+            }
+
             num_total_labels++;
+
+            lbl_pos = next_lbl_pos + 1;
         }
+
+        cur_pos = next_pos + 1;
+
         num_pts++;
     }
 
     diskann::cout << "Labels file metadata: num_points: " << num_pts << ", #total_labels: " << num_total_labels
                   << std::endl;
-    reset_stream_for_reading(infile);
 }
 
 template <typename T, typename LabelT>
@@ -654,44 +673,82 @@ inline bool PQFlashIndex<T, LabelT>::point_has_label(uint32_t point_id, LabelT l
 template <typename T, typename LabelT>
 void PQFlashIndex<T, LabelT>::parse_label_file(std::basic_istream<char> &infile, size_t &num_points_labels)
 {
-    std::string line, token;
+    infile.seekg(0, std::ios::end);
+    size_t file_size = infile.tellg();
+
+    std::string buffer(file_size, ' ');
+
+    infile.seekg(0, std::ios::beg);
+    infile.read(&buffer[0], file_size);
+
+    std::string line;
     uint32_t line_cnt = 0;
 
     uint32_t num_pts_in_label_file;
     uint32_t num_total_labels;
-    get_label_file_metadata(infile, num_pts_in_label_file, num_total_labels);
+    get_label_file_metadata(buffer, num_pts_in_label_file, num_total_labels);
 
     _pts_to_label_offsets = new uint32_t[num_pts_in_label_file];
     _pts_to_label_counts = new uint32_t[num_pts_in_label_file];
     _pts_to_labels = new LabelT[num_total_labels];
     uint32_t labels_seen_so_far = 0;
 
-    while (std::getline(infile, line))
+    std::string label_str;
+    size_t cur_pos = 0;
+    size_t next_pos = 0;
+    while (cur_pos < file_size && cur_pos != std::string::npos)
     {
-        std::istringstream iss(line);
-        std::vector<uint32_t> lbls(0);
+        next_pos = buffer.find('\n', cur_pos);
+        if (next_pos == std::string::npos)
+        {
+            break;
+        }
 
         _pts_to_label_offsets[line_cnt] = labels_seen_so_far;
         uint32_t &num_lbls_in_cur_pt = _pts_to_label_counts[line_cnt];
         num_lbls_in_cur_pt = 0;
-        getline(iss, token, '\t');
-        std::istringstream new_iss(token);
-        while (getline(new_iss, token, ','))
+
+        size_t lbl_pos = cur_pos;
+        size_t next_lbl_pos = 0;
+        while (lbl_pos < next_pos && lbl_pos != std::string::npos)
         {
-            token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
-            token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-            LabelT token_as_num = (LabelT)std::stoul(token);
+            next_lbl_pos = buffer.find(',', lbl_pos);
+            if (next_lbl_pos == std::string::npos) // the last label in the whole file
+            {
+                next_lbl_pos = next_pos;
+            }
+
+            if (next_lbl_pos > next_pos) // the last label in one line, just read to the end
+            {
+                next_lbl_pos = next_pos;
+            }
+
+            label_str.assign(buffer.c_str() + lbl_pos, next_lbl_pos - lbl_pos);
+            if (label_str[label_str.length() - 1] == '\t') // '\t' won't exist in label file?
+            {
+                label_str.erase(label_str.length() - 1);
+            }
+
+            LabelT token_as_num = (LabelT)std::stoul(label_str);
             _pts_to_labels[labels_seen_so_far++] = (LabelT)token_as_num;
             num_lbls_in_cur_pt++;
+
+            // move to next label
+            lbl_pos = next_lbl_pos + 1;
         }
 
+        // move to next line
+        cur_pos = next_pos + 1;
+
         if (num_lbls_in_cur_pt == 0)
         {
             diskann::cout << "No label found for point " << line_cnt << std::endl;
             exit(-1);
         }
+
         line_cnt++;
     }
+
     num_points_labels = line_cnt;
     reset_stream_for_reading(infile);
 }
@@ -784,7 +841,7 @@ int PQFlashIndex<T, LabelT>::load_from_separate_paths(uint32_t num_threads, cons
 #else
     if (file_exists(labels_file))
     {
-        std::ifstream infile(labels_file);
+        std::ifstream infile(labels_file, std::ios::binary);
         if (infile.fail())
         {
             throw diskann::ANNException(std::string("Failed to open file ") + labels_file, -1);