refactor(search_family): Add Aggregator class (#4290)

* refactor(search_family): Add Aggregator class

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* fix(aggregator_test): Fix tests failing

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* refactor: address comments

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* refactor: Restore the previous comment

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* refactor: address comments 2

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* refactor: address comments 3

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

* fix(aggregator): Simplify comparator for the case when one of the values is not present

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

---------

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

Author: BagritsevichStepan

src/server/search/aggregator.cc

@@ -10,62 +10,99 @@ namespace dfly::aggregate {
 
 namespace {
 
-struct GroupStep {
-  PipelineResult operator()(PipelineResult result) {
-    // Separate items into groups
-    absl::flat_hash_map<absl::FixedArray<Value>, std::vector<DocValues>> groups;
-    for (auto& value : result.values) {
-      groups[Extract(value)].push_back(std::move(value));
-    }
+using ValuesList = absl::FixedArray<Value>;
 
-    // Restore DocValues and apply reducers
-    std::vector<DocValues> out;
-    while (!groups.empty()) {
-      auto node = groups.extract(groups.begin());
-      DocValues doc = Unpack(std::move(node.key()));
-      for (auto& reducer : reducers_) {
-        doc[reducer.result_field] = reducer.func({reducer.source_field, node.mapped()});
-      }
-      out.push_back(std::move(doc));
-    }
+ValuesList ExtractFieldsValues(const DocValues& dv, absl::Span<const std::string> fields) {
+  ValuesList out(fields.size());
+  for (size_t i = 0; i < fields.size(); i++) {
+    auto it = dv.find(fields[i]);
+    out[i] = (it != dv.end()) ? it->second : Value{};
+  }
+  return out;
+}
 
-    absl::flat_hash_set<std::string> fields_to_print;
-    fields_to_print.reserve(fields_.size() + reducers_.size());
+DocValues PackFields(ValuesList values, absl::Span<const std::string> fields) {
+  DCHECK_EQ(values.size(), fields.size());
+  DocValues out;
+  for (size_t i = 0; i < fields.size(); i++)
+    out[fields[i]] = std::move(values[i]);
+  return out;
+}
 
-    for (auto& field : fields_) {
-      fields_to_print.insert(std::move(field));
-    }
-    for (auto& reducer : reducers_) {
-      fields_to_print.insert(std::move(reducer.result_field));
-    }
+const Value kEmptyValue = Value{};
+
+}  // namespace
 
-    return {std::move(out), std::move(fields_to_print)};
+void Aggregator::DoGroup(absl::Span<const std::string> fields, absl::Span<const Reducer> reducers) {
+  // Separate items into groups
+  absl::flat_hash_map<ValuesList, std::vector<DocValues>> groups;
+  for (auto& value : result.values) {
+    groups[ExtractFieldsValues(value, fields)].push_back(std::move(value));
   }
 
-  absl::FixedArray<Value> Extract(const DocValues& dv) {
-    absl::FixedArray<Value> out(fields_.size());
-    for (size_t i = 0; i < fields_.size(); i++) {
-      auto it = dv.find(fields_[i]);
-      out[i] = (it != dv.end()) ? it->second : Value{};
+  // Restore DocValues and apply reducers
+  auto& values = result.values;
+  values.clear();
+  values.reserve(groups.size());
+  while (!groups.empty()) {
+    auto node = groups.extract(groups.begin());
+    DocValues doc = PackFields(std::move(node.key()), fields);
+    for (auto& reducer : reducers) {
+      doc[reducer.result_field] = reducer.func({reducer.source_field, node.mapped()});
     }
-    return out;
+    values.push_back(std::move(doc));
   }
 
-  DocValues Unpack(absl::FixedArray<Value>&& values) {
-    DCHECK_EQ(values.size(), fields_.size());
-    DocValues out;
-    for (size_t i = 0; i < fields_.size(); i++)
-      out[fields_[i]] = std::move(values[i]);
-    return out;
+  auto& fields_to_print = result.fields_to_print;
+  fields_to_print.clear();
+  fields_to_print.reserve(fields.size() + reducers.size());
+
+  for (auto& field : fields) {
+    fields_to_print.insert(field);
   }
+  for (auto& reducer : reducers) {
+    fields_to_print.insert(reducer.result_field);
+  }
+}
 
-  std::vector<std::string> fields_;
-  std::vector<Reducer> reducers_;
-};
+void Aggregator::DoSort(std::string_view field, bool descending) {
+  /*
+    Comparator for sorting DocValues by field.
+    If some of the fields is not present in the DocValues, comparator returns:
+    1. l_it == l.end() && r_it != r.end()
+      asc -> false
+      desc -> false
+    2. l_it != l.end() && r_it == r.end()
+      asc -> true
+      desc -> true
+    3. l_it == l.end() && r_it == r.end()
+      asc -> false
+      desc -> false
+  */
+  auto comparator = [&](const DocValues& l, const DocValues& r) {
+    auto l_it = l.find(field);
+    auto r_it = r.find(field);
+
+    // If some of the values is not present
+    if (l_it == l.end() || r_it == r.end()) {
+      return l_it != l.end();
+    }
 
-const Value kEmptyValue = Value{};
+    auto& lv = l_it->second;
+    auto& rv = r_it->second;
+    return !descending ? lv < rv : lv > rv;
+  };
 
-}  // namespace
+  std::sort(result.values.begin(), result.values.end(), std::move(comparator));
+
+  result.fields_to_print.insert(field);
+}
+
+void Aggregator::DoLimit(size_t offset, size_t num) {
+  auto& values = result.values;
+  values.erase(values.begin(), values.begin() + std::min(offset, values.size()));
+  values.resize(std::min(num, values.size()));
+}
 
 const Value& ValueIterator::operator*() const {
   auto it = values_.front().find(field_);
@@ -109,48 +146,30 @@ Reducer::Func FindReducerFunc(ReducerFunc name) {
   return nullptr;
 }
 
-PipelineStep MakeGroupStep(absl::Span<const std::string_view> fields,
-                           std::vector<Reducer> reducers) {
-  return GroupStep{std::vector<std::string>(fields.begin(), fields.end()), std::move(reducers)};
+AggregationStep MakeGroupStep(std::vector<std::string> fields, std::vector<Reducer> reducers) {
+  return [fields = std::move(fields), reducers = std::move(reducers)](Aggregator* aggregator) {
+    aggregator->DoGroup(fields, reducers);
+  };
 }
 
-PipelineStep MakeSortStep(std::string_view field, bool descending) {
-  return [field = std::string(field), descending](PipelineResult result) -> PipelineResult {
-    auto& values = result.values;
-
-    std::sort(values.begin(), values.end(), [field](const DocValues& l, const DocValues& r) {
-      auto it1 = l.find(field);
-      auto it2 = r.find(field);
-      return it1 == l.end() || (it2 != r.end() && it1->second < it2->second);
-    });
-
-    if (descending) {
-      std::reverse(values.begin(), values.end());
-    }
-
-    result.fields_to_print.insert(field);
-    return result;
+AggregationStep MakeSortStep(std::string field, bool descending) {
+  return [field = std::move(field), descending](Aggregator* aggregator) {
+    aggregator->DoSort(field, descending);
   };
 }
 
-PipelineStep MakeLimitStep(size_t offset, size_t num) {
-  return [offset, num](PipelineResult result) {
-    auto& values = result.values;
-    values.erase(values.begin(), values.begin() + std::min(offset, values.size()));
-    values.resize(std::min(num, values.size()));
-    return result;
-  };
+AggregationStep MakeLimitStep(size_t offset, size_t num) {
+  return [=](Aggregator* aggregator) { aggregator->DoLimit(offset, num); };
 }
 
-PipelineResult Process(std::vector<DocValues> values,
-                       absl::Span<const std::string_view> fields_to_print,
-                       absl::Span<const PipelineStep> steps) {
-  PipelineResult result{std::move(values), {fields_to_print.begin(), fields_to_print.end()}};
+AggregationResult Process(std::vector<DocValues> values,
+                          absl::Span<const std::string_view> fields_to_print,
+                          absl::Span<const AggregationStep> steps) {
+  Aggregator aggregator{std::move(values), {fields_to_print.begin(), fields_to_print.end()}};
   for (auto& step : steps) {
-    PipelineResult step_result = step(std::move(result));
-    result = std::move(step_result);
+    step(&aggregator);
   }
-  return result;
+  return aggregator.result;
 }
 
 }  // namespace dfly::aggregate

src/server/search/aggregator.h

@@ -17,19 +17,31 @@
 
 namespace dfly::aggregate {
 
+struct Reducer;
+
 using Value = ::dfly::search::SortableValue;
-using DocValues = absl::flat_hash_map<std::string, Value>;  // documents sent through the pipeline
 
-struct PipelineResult {
+// DocValues sent through the pipeline
+// TODO: Replace DocValues with compact linear search map instead of hash map
+using DocValues = absl::flat_hash_map<std::string_view, Value>;
+
+struct AggregationResult {
   // Values to be passed to the next step
-  // TODO: Replace DocValues with compact linear search map instead of hash map
   std::vector<DocValues> values;
 
   // Fields from values to be printed
-  absl::flat_hash_set<std::string> fields_to_print;
+  absl::flat_hash_set<std::string_view> fields_to_print;
+};
+
+struct Aggregator {
+  void DoGroup(absl::Span<const std::string> fields, absl::Span<const Reducer> reducers);
+  void DoSort(std::string_view field, bool descending = false);
+  void DoLimit(size_t offset, size_t num);
+
+  AggregationResult result;
 };
 
-using PipelineStep = std::function<PipelineResult(PipelineResult)>;  // Group, Sort, etc.
+using AggregationStep = std::function<void(Aggregator*)>;  // Group, Sort, etc.
 
 // Iterator over Span<DocValues> that yields doc[field] or monostate if not present.
 // Extra clumsy for STL compatibility!
@@ -79,18 +91,17 @@ enum class ReducerFunc { COUNT, COUNT_DISTINCT, SUM, AVG, MAX, MIN };
 Reducer::Func FindReducerFunc(ReducerFunc name);
 
 // Make `GROUPBY [fields...]`  with REDUCE step
-PipelineStep MakeGroupStep(absl::Span<const std::string_view> fields,
-                           std::vector<Reducer> reducers);
+AggregationStep MakeGroupStep(std::vector<std::string> fields, std::vector<Reducer> reducers);
 
 // Make `SORTBY field [DESC]` step
-PipelineStep MakeSortStep(std::string_view field, bool descending = false);
+AggregationStep MakeSortStep(std::string field, bool descending = false);
 
 // Make `LIMIT offset num` step
-PipelineStep MakeLimitStep(size_t offset, size_t num);
+AggregationStep MakeLimitStep(size_t offset, size_t num);
 
 // Process values with given steps
-PipelineResult Process(std::vector<DocValues> values,
-                       absl::Span<const std::string_view> fields_to_print,
-                       absl::Span<const PipelineStep> steps);
+AggregationResult Process(std::vector<DocValues> values,
+                          absl::Span<const std::string_view> fields_to_print,
+                          absl::Span<const AggregationStep> steps);
 
 }  // namespace dfly::aggregate

src/server/search/aggregator_test.cc

@@ -10,13 +10,15 @@ namespace dfly::aggregate {
 
 using namespace std::string_literals;
 
+using StepsList = std::vector<AggregationStep>;
+
 TEST(AggregatorTest, Sort) {
   std::vector<DocValues> values = {
       DocValues{{"a", 1.0}},
       DocValues{{"a", 0.5}},
       DocValues{{"a", 1.5}},
   };
-  PipelineStep steps[] = {MakeSortStep("a", false)};
+  StepsList steps = {MakeSortStep("a", false)};
 
   auto result = Process(values, {"a"}, steps);
 
@@ -32,7 +34,8 @@ TEST(AggregatorTest, Limit) {
       DocValues{{"i", 3.0}},
       DocValues{{"i", 4.0}},
   };
-  PipelineStep steps[] = {MakeLimitStep(1, 2)};
+
+  StepsList steps = {MakeLimitStep(1, 2)};
 
   auto result = Process(values, {"i"}, steps);
 
@@ -49,8 +52,8 @@ TEST(AggregatorTest, SimpleGroup) {
       DocValues{{"i", 4.0}, {"tag", "even"}},
   };
 
-  std::string_view fields[] = {"tag"};
-  PipelineStep steps[] = {MakeGroupStep(fields, {})};
+  std::vector<std::string> fields = {"tag"};
+  StepsList steps = {MakeGroupStep(std::move(fields), {})};
 
   auto result = Process(values, {"i", "tag"}, steps);
   EXPECT_EQ(result.values.size(), 2);
@@ -72,13 +75,14 @@ TEST(AggregatorTest, GroupWithReduce) {
     });
   }
 
-  std::string_view fields[] = {"tag"};
+  std::vector<std::string> fields = {"tag"};
   std::vector<Reducer> reducers = {
       Reducer{"", "count", FindReducerFunc(ReducerFunc::COUNT)},
       Reducer{"i", "sum-i", FindReducerFunc(ReducerFunc::SUM)},
       Reducer{"half-i", "distinct-hi", FindReducerFunc(ReducerFunc::COUNT_DISTINCT)},
       Reducer{"null-field", "distinct-null", FindReducerFunc(ReducerFunc::COUNT_DISTINCT)}};
-  PipelineStep steps[] = {MakeGroupStep(fields, std::move(reducers))};
+
+  StepsList steps = {MakeGroupStep(std::move(fields), std::move(reducers))};
 
   auto result = Process(values, {"i", "half-i", "tag"}, steps);
   EXPECT_EQ(result.values.size(), 2);

src/server/search/doc_index.h

@@ -168,7 +168,7 @@ struct AggregateParams {
   search::QueryParams params;
 
   std::optional<SearchFieldsList> load_fields;
-  std::vector<aggregate::PipelineStep> steps;
+  std::vector<aggregate::AggregationStep> steps;
 };
 
 // Stores basic info about a document index.