feat(search_family): Add basic support for the FT.CONFIG

Stepan Bagritsevich · Stepan Bagritsevich · commit fee09baf4ce5 · 2025-03-16T21:34:32.000+01:00
fixes #4352 Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>
diff --git a/src/server/search/search_family.cc b/src/server/search/search_family.cc
@@ -599,6 +599,111 @@ void SearchReply(const SearchParams& params, std::optional<search::AggregationIn
   }
 }
 
+constexpr std::string_view kSearchLimit = "MAXSEARCHRESULTS"sv;
+constexpr std::string_view kAggregateLimit = "MAXAGGREGATERESULTS"sv;
+
+// Do not forget to update SearchFamily::config_values_ after adding new option
+constexpr SearchFamily::ConfigOptionsMap<std::string_view> kConfigOptionsHelp = {{
+    {kSearchLimit, "Maximum number of results from ft.search command"sv},
+    {kAggregateLimit, "Maximum number of results from ft.aggregate command"sv},
+}};
+
+template <typename V>
+std::optional<V> FindOptionsMapValue(const SearchFamily::ConfigOptionsMap<V>& options,
+                                     std::string_view name) {
+  auto it = std::find_if(options.begin(), options.end(), [name](const auto& opt) {
+    return absl::EqualsIgnoreCase(opt.first, name);
+  });
+  if (it != options.end()) {
+    return it->second;
+  }
+  return std::nullopt;
+}
+
+void FtConfigHelp(CmdArgParser* parser, RedisReplyBuilder* rb, SearchFamily::Config* config) {
+  string_view option = parser->Next();
+
+  auto send_value = [&](string_view option_name) {
+    auto value = FindOptionsMapValue(*config, option_name);
+    DCHECK(value.has_value());
+    rb->SendLong(value.value());
+  };
+
+  if (option == "*"sv) {
+    rb->StartArray(kConfigOptionsHelp.size());
+    for (const auto& option_help : kConfigOptionsHelp) {
+      rb->StartArray(5);
+      rb->SendBulkString(option_help.first);
+      rb->SendBulkString("Description"sv);
+      rb->SendBulkString(option_help.second);
+      rb->SendBulkString("Value"sv);
+      send_value(option_help.first);
+    }
+    return;
+  }
+
+  auto option_description = FindOptionsMapValue(kConfigOptionsHelp, option);
+  if (option_description) {
+    rb->StartArray(1);
+    rb->StartArray(5);
+    rb->SendBulkString(absl::AsciiStrToUpper(option));
+    rb->SendBulkString("Description"sv);
+    rb->SendBulkString(option_description.value());
+    rb->SendBulkString("Value"sv);
+    send_value(option);
+    return;
+  }
+
+  LOG(WARNING) << "Unknown configuration option: " << option;
+  rb->SendEmptyArray();
+}
+
+void FtConfigGet(CmdArgParser* parser, RedisReplyBuilder* rb, SearchFamily::Config* config) {
+  string_view option = parser->Next();
+
+  if (option == "*"sv) {
+    rb->StartArray(config->size());
+    for (const auto& option_help : *config) {
+      rb->StartArray(2);
+      rb->SendBulkString(option_help.first);
+      rb->SendLong(option_help.second);
+    }
+    return;
+  }
+
+  auto option_value = FindOptionsMapValue(*config, option);
+  if (option_value) {
+    rb->StartArray(1);
+    rb->StartArray(2);
+    rb->SendBulkString(absl::AsciiStrToUpper(option));
+    rb->SendLong(option_value.value());
+    return;
+  }
+
+  LOG(WARNING) << "Unknown configuration option: " << option;
+  rb->SendEmptyArray();
+}
+
+void FtConfigSet(CmdArgParser* parser, RedisReplyBuilder* rb, SearchFamily::Config* config) {
+  string_view option = parser->Next();
+  uint64_t value = parser->Next<uint64_t>();
+
+  if (auto err = parser->Error(); err)
+    return rb->SendError(err->MakeReply());
+
+  auto it = std::find_if(config->begin(), config->end(), [option_name = option](const auto& opt) {
+    return absl::EqualsIgnoreCase(opt.first, option_name);
+  });
+
+  if (it != config->end()) {
+    LOG(INFO) << "Setting " << option << " to " << value;
+    it->second = value;
+    rb->SendOk();
+  } else {
+    rb->SendError("Invalid option"sv);
+  }
+}
+
 }  // namespace
 
 void SearchFamily::FtCreate(CmdArgList args, const CommandContext& cmd_cntx) {
@@ -815,6 +920,13 @@ void SearchFamily::FtSearch(CmdArgList args, const CommandContext& cmd_cntx) {
   if (!search_algo.Init(query_str, &params->query_params, sort_opt))
     return builder->SendError("Query syntax error");
 
+  {
+    util::fb2::LockGuard lg{config_mu_};
+    auto search_limit = FindOptionsMapValue(config_, kSearchLimit);
+    DCHECK(search_limit.has_value());
+    params->limit_total = std::min(params->limit_total, search_limit.value());
+  }
+
   // Because our coordinator thread may not have a shard, we can't check ahead if the index exists.
   atomic<bool> index_not_found{false};
   vector<SearchResult> docs(shard_set->size());
@@ -966,6 +1078,31 @@ void SearchFamily::FtProfile(CmdArgList args, const CommandContext& cmd_cntx) {
   }
 }
 
+// Do not forget to update kConfigOptionsHelp after adding new option
+SearchFamily::Config SearchFamily::config_ = {{
+    {kSearchLimit, 10000},
+    {kAggregateLimit, 10000},
+}};
+
+util::fb2::Mutex SearchFamily::config_mu_;
+
+void SearchFamily::FtConfig(CmdArgList args, const CommandContext& cmd_cntx) {
+  CmdArgParser parser{args};
+  auto* rb = static_cast<RedisReplyBuilder*>(cmd_cntx.rb);
+
+  auto func = parser.TryMapNext("GET", &FtConfigGet, "SET", &FtConfigSet, "HELP", &FtConfigHelp);
+
+  if (func) {
+    util::fb2::LockGuard lg{config_mu_};
+    return func.value()(&parser, rb, &config_);
+  } else {
+    return rb->SendError("Unknown subcommand");
+  }
+
+  static_assert(config_.size() == kConfigOptionsHelp.size(),
+                "SearchFamily::config_values_ and kConfigOptionsHelp must have the same size.");
+}
+
 void SearchFamily::FtTagVals(CmdArgList args, const CommandContext& cmd_cntx) {
   string_view index_name = ArgS(args, 0);
   string_view field_name = ArgS(args, 1);
@@ -1052,11 +1189,24 @@ void SearchFamily::FtAggregate(CmdArgList args, const CommandContext& cmd_cntx)
   auto* rb = static_cast<RedisReplyBuilder*>(cmd_cntx.rb);
   auto sortable_value_sender = SortableValueSender(rb);
 
-  const size_t result_size = agg_results.values.size();
+  size_t result_size = agg_results.values.size();
+
+  {
+    util::fb2::LockGuard lg{config_mu_};
+    auto aggregate_limit = FindOptionsMapValue(config_, kAggregateLimit);
+    DCHECK(aggregate_limit.has_value());
+    result_size = std::min(result_size, aggregate_limit.value());
+  }
+
   rb->StartArray(result_size + 1);
   rb->SendLong(result_size);
 
   for (const auto& value : agg_results.values) {
+    if (result_size == 0) {
+      break;
+    }
+    result_size--;
+
     size_t fields_count = 0;
     for (const auto& field : agg_results.fields_to_print) {
       if (value.find(field) != value.end()) {
@@ -1089,18 +1239,19 @@ void SearchFamily::Register(CommandRegistry* registry) {
       CO::NO_KEY_TRANSACTIONAL | CO::NO_KEY_TX_SPAN_ALL | CO::NO_AUTOJOURNAL;
 
   registry->StartFamily();
-  *registry << CI{"FT.CREATE", CO::WRITE | CO::GLOBAL_TRANS, -2, 0, 0, acl::FT_SEARCH}.HFUNC(
-                   FtCreate)
-            << CI{"FT.ALTER", CO::WRITE | CO::GLOBAL_TRANS, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtAlter)
-            << CI{"FT.DROPINDEX", CO::WRITE | CO::GLOBAL_TRANS, -2, 0, 0, acl::FT_SEARCH}.HFUNC(
-                   FtDropIndex)
-            << CI{"FT.INFO", kReadOnlyMask, 2, 0, 0, acl::FT_SEARCH}.HFUNC(FtInfo)
-            // Underscore same as in RediSearch because it's "temporary" (long time already)
-            << CI{"FT._LIST", kReadOnlyMask, 1, 0, 0, acl::FT_SEARCH}.HFUNC(FtList)
-            << CI{"FT.SEARCH", kReadOnlyMask, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtSearch)
-            << CI{"FT.AGGREGATE", kReadOnlyMask, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtAggregate)
-            << CI{"FT.PROFILE", kReadOnlyMask, -4, 0, 0, acl::FT_SEARCH}.HFUNC(FtProfile)
-            << CI{"FT.TAGVALS", kReadOnlyMask, 3, 0, 0, acl::FT_SEARCH}.HFUNC(FtTagVals);
+  *registry
+      << CI{"FT.CREATE", CO::WRITE | CO::GLOBAL_TRANS, -2, 0, 0, acl::FT_SEARCH}.HFUNC(FtCreate)
+      << CI{"FT.ALTER", CO::WRITE | CO::GLOBAL_TRANS, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtAlter)
+      << CI{"FT.DROPINDEX", CO::WRITE | CO::GLOBAL_TRANS, -2, 0, 0, acl::FT_SEARCH}.HFUNC(
+             FtDropIndex)
+      << CI{"FT.CONFIG", CO::WRITE | CO::GLOBAL_TRANS, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtConfig)
+      << CI{"FT.INFO", kReadOnlyMask, 2, 0, 0, acl::FT_SEARCH}.HFUNC(FtInfo)
+      // Underscore same as in RediSearch because it's "temporary" (long time already)
+      << CI{"FT._LIST", kReadOnlyMask, 1, 0, 0, acl::FT_SEARCH}.HFUNC(FtList)
+      << CI{"FT.SEARCH", kReadOnlyMask, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtSearch)
+      << CI{"FT.AGGREGATE", kReadOnlyMask, -3, 0, 0, acl::FT_SEARCH}.HFUNC(FtAggregate)
+      << CI{"FT.PROFILE", kReadOnlyMask, -4, 0, 0, acl::FT_SEARCH}.HFUNC(FtProfile)
+      << CI{"FT.TAGVALS", kReadOnlyMask, 3, 0, 0, acl::FT_SEARCH}.HFUNC(FtTagVals);
 }
 
 }  // namespace dfly
diff --git a/src/server/search/search_family.h b/src/server/search/search_family.h
@@ -18,6 +18,16 @@ class CommandRegistry;
 struct CommandContext;
 
 class SearchFamily {
+ public:
+  static void Register(CommandRegistry* registry);
+
+  static inline constexpr size_t kOptionsCount = 2;
+
+  template <typename V>
+  using ConfigOptionsMap = std::array<std::pair<std::string_view, V>, kOptionsCount>;
+  using Config = ConfigOptionsMap<uint64_t>;
+
+ private:
   using SinkReplyBuilder = facade::SinkReplyBuilder;
 
   static void FtCreate(CmdArgList args, const CommandContext& cmd_cntx);
@@ -27,11 +37,12 @@ class SearchFamily {
   static void FtList(CmdArgList args, const CommandContext& cmd_cntx);
   static void FtSearch(CmdArgList args, const CommandContext& cmd_cntx);
   static void FtProfile(CmdArgList args, const CommandContext& cmd_cntx);
+  static void FtConfig(CmdArgList args, const CommandContext& cmd_cntx);
   static void FtAggregate(CmdArgList args, const CommandContext& cmd_cntx);
   static void FtTagVals(CmdArgList args, const CommandContext& cmd_cntx);
 
- public:
-  static void Register(CommandRegistry* registry);
+  static util::fb2::Mutex config_mu_;
+  static Config config_;  // guarded by config_mu_
 };
 
 }  // namespace dfly
diff --git a/src/server/search/search_family_test.cc b/src/server/search/search_family_test.cc
@@ -2014,4 +2014,76 @@ TEST_F(SearchFamilyTest, InvalidCreateOptions) {
   EXPECT_THAT(resp, ErrArg(kInvalidIntErr));
 }
 
+TEST_F(SearchFamilyTest, ConfigMaxSearchResults) {
+  Run({"HSET", "doc1", "title", "hello world1"});
+  Run({"HSET", "doc2", "title", "hello world2"});
+  Run({"FT.CREATE", "index", "ON", "HASH", "SCHEMA", "title", "TEXT"});
+
+  auto resp = Run({"FT.SEARCH", "index", "*", "NOCONTENT"});
+  EXPECT_THAT(resp, IsUnordArray(IntArg(2), "doc1", "doc2"));
+
+  resp = Run({"FT.CONFIG", "GET", "MAXSEARCHRESULTS"});
+  EXPECT_THAT(resp, IsArray("MAXSEARCHRESULTS", IntArg(10000)));
+
+  resp = Run({"FT.CONFIG", "SET", "MAXSEARCHRESULTS", "1"});
+  EXPECT_EQ(resp, "OK");
+
+  resp = Run({"FT.SEARCH", "index", "*", "NOCONTENT"});
+  EXPECT_THAT(resp, IsUnordArray(IntArg(2), AnyOf("doc1", "doc2")));
+
+  resp = Run({"FT.CONFIG", "GET", "MAXSEARCHRESULTS"});
+  EXPECT_THAT(resp, IsArray("MAXSEARCHRESULTS", IntArg(1)));
+}
+
+TEST_F(SearchFamilyTest, ConfigMaxAggregateResult) {
+  Run({"HSET", "doc1", "title", "hello world1"});
+  Run({"HSET", "doc2", "title", "hello world2"});
+  Run({"FT.CREATE", "index", "ON", "HASH", "SCHEMA", "title", "TEXT"});
+
+  auto resp = Run({"FT.AGGREGATE", "index", "*", "GROUPBY", "1", "@title", "REDUCE", "COUNT", "0",
+                   "AS", "count"});
+  EXPECT_THAT(resp, IsUnordArrayWithSize(IsMap("title", "hello world1", "count", "1"),
+                                         IsMap("title", "hello world2", "count", "1")));
+
+  resp = Run({"FT.CONFIG", "GET", "MAXAGGREGATERESULTS"});
+  EXPECT_THAT(resp, IsArray("MAXAGGREGATERESULTS", IntArg(10000)));
+
+  resp = Run({"FT.CONFIG", "SET", "MAXAGGREGATERESULTS", "1"});
+  EXPECT_EQ(resp, "OK");
+
+  resp = Run({"FT.AGGREGATE", "index", "*", "GROUPBY", "1", "@title", "REDUCE", "COUNT", "0", "AS",
+              "count"});
+  EXPECT_THAT(resp, IsUnordArrayWithSize(AnyOf(IsMap("title", "hello world1", "count", "1"),
+                                               IsMap("title", "hello world2", "count", "1"))));
+
+  resp = Run({"FT.CONFIG", "GET", "MAXAGGREGATERESULTS"});
+  EXPECT_THAT(resp, IsArray("MAXAGGREGATERESULTS", IntArg(1)));
+}
+
+TEST_F(SearchFamilyTest, InvalidConfigOptions) {
+  // Test with an invalid argument
+  auto resp = Run({"FT.CONFIG", "INVALIDARG", "INVLIDARG"});
+  EXPECT_THAT(resp, ErrArg("Unknown subcommand"));
+
+  // Test with an invalid argument
+  resp = Run({"FT.CONFIG", "GET", "INVALIDARG"});
+  EXPECT_THAT(resp, IsArray());
+
+  // Test with an invalid argument
+  resp = Run({"FT.CONFIG", "SET", "INVALIDARG"});
+  EXPECT_THAT(resp, ErrArg(kSyntaxErr));
+
+  // Test with an invalid argument
+  resp = Run({"FT.CONFIG", "SET", "INVALIDARG", "5"});
+  EXPECT_THAT(resp, ErrArg("Invalid option"));
+
+  // Test with an invalid value
+  resp = Run({"FT.CONFIG", "SET", "MAXSEARCHRESULTS", "not_a_number"});
+  EXPECT_THAT(resp, ErrArg(kInvalidIntErr));
+
+  // Test with an invalid argument
+  resp = Run({"FT.CONFIG", "HELP", "INVALIDARG"});
+  EXPECT_THAT(resp, IsArray());
+}
+
 }  // namespace dfly
