feat(search): Add support for multiple prefixes in search indexes (#5985)

* feat(search): Add support for multiple prefixes in search indexes

* fix: failing tests

* fix: review comments

Author: vyavdoshenko

src/server/search/doc_index.cc

@@ -35,11 +35,9 @@ void TraverseAllMatching(const DocIndex& index, const OpArgs& op_args, F&& f) {
   string scratch;
   auto cb = [&](PrimeTable::iterator it) {
     const PrimeValue& pv = it->second;
-    if (pv.ObjType() != index.GetObjCode())
-      return;
-
     string_view key = it->first.GetSlice(&scratch);
-    if (key.rfind(index.prefix, 0) != 0)
+
+    if (!index.Matches(key, pv.ObjType()))
       return;
 
     auto accessor = GetAccessor(op_args.db_cntx, pv);
@@ -149,9 +147,13 @@ string DocIndexInfo::BuildRestoreCommand() const {
   // ON HASH/JSON
   absl::StrAppend(&out, "ON", " ", base_index.type == DocIndex::HASH ? "HASH" : "JSON");
 
-  // optional PREFIX 1 *prefix*
-  if (!base_index.prefix.empty())
-    absl::StrAppend(&out, " PREFIX", " 1 ", base_index.prefix);
+  // optional PREFIX count *prefix1* *prefix2* ...
+  if (!base_index.prefixes.empty()) {
+    absl::StrAppend(&out, " PREFIX", " ", base_index.prefixes.size());
+    for (const auto& prefix : base_index.prefixes) {
+      absl::StrAppend(&out, " ", prefix);
+    }
+  }
 
   // STOPWORDS
   absl::StrAppend(&out, " STOPWORDS ", base_index.options.stopwords.size());
@@ -247,7 +249,18 @@ uint8_t DocIndex::GetObjCode() const {
 }
 
 bool DocIndex::Matches(string_view key, unsigned obj_code) const {
-  return obj_code == GetObjCode() && key.rfind(prefix, 0) == 0;
+  if (obj_code != GetObjCode())
+    return false;
+
+  // Empty prefixes means match all keys
+  if (prefixes.empty())
+    return true;
+
+  for (const auto& prefix : prefixes) {
+    if (key.rfind(prefix, 0) == 0)
+      return true;
+  }
+  return false;
 }
 
 ShardDocIndex::ShardDocIndex(shared_ptr<const DocIndex> index)
@@ -269,7 +282,8 @@ void ShardDocIndex::Rebuild(const OpArgs& op_args, PMR_NS::memory_resource* mr)
 
   indices_->FinalizeInitialization();
 
-  VLOG(1) << "Indexed " << key_index_.Size() << " docs on " << base_->prefix;
+  VLOG(1) << "Indexed " << key_index_.Size()
+          << " docs on prefixes: " << absl::StrJoin(base_->prefixes, ", ");
 }
 
 void ShardDocIndex::RebuildForGroup(const OpArgs& op_args, const std::string_view& group_id,

src/server/search/doc_index.h

@@ -194,7 +194,7 @@ struct DocIndex {
 
   search::Schema schema;
   search::IndicesOptions options{};
-  std::string prefix{};
+  std::vector<std::string> prefixes{};
   DataType type{HASH};
 };
 

src/server/search/search_family.cc

@@ -211,10 +211,11 @@ ParseResult<bool> ParseOnOption(CmdArgParser* parser, DocIndex* index) {
 
 // PREFIX count prefix [prefix ...]
 ParseResult<bool> ParsePrefix(CmdArgParser* parser, DocIndex* index) {
-  if (!parser->Check("1")) {
-    return CreateSyntaxError("Multiple prefixes are not supported"sv);
+  size_t count = parser->Next<size_t>();
+  index->prefixes.reserve(count);
+  for (size_t i = 0; i < count; i++) {
+    index->prefixes.push_back(parser->Next<std::string>());
   }
-  index->prefix = parser->Next<std::string>();
   return true;
 }
 
@@ -1217,8 +1218,11 @@ void SearchFamily::FtInfo(CmdArgList args, const CommandContext& cmd_cntx) {
     rb->StartCollection(3, RedisReplyBuilder::MAP);
     rb->SendSimpleString("key_type");
     rb->SendSimpleString(info.base_index.type == DocIndex::JSON ? "JSON" : "HASH");
-    rb->SendSimpleString("prefix");
-    rb->SendSimpleString(info.base_index.prefix);
+    rb->SendSimpleString("prefixes");
+    rb->StartArray(info.base_index.prefixes.size());
+    for (const auto& prefix : info.base_index.prefixes) {
+      rb->SendBulkString(prefix);
+    }
     rb->SendSimpleString("default_score");
     rb->SendLong(1);
   }

src/server/search/search_family_test.cc

@@ -290,11 +290,12 @@ TEST_F(SearchFamilyTest, InfoIndex) {
   }
 
   auto info = Run({"ft.info", "idx-1"});
-  EXPECT_THAT(info,
-              IsArray(_, _, _, IsArray("key_type", "HASH", "prefix", "doc-", "default_score", 1),
-                      "index_options", RespArray(IsEmpty()), "attributes",
-                      IsArray(IsArray("identifier", "name", "attribute", "name", "type", "TEXT")),
-                      "num_docs", IntArg(15)));
+  EXPECT_THAT(
+      info,
+      IsArray(_, _, _, IsArray("key_type", "HASH", "prefixes", IsArray("doc-"), "default_score", 1),
+              "index_options", RespArray(IsEmpty()), "attributes",
+              IsArray(IsArray("identifier", "name", "attribute", "name", "type", "TEXT")),
+              "num_docs", IntArg(15)));
 }
 
 TEST_F(SearchFamilyTest, Stats) {

tests/dragonfly/search_test.py

@@ -128,15 +128,29 @@ async def test_management(async_client: aioredis.Redis):
     assert sorted(await async_client.execute_command("FT._LIST")) == ["i1", "i2"]
 
     i1info = await i1.info()
-    assert i1info["index_definition"] == ["key_type", "HASH", "prefix", "p1", "default_score", 1]
+    assert i1info["index_definition"] == [
+        "key_type",
+        "HASH",
+        "prefixes",
+        ["p1"],
+        "default_score",
+        1,
+    ]
     assert i1info["num_docs"] == 10
     assert sorted(i1info["attributes"]) == [
         ["identifier", "f1", "attribute", "f1", "type", "TEXT"],
         ["identifier", "f2", "attribute", "f2", "type", "NUMERIC", "SORTABLE", "blocksize", "7000"],
     ]
 
     i2info = await i2.info()
-    assert i2info["index_definition"] == ["key_type", "HASH", "prefix", "p2", "default_score", 1]
+    assert i2info["index_definition"] == [
+        "key_type",
+        "HASH",
+        "prefixes",
+        ["p2"],
+        "default_score",
+        1,
+    ]
     assert i2info["num_docs"] == 15
     assert sorted(i2info["attributes"]) == [
         [