feat(search): Multishard cutoffs

.github/actions/regression-tests/action.yml

@@ -26,7 +26,7 @@ runs:
   # timeout-minutes: 20
   steps:
     - name: Run PyTests
-      id: first
+      id: main
       shell: bash
       run: |
         ls -l ${GITHUB_WORKSPACE}/
@@ -37,7 +37,7 @@ runs:
         export DRAGONFLY_PATH="${GITHUB_WORKSPACE}/${{inputs.build-folder-name}}/${{inputs.dfly-executable}}"
         export UBSAN_OPTIONS=print_stacktrace=1:halt_on_error=1 # to crash on errors
 
-        timeout 20m pytest -m "${{inputs.filter}}" --durations=10 --color=yes --json-report --json-report-file=report.json dragonfly --ignore=dragonfly/replication_test.py --log-cli-level=INFO || code=$?
+        timeout 40m pytest -m "${{inputs.filter}}" --durations=10 --color=yes --json-report --json-report-file=report.json dragonfly --log-cli-level=INFO || code=$?
 
         # timeout returns 124 if we exceeded the timeout duration
         if [[ $code -eq 124 ]]; then
@@ -50,32 +50,6 @@ runs:
           exit 1
         fi
 
-    - name: Run PyTests replication test
-      id: second
-      if: ${{ inputs.run-only-on-ubuntu-latest == 'true' || (inputs.run-only-on-ubuntu-latest == 'false' && matrix.runner == 'ubuntu-latest') }}
-      shell: bash
-      run: |
-        echo "Running PyTests replication test"
-        cd ${GITHUB_WORKSPACE}/tests
-        # used by PyTests
-        export DRAGONFLY_PATH="${GITHUB_WORKSPACE}/${{inputs.build-folder-name}}/${{inputs.dfly-executable}}"
-
-
-        timeout 20m pytest -m "${{inputs.filter}}" --durations=10 --color=yes --json-report \
-          --json-report-file=rep1_report.json dragonfly/replication_test.py --log-cli-level=INFO \
-          --df alsologtostderr $1 $2 || code=$?
-
-        # timeout returns 124 if we exceeded the timeout duration
-        if [[ $code -eq 124 ]]; then
-          echo "TIMEDOUT=1">> "$GITHUB_OUTPUT"
-          exit 1
-        fi
-
-        # when a test fails in pytest it returns 1 but there are other return codes as well so we just check if the code is non zero
-        if [[ $code -ne 0 ]]; then
-          exit 1
-        fi
-
     - name: Print last log on timeout
       if: failure()
       shell: bash
@@ -106,14 +80,7 @@ runs:
         }
         cd ${GITHUB_WORKSPACE}/tests
         failed_tests=""
-        # The order in of if is important, and expected to be the oposite order of the pytest run.
-        # As github runner will not run the next step if the pytest failed, we start from the last
-        # report file and if exist we get the failed test from the pytest run, if there are any.
-        if [ -f rep2_report.json ]; then
-          failed_tests=$(get_failed_tests rep2_report.json)
-        elif [ -f rep1_report.json ]; then
-          failed_tests=$(get_failed_tests rep1_report.json)
-        elif [ -f report.json ]; then
+        if [ -f report.json ]; then
           failed_tests=$(get_failed_tests report.json)
         fi
 

.github/workflows/regression-tests.yml

@@ -41,7 +41,6 @@ jobs:
         with:
           dfly-executable: dragonfly
           gspace-secret: ${{ secrets.GSPACES_BOT_DF_BUILD }}
-          run-only-on-ubuntu-latest: false
           build-folder-name: build
           # This expression serves as a ternary operator, i.e. if the condition holds it returns
           # 'not NON_EXISTING_MARK' otherwise not opt_only.

.github/workflows/release.yml

@@ -131,7 +131,6 @@ jobs:
       with:
         dfly-executable: dragonfly-x86_64
         gspace-secret: ${{ secrets.GSPACES_BOT_DF_BUILD }}
-        run-only-on-ubuntu-latest: true
         build-folder-name: ${{ env.RELEASE_DIR }}
     - name: Save artifacts
       run: |

src/core/bloom.cc

@@ -8,6 +8,7 @@
 #include <absl/numeric/bits.h>
 #include <xxhash.h>
 
+#include <algorithm>
 #include <cmath>
 
 #include "base/logging.h"

src/core/search/ast_expr.cc

@@ -58,12 +58,13 @@ AstTagsNode::AstTagsNode(AstExpr&& l, std::string tag) {
 }
 
 AstKnnNode::AstKnnNode(uint32_t limit, std::string_view field, OwnedFtVector vec,
-                       std::string_view score_alias)
+                       std::string_view score_alias, std::optional<size_t> ef_runtime)
     : filter{nullptr},
       limit{limit},
       field{field.substr(1)},
       vec{std::move(vec)},
-      score_alias{score_alias} {
+      score_alias{score_alias},
+      ef_runtime{ef_runtime} {
 }
 
 AstKnnNode::AstKnnNode(AstNode&& filter, AstKnnNode&& self) {
@@ -72,3 +73,9 @@ AstKnnNode::AstKnnNode(AstNode&& filter, AstKnnNode&& self) {
 }
 
 }  // namespace dfly::search
+
+namespace std {
+ostream& operator<<(ostream& os, optional<size_t> o) {
+  return os;
+}
+}  // namespace std

src/core/search/ast_expr.h

@@ -74,7 +74,8 @@ struct AstTagsNode {
 struct AstKnnNode {
   AstKnnNode() = default;
   AstKnnNode(uint32_t limit, std::string_view field, OwnedFtVector vec,
-             std::string_view score_alias);
+             std::string_view score_alias, std::optional<size_t> ef_runtime);
+
   AstKnnNode(AstNode&& sub, AstKnnNode&& self);
 
   friend std::ostream& operator<<(std::ostream& stream, const AstKnnNode& matrix) {
@@ -86,6 +87,7 @@ struct AstKnnNode {
   std::string field;
   OwnedFtVector vec;
   std::string score_alias;
+  std::optional<float> ef_runtime;
 };
 
 struct AstSortNode {
@@ -108,11 +110,18 @@ struct AstNode : public NodeVariants {
   const NodeVariants& Variant() const& {
     return *this;
   }
+
+  // Aggregations: KNN, SORTBY. They reorder result sets and optionally reduce them.
+  bool IsAggregation() const {
+    return std::holds_alternative<AstKnnNode>(Variant());
+  }
 };
 
 using AstExpr = AstNode;
 
 }  // namespace search
 }  // namespace dfly
 
-namespace std {}  // namespace std
+namespace std {
+ostream& operator<<(ostream& os, optional<size_t> o);
+}

src/core/search/indices.cc

@@ -195,6 +195,9 @@ const float* FlatVectorIndex::Get(DocId doc) const {
 }
 
 struct HnswlibAdapter {
+  // Default setting of hnswlib/hnswalg
+  constexpr static size_t kDefaultEfRuntime = 10;
+
   HnswlibAdapter(const SchemaField::VectorParams& params)
       : space_{MakeSpace(params.dim, params.sim)}, world_{GetSpacePtr(),
                                                           params.capacity,
@@ -214,11 +217,13 @@ struct HnswlibAdapter {
     world_.markDelete(id);
   }
 
-  vector<pair<float, DocId>> Knn(float* target, size_t k) {
+  vector<pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef) {
+    world_.setEf(ef.value_or(kDefaultEfRuntime));
     return QueueToVec(world_.searchKnn(target, k));
   }
 
-  vector<pair<float, DocId>> Knn(float* target, size_t k, const vector<DocId>& allowed) {
+  vector<pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
+                                 const vector<DocId>& allowed) {
     struct BinsearchFilter : hnswlib::BaseFilterFunctor {
       virtual bool operator()(hnswlib::labeltype id) {
         return binary_search(allowed->begin(), allowed->end(), id);
@@ -229,6 +234,7 @@ struct HnswlibAdapter {
       const vector<DocId>* allowed;
     };
 
+    world_.setEf(ef.value_or(kDefaultEfRuntime));
     BinsearchFilter filter{&allowed};
     return QueueToVec(world_.searchKnn(target, k, &filter));
   }
@@ -276,12 +282,14 @@ void HnswVectorIndex::Add(DocId id, DocumentAccessor* doc, string_view field) {
     adapter_->Add(ptr.get(), id);
 }
 
-std::vector<std::pair<float, DocId>> HnswVectorIndex::Knn(float* target, size_t k) const {
-  return adapter_->Knn(target, k);
+std::vector<std::pair<float, DocId>> HnswVectorIndex::Knn(float* target, size_t k,
+                                                          std::optional<size_t> ef) const {
+  return adapter_->Knn(target, k, ef);
 }
 std::vector<std::pair<float, DocId>> HnswVectorIndex::Knn(float* target, size_t k,
+                                                          std::optional<size_t> ef,
                                                           const std::vector<DocId>& allowed) const {
-  return adapter_->Knn(target, k, allowed);
+  return adapter_->Knn(target, k, ef, allowed);
 }
 
 void HnswVectorIndex::Remove(DocId id, DocumentAccessor* doc, string_view field) {

src/core/search/indices.h

@@ -136,8 +136,8 @@ struct HnswVectorIndex : public BaseVectorIndex {
   void Add(DocId id, DocumentAccessor* doc, std::string_view field) override;
   void Remove(DocId id, DocumentAccessor* doc, std::string_view field) override;
 
-  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k) const;
-  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k,
+  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef) const;
+  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
                                            const std::vector<DocId>& allowed) const;
 
  private:

src/core/search/lexer.lex

@@ -64,6 +64,7 @@ term_char [_]|\w
 "|"            return Parser::make_OR_OP (loc());
 "KNN"          return Parser::make_KNN (loc());
 "AS"           return Parser::make_AS (loc());
+"EF_RUNTIME"   return Parser::make_EF_RUNTIME (loc());
 
 [0-9]{1,9}                     return make_UINT32(matched_view(), loc());
 [+-]?(([0-9]*[.])?[0-9]+|inf)  return make_DOUBLE(matched_view(), loc());

src/core/search/parser.y

@@ -34,6 +34,7 @@
 #define yylex driver->scanner()->Lex
 
 using namespace std;
+
 }
 
 %parse-param { QueryDriver *driver  }
@@ -46,18 +47,19 @@ using namespace std;
 %define api.token.prefix {TOK_}
 
 %token
-  LPAREN   "("
-  RPAREN   ")"
-  STAR     "*"
-  ARROW    "=>"
-  COLON    ":"
-  LBRACKET "["
-  RBRACKET "]"
-  LCURLBR  "{"
-  RCURLBR  "}"
-  OR_OP    "|"
-  KNN      "KNN"
-  AS       "AS"
+  LPAREN      "("
+  RPAREN      ")"
+  STAR        "*"
+  ARROW       "=>"
+  COLON       ":"
+  LBRACKET    "["
+  RBRACKET    "]"
+  LCURLBR     "{"
+  RCURLBR     "}"
+  OR_OP       "|"
+  KNN         "KNN"
+  AS          "AS"
+  EF_RUNTIME  "EF_RUNTIME"
 ;
 
 %token AND_OP
@@ -81,6 +83,7 @@ using namespace std;
 
 %nterm <AstKnnNode> knn_query
 %nterm <std::string> opt_knn_alias
+%nterm <std::optional<size_t>> opt_ef_runtime
 
 %printer { yyo << $$; } <*>;
 
@@ -93,13 +96,17 @@ final_query:
       { driver->Set(AstKnnNode(std::move($1), std::move($3))); }
 
 knn_query:
-  LBRACKET KNN UINT32 FIELD TERM opt_knn_alias RBRACKET
-    { $$ = AstKnnNode($3, $4, BytesToFtVector($5), $6); }
+  LBRACKET KNN UINT32 FIELD TERM opt_knn_alias opt_ef_runtime RBRACKET
+    { $$ = AstKnnNode($3, $4, BytesToFtVector($5), $6, $7); }
 
 opt_knn_alias:
   AS TERM { $$ = std::move($2); }
   | { $$ = std::string{}; }
 
+opt_ef_runtime:
+  /* empty */ { $$ = std::nullopt; }
+  | EF_RUNTIME UINT32 { $$ = $2; }
+
 filter:
   search_expr               { $$ = std::move($1); }
   | STAR                    { $$ = AstStarNode(); }
@@ -174,5 +181,5 @@ tag_list:
 void
 dfly::search::Parser::error(const location_type& l, const string& m)
 {
-  cerr << l << ": " << m << '\n';
+  driver->Error(l, m);
 }

src/core/search/query_driver.cc

@@ -18,6 +18,10 @@ void QueryDriver::ResetScanner() {
   scanner_->SetParams(params_);
 }
 
+void QueryDriver::Error(const Parser::location_type& loc, std::string_view msg) {
+  LOG(ERROR) << "Parse error " << loc << ": " << msg;
+}
+
 }  // namespace search
 
 }  // namespace dfly

src/core/search/query_driver.h

@@ -52,6 +52,8 @@ class QueryDriver {
     return scanner_.get();
   }
 
+  void Error(const Parser::location_type& loc, std::string_view msg);
+
  public:
   Parser::location_type location;
 

src/core/search/search.cc

@@ -353,9 +353,10 @@ struct BasicSearch {
 
   void SearchKnnHnsw(HnswVectorIndex* vec_index, const AstKnnNode& knn, IndexResult&& sub_results) {
     if (indices_->GetAllDocs().size() == sub_results.Size())
-      knn_distances_ = vec_index->Knn(knn.vec.first.get(), knn.limit);
+      knn_distances_ = vec_index->Knn(knn.vec.first.get(), knn.limit, knn.ef_runtime);
     else
-      knn_distances_ = vec_index->Knn(knn.vec.first.get(), knn.limit, sub_results.Take());
+      knn_distances_ =
+          vec_index->Knn(knn.vec.first.get(), knn.limit, knn.ef_runtime, sub_results.Take());
   }
 
   // [KNN limit @field vec]: Compute distance from `vec` to all vectors keep closest `limit`
@@ -420,6 +421,7 @@ struct BasicSearch {
         profile_builder_ ? make_optional(profile_builder_->Take()) : nullopt;
 
     size_t total = result.Size();
+
     return SearchResult{total,
                         max(total, preagg_total_),
                         result.Take(limit_),
@@ -428,6 +430,7 @@ struct BasicSearch {
                         std::move(error_)};
   }
 
+ private:
   const FieldIndices* indices_;
   size_t limit_;
 
@@ -622,4 +625,8 @@ void SearchAlgorithm::EnableProfiling() {
   profiling_enabled_ = true;
 }
 
+bool SearchAlgorithm::IsProfilingEnabled() const {
+  return profiling_enabled_;
+}
+
 }  // namespace dfly::search

src/core/search/search.h

@@ -144,6 +144,7 @@ class SearchAlgorithm {
   std::optional<AggregationInfo> HasAggregation() const;
 
   void EnableProfiling();
+  bool IsProfilingEnabled() const;
 
  private:
   bool profiling_enabled_ = false;

src/core/search/search_parser_test.cc

@@ -149,4 +149,24 @@ TEST_F(SearchParserTest, KNN) {
   NEXT_TOK(TOK_LBRACKET);
 }
 
+TEST_F(SearchParserTest, KNNfull) {
+  SetInput("*=>[KNN 1 @vector field_vec AS vec_sort EF_RUNTIME 15]");
+  NEXT_TOK(TOK_STAR);
+  NEXT_TOK(TOK_ARROW);
+  NEXT_TOK(TOK_LBRACKET);
+
+  NEXT_TOK(TOK_KNN);
+  NEXT_EQ(TOK_UINT32, uint32_t, 1);
+  NEXT_TOK(TOK_FIELD);
+  NEXT_TOK(TOK_TERM);
+
+  NEXT_TOK(TOK_AS);
+  NEXT_EQ(TOK_TERM, string, "vec_sort");
+
+  NEXT_TOK(TOK_EF_RUNTIME);
+  NEXT_EQ(TOK_UINT32, uint32_t, 15);
+
+  NEXT_TOK(TOK_RBRACKET);
+}
+
 }  // namespace dfly::search