fix to reset stream after reads

microsoft · Oct 12, 2023 · 4031bc2 · 4031bc2
1 parent 381bae6
commit 4031bc2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h
@@ -122,6 +122,7 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
                                                    uint32_t &num_total_labels);
     DISKANN_DLLEXPORT void generate_random_labels(std::vector<LabelT> &labels, const uint32_t num_labels,
                                                   const uint32_t nthreads);
+    void reset_stream_for_reading(std::basic_istream<char> &infile);
 
     // sector # on disk where node_id is present with in the graph part
     DISKANN_DLLEXPORT uint64_t get_node_sector(uint64_t node_id);

diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
@@ -602,6 +602,13 @@ LabelT PQFlashIndex<T, LabelT>::get_converted_label(const std::string &filter_la
     throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__);
 }
 
+template <typename T, typename LabelT>
+void PQFlashIndex<T, LabelT>::reset_stream_for_reading(std::basic_istream<char> &infile)
+{
+    infile.clear();
+    infile.seekg(0);
+}
+
 template <typename T, typename LabelT>
 void PQFlashIndex<T, LabelT>::get_label_file_metadata(std::basic_istream<char> &infile, uint32_t &num_pts,
                                                       uint32_t &num_total_labels)
@@ -624,7 +631,7 @@ void PQFlashIndex<T, LabelT>::get_label_file_metadata(std::basic_istream<char> &
 
     diskann::cout << "Labels file metadata: num_points: " << num_pts << ", #total_labels: " << num_total_labels
                   << std::endl;
-    infile.seekg(0);
+    reset_stream_for_reading(infile);
 }
 
 template <typename T, typename LabelT>
@@ -685,8 +692,8 @@ void PQFlashIndex<T, LabelT>::parse_label_file(std::basic_istream<char> &infile,
         }
         line_cnt++;
     }
-    infile.seekg(0);
     num_points_labels = line_cnt;
+    reset_stream_for_reading(infile);
 }
 
 template <typename T, typename LabelT> void PQFlashIndex<T, LabelT>::set_universal_label(const LabelT &label)