Stats tweaks and KNN benchmark

1 files changed, 5 insertions, 1 deletions

diff --git a/benchmarks/include/file_util.h b/benchmarks/include/file_util.h
index a159dbe..39b297b 100644
--- a/benchmarks/include/file_util.h
+++ b/benchmarks/include/file_util.h
@@ -149,7 +149,7 @@ static std::vector<QP> read_binary_knn_queries(std::string fname, size_t k,
 }
 
 template <typename QP>
-static std::vector<QP> read_knn_queries(std::string fname, size_t k) {
+static std::vector<QP> read_knn_queries(std::string fname, size_t k, size_t n=0) {
   std::vector<QP> queries;
 
   FILE *qf = fopen(fname.c_str(), "r");
@@ -173,6 +173,10 @@ static std::vector<QP> read_knn_queries(std::string fname, size_t k) {
 
     query.k = k;
     queries.emplace_back(query);
+
+    if (n != 0 && queries.size() == n) {
+      break;
+    }
   }
 
   free(line);