feat(search): Defragment search indices (#6144)

* core: Add quota check to page usage
* core/search: Add utility functions to search/base
* core/search: Add defragment support to blocklist,sortedvector
* core/search: Introduce a defragment utility
* search/core: Add defragmentation support for TagIndex
* core/search: Add defragmentation support in FieldIndices
* server/search: Add defragmentation support for shard doc indices
* tests: Unit tests for defragmentation of indices and quota checks

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/core/page_usage/page_usage_stats.cc

@@ -12,6 +12,8 @@
 
 #include <string>
 
+#include "base/cycle_clock.h"
+
 extern "C" {
 #include <unistd.h>
 
@@ -156,8 +158,11 @@ CollectedPageStats PageUsage::UniquePages::CollectedStats() const {
       .shard_wide_summary = {}};
 }
 
-PageUsage::PageUsage(CollectPageStats collect_stats, float threshold)
-    : collect_stats_{collect_stats}, threshold_{threshold} {
+PageUsage::PageUsage(CollectPageStats collect_stats, float threshold, uint64_t quota_usec)
+    : collect_stats_{collect_stats},
+      threshold_{threshold},
+      quota_cycles_{quota_usec == kMaxQuota ? kMaxQuota : base::CycleClock::FromUsec(quota_usec)},
+      start_cycles_(base::CycleClock::Now()) {
 }
 
 bool PageUsage::IsPageForObjectUnderUtilized(void* object) {
@@ -179,4 +184,12 @@ bool PageUsage::ConsumePageStats(mi_page_usage_stats_t stat) {
   return force_reallocate_ || should_reallocate;
 }
 
+bool PageUsage::QuotaDepleted() const {
+  if (quota_cycles_ == kMaxQuota) {
+    return false;
+  }
+
+  return base::CycleClock::Now() - start_cycles_ >= quota_cycles_;
+}
+
 }  // namespace dfly

src/core/page_usage/page_usage_stats.h

@@ -42,7 +42,9 @@ struct CollectedPageStats {
 
 class PageUsage {
  public:
-  PageUsage(CollectPageStats collect_stats, float threshold);
+  static constexpr uint64_t kMaxQuota = std::numeric_limits<uint64_t>::max();
+
+  PageUsage(CollectPageStats collect_stats, float threshold, uint64_t quota_usec = kMaxQuota);
 
   bool IsPageForObjectUnderUtilized(void* object);
 
@@ -66,6 +68,8 @@ class PageUsage {
     force_reallocate_ = force_reallocate;
   }
 
+  bool QuotaDepleted() const;
+
  private:
   CollectPageStats collect_stats_{CollectPageStats::NO};
   float threshold_;
@@ -93,6 +97,9 @@ class PageUsage {
 
   // For use in testing, forces reallocate check to always return true
   bool force_reallocate_{false};
+
+  uint64_t quota_cycles_;
+  uint64_t start_cycles_;
 };
 
 }  // namespace dfly

src/core/page_usage_stats_test.cc

@@ -12,11 +12,14 @@
 #include "core/compact_object.h"
 #include "core/qlist.h"
 #include "core/score_map.h"
+#include "core/search/block_list.h"
+#include "core/search/search.h"
 #include "core/small_string.h"
 #include "core/sorted_map.h"
 #include "core/string_map.h"
 #include "core/string_set.h"
 #include "redis/redis_aux.h"
+#include "util/fibers/fibers.h"
 
 extern "C" {
 #include "redis/zmalloc.h"
@@ -25,6 +28,7 @@ extern "C" {
 ABSL_DECLARE_FLAG(bool, experimental_flat_json);
 
 using namespace dfly;
+using namespace std::chrono_literals;
 
 class PageUsageStatsTest : public ::testing::Test {
  protected:
@@ -213,3 +217,197 @@ TEST_F(PageUsageStatsTest, JSONCons) {
   EXPECT_EQ(json_obj->at("count").as_integer<uint8_t>(), 1);
   EXPECT_EQ(json_obj->at("checked").as_bool(), false);
 }
+
+TEST_F(PageUsageStatsTest, QuotaChecks) {
+  {
+    PageUsage p{CollectPageStats::NO, 0};
+    EXPECT_FALSE(p.QuotaDepleted());
+  }
+  {
+    PageUsage p{CollectPageStats::NO, 0, 4};
+    util::ThisFiber::SleepFor(5us);
+    EXPECT_TRUE(p.QuotaDepleted());
+  }
+}
+
+TEST_F(PageUsageStatsTest, BlockList) {
+  search::BlockList<search::SortedVector<search::DocId>> bl{&m_, 20};
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  // empty list
+  auto result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+
+  // single item will move twice, once for the blocklist and once for the sorted vector
+  bl.Insert(1);
+  result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 2);
+
+  // quota depleted without defragmentation
+  PageUsage p_zero{CollectPageStats::NO, 0.1, 0};
+  p_zero.SetForceReallocate(true);
+  result = bl.Defragment(&p_zero);
+  EXPECT_TRUE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+}
+
+TEST_F(PageUsageStatsTest, BlockListDefragmentResumes) {
+  search::BlockList<search::SortedVector<search::DocId>> bl{&m_, 20};
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  for (size_t i = 0; i < 1000; ++i) {
+    bl.Insert(i);
+  }
+
+  PageUsage p_small_quota{CollectPageStats::NO, 0.1, 10};
+  p_small_quota.SetForceReallocate(true);
+  util::ThisFiber::SleepFor(10us);
+  auto result = bl.Defragment(&p_small_quota);
+  EXPECT_TRUE(result.quota_depleted);
+  EXPECT_GE(result.objects_moved, 0);
+
+  result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_GT(result.objects_moved, 0);
+}
+
+TEST_F(PageUsageStatsTest, BlockListWithPairs) {
+  search::BlockList<search::SortedVector<std::pair<search::DocId, double>>> bl{&m_, 20};
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  for (size_t i = 0; i < 100; ++i) {
+    bl.Insert({i, i * 1.1});
+  }
+
+  const auto result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_GT(result.objects_moved, 0);
+}
+
+TEST_F(PageUsageStatsTest, BlockListWithNonDefragmentableContainer) {
+  search::BlockList<search::CompressedSortedSet> bl{&m_, 20};
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  // empty list
+  auto result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+
+  // will reallocate once for the blocklist, the inner sorted set will be skipped
+  bl.Insert(1);
+  result = bl.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 1);
+}
+
+class MockDocument final : public search::DocumentAccessor {
+ public:
+  MockDocument() {
+    words.reserve(1000);
+    for (size_t i = 0; i < 1000; ++i) {
+      words.push_back(absl::StrFormat("word-%d", i));
+    }
+  }
+
+  std::optional<StringList> GetStrings(std::string_view active_field) const override {
+    return {{words[absl::GetCurrentTimeNanos() % words.size()]}};
+  }
+  std::optional<VectorInfo> GetVector(std::string_view active_field) const override {
+    return std::nullopt;
+  }
+  std::optional<NumsList> GetNumbers(std::string_view active_field) const override {
+    return {{1, 2, 3, 4}};
+  }
+  std::optional<StringList> GetTags(std::string_view active_field) const override {
+    return {{words[absl::GetCurrentTimeNanos() % words.size()]}};
+  }
+
+  std::vector<std::string> words;
+};
+
+TEST_F(PageUsageStatsTest, DefragmentTagIndex) {
+  search::Schema schema;
+  schema.fields["field_name"] =
+      search::SchemaField{search::SchemaField::TAG, 0, "fn", search::SchemaField::TagParams{}};
+  search::FieldIndices index{schema, {}, &m_, nullptr};
+
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  // Empty index
+  search::DefragmentResult result = index.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+
+  const MockDocument md;
+  index.Add(1, md);
+
+  result = index.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  // single doc with single term returned by `GetTags` should result in two reallocations.
+  EXPECT_EQ(result.objects_moved, 2);
+
+  PageUsage p_zero{CollectPageStats::NO, 0.1, 0};
+  p_zero.SetForceReallocate(true);
+  result = index.Defragment(&p_zero);
+  EXPECT_TRUE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+}
+
+TEST_F(PageUsageStatsTest, TagIndexDefragResumeWithChanges) {
+  search::Schema schema;
+  schema.fields["field_name"] =
+      search::SchemaField{search::SchemaField::TAG, 0, "fn", search::SchemaField::TagParams{}};
+  search::FieldIndices index{schema, {}, &m_, nullptr};
+
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  const MockDocument md;
+  for (size_t i = 0; i < 100; ++i) {
+    index.Add(i, md);
+  }
+
+  PageUsage p_small_quota{CollectPageStats::NO, 0.1, 10};
+  p_small_quota.SetForceReallocate(true);
+  search::DefragmentResult result = index.Defragment(&p_small_quota);
+  EXPECT_TRUE(result.quota_depleted);
+  EXPECT_GE(result.objects_moved, 0);
+
+  index.Remove(99, md);
+
+  for (size_t i = 200; i < 300; ++i) {
+    index.Add(i, md);
+  }
+
+  result = index.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_GT(result.objects_moved, 0);
+}
+
+TEST_F(PageUsageStatsTest, DefragmentIndexWithNonDefragmentableFields) {
+  search::Schema schema;
+  schema.fields["text"] =
+      search::SchemaField{search::SchemaField::TEXT, 0, "fn", search::SchemaField::TextParams{}};
+  schema.fields["num"] = search::SchemaField{search::SchemaField::NUMERIC, 0, "fn",
+                                             search::SchemaField::NumericParams{}};
+  search::IndicesOptions options{{}};
+  search::FieldIndices index{schema, options, &m_, nullptr};
+
+  PageUsage p{CollectPageStats::NO, 0.1};
+  p.SetForceReallocate(true);
+
+  const MockDocument md;
+  index.Add(1, md);
+
+  // Unsupported index types will skip defragmenting themselves
+  const search::DefragmentResult result = index.Defragment(&p);
+  EXPECT_FALSE(result.quota_depleted);
+  EXPECT_EQ(result.objects_moved, 0);
+}

src/core/search/CMakeLists.txt

@@ -9,7 +9,7 @@ add_library(dfly_search_core ast_expr.cc base.cc hnsw_index.cc query_driver.cc s
             indices.cc sort_indices.cc vector_utils.cc compressed_sorted_set.cc block_list.cc
             range_tree.cc synonyms.cc ${gen_dir}/parser.cc ${gen_dir}/lexer.cc)
 
-target_link_libraries(dfly_search_core base redis_lib absl::strings
+target_link_libraries(dfly_search_core dfly_page_usage base redis_lib absl::strings
   TRDP::reflex TRDP::uni-algo TRDP::hnswlib Boost::headers)
 
 if(WITH_SIMSIMD)

src/core/search/base.cc

@@ -25,4 +25,10 @@ std::optional<double> ParseNumericField(std::string_view value) {
   return std::nullopt;
 }
 
+DefragmentResult& DefragmentResult::Merge(DefragmentResult&& other) {
+  quota_depleted |= other.quota_depleted;
+  objects_moved += other.objects_moved;
+  return *this;
+}
+
 }  // namespace dfly::search

src/core/search/base.h

@@ -14,8 +14,18 @@
 #include <string_view>
 #include <vector>
 
+namespace dfly {
+class PageUsage;
+}
+
 namespace dfly::search {
 
+struct DefragmentResult {
+  bool quota_depleted{false};
+  size_t objects_moved{0};
+  DefragmentResult& Merge(DefragmentResult&& other);
+};
+
 using DocId = uint32_t;
 using GlobalDocId = uint64_t;
 using ShardId = uint16_t;
@@ -101,6 +111,11 @@ struct BaseIndex {
      Some indices may need to finalize internal structures. See RangeTree for example. */
   virtual void FinalizeInitialization() {
   }
+
+  // Defragments the index by moving objects in underutilized pages to the current malloc page.
+  virtual DefragmentResult Defragment(PageUsage* page_usage) {
+    return DefragmentResult{.quota_depleted = false, .objects_moved = 0};
+  }
 };
 
 // Base class for type-specific sorting indices.