feat(search_family): Add LOAD_FROM option to the FT.AGGREGATE command

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

Author: BagritsevichStepan

src/server/search/aggregator.h

@@ -27,7 +27,7 @@ using Value = ::dfly::search::SortableValue;
 
 // 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>;
+using DocValues = absl::flat_hash_map<std::string, Value>;
 
 struct AggregationResult {
   // Values to be passed to the next step

src/server/search/doc_index.cc

@@ -492,6 +492,115 @@ vector<SearchDocData> ShardDocIndex::SearchForAggregator(
   return out;
 }
 
+join::Vector<join::OwnedEntry> ShardDocIndex::PreagregateDataForJoin(
+    const OpArgs& op_args, absl::Span<const std::string_view> join_fields,
+    search::SearchAlgorithm* search_algo) const {
+  auto search_results = search_algo->Search(&*indices_);
+
+  // First filter out sortable and non-sortable fields
+  // We will load them in different ways
+  const size_t fields_count = join_fields.size();
+  std::vector<bool> is_sortable_field(fields_count);
+  std::vector<FieldReference> basic_fields;
+  basic_fields.reserve(fields_count);
+  for (size_t i = 0; i < fields_count; ++i) {
+    bool is_sortable = IsSortableField(join_fields[i], base_->schema);
+    is_sortable_field[i] = is_sortable;
+    if (!is_sortable) {
+      basic_fields.emplace_back(join_fields[i]);
+    }
+  }
+
+  join::Vector<join::OwnedEntry> result;
+  result.reserve(search_results.ids.size());
+
+  const ShardId shard_id = op_args.shard->shard_id();
+  for (DocId doc : search_results.ids) {
+    auto entry = LoadEntry(doc, op_args);
+    if (!entry)
+      continue;
+
+    auto& [key, accessor] = *entry;
+
+    SearchDocData loaded_basic_fields = accessor->Serialize(base_->schema, basic_fields);
+
+    bool insert_key = true;
+    join::Vector<join::OwnedJoinableValue> join_fields_values(fields_count);
+    for (size_t i = 0; i < fields_count; ++i) {
+      search::SortableValue value;
+      if (is_sortable_field[i]) {
+        value = indices_->GetSortIndexValue(doc, join_fields[i]);
+      } else {
+        value = loaded_basic_fields[join_fields[i]];
+      }
+
+      auto copy = [&](auto&& v) {
+        using T = std::decay_t<decltype(v)>;
+        if constexpr (!std::is_same_v<T, std::monostate>) {
+          join_fields_values[i] = v;
+        } else {
+          // If the value is nil, we skip this key
+          insert_key = false;
+        }
+      };
+
+      std::visit(std::move(copy), value);
+    }
+
+    if (insert_key) {
+      result.emplace_back(std::piecewise_construct, std::forward_as_tuple(shard_id, doc),
+                          std::forward_as_tuple(std::make_move_iterator(join_fields_values.begin()),
+                                                std::make_move_iterator(join_fields_values.end())));
+    }
+  }
+
+  return result;
+}
+
+ShardDocIndex::FieldsValuesPerDocId ShardDocIndex::LoadKeysData(
+    const OpArgs& op_args, const absl::flat_hash_set<search::DocId>& doc_ids,
+    absl::Span<const std::string_view> fields_to_load) const {
+  const size_t fields_count = fields_to_load.size();
+  std::vector<bool> is_sortable_field(fields_count);
+  std::vector<FieldReference> basic_fields;
+  basic_fields.reserve(fields_count);
+  for (size_t i = 0; i < fields_count; ++i) {
+    bool is_sortable = IsSortableField(fields_to_load[i], base_->schema);
+    is_sortable_field[i] = is_sortable;
+    if (!is_sortable) {
+      basic_fields.emplace_back(fields_to_load[i]);
+    }
+  }
+
+  FieldsValuesPerDocId result;
+  result.reserve(doc_ids.size());
+
+  for (DocId doc : doc_ids) {
+    auto entry = LoadEntry(doc, op_args);
+    if (!entry)
+      continue;
+
+    auto& [key, accessor] = *entry;
+
+    SearchDocData loaded_basic_fields = accessor->Serialize(base_->schema, basic_fields);
+
+    FieldsValues fields_values(fields_count);
+    for (size_t i = 0; i < fields_count; ++i) {
+      if (is_sortable_field[i]) {
+        fields_values[i] = indices_->GetSortIndexValue(doc, fields_to_load[i]);
+      } else {
+        fields_values[i] = loaded_basic_fields[fields_to_load[i]];
+      }
+    }
+
+    result.emplace(std::piecewise_construct, std::forward_as_tuple(doc),
+                   std::forward_as_tuple(std::make_move_iterator(fields_values.begin()),
+                                         std::make_move_iterator(fields_values.end())));
+  }
+
+  return result;
+}
+
 DocIndexInfo ShardDocIndex::GetInfo() const {
   return {*base_, key_index_.Size()};
 }

src/server/search/doc_index.h

@@ -20,6 +20,7 @@
 #include "core/search/synonyms.h"
 #include "server/common.h"
 #include "server/search/aggregator.h"
+#include "server/search/index_join.h"
 #include "server/table.h"
 
 namespace dfly {
@@ -66,6 +67,10 @@ struct FieldReference {
     return (is_json && IsJsonPath(name_)) ? name_ : schema.LookupAlias(name_);
   }
 
+  std::string_view Name() const {
+    return name_;
+  }
+
   std::string_view OutputName() const {
     return alias_.empty() ? name_ : alias_;
   }
@@ -122,9 +127,31 @@ struct SearchParams {
 };
 
 struct AggregateParams {
+  struct JoinParams {
+    // Fist field is the index name, second is the field name.
+    using Field = std::pair<std::string, std::string>;
+
+    struct Condition {
+      Condition(std::string_view field_, std::string_view foreign_index_,
+                std::string_view foreign_field_)
+          : field{field_}, foreign_field{Field{foreign_index_, foreign_field_}} {
+      }
+
+      std::string field;
+      Field foreign_field;
+    };
+
+    std::string index;
+    std::string index_alias;
+    std::vector<Condition> conditions;
+    std::string query = "*";
+  };
+
   std::string_view index, query;
   search::QueryParams params;
 
+  std::vector<JoinParams> joins;
+
   std::optional<std::vector<FieldReference>> load_fields;
   std::vector<aggregate::AggregationStep> steps;
 };
@@ -160,6 +187,9 @@ class ShardDocIndex {
   friend class ShardDocIndices;
   using DocId = search::DocId;
 
+  // Used in FieldsValuesPerDocId to store values for each field per document
+  using FieldsValues = absl::InlinedVector<search::SortableValue, 4>;
+
   // DocKeyIndex manages mapping document keys to ids and vice versa through a simple interface.
   struct DocKeyIndex {
     DocId Add(std::string_view key);
@@ -188,6 +218,16 @@ class ShardDocIndex {
                                                  const AggregateParams& params,
                                                  search::SearchAlgorithm* search_algo) const;
 
+  // Methods needed for join operation
+  join::Vector<join::OwnedEntry> PreagregateDataForJoin(
+      const OpArgs& op_args, absl::Span<const std::string_view> join_fields,
+      search::SearchAlgorithm* search_algo) const;
+
+  using FieldsValuesPerDocId = absl::flat_hash_map<DocId, FieldsValues>;
+  FieldsValuesPerDocId LoadKeysData(const OpArgs& op_args,
+                                    const absl::flat_hash_set<search::DocId>& doc_ids,
+                                    absl::Span<const std::string_view> fields_to_load) const;
+
   // Return whether base index matches
   bool Matches(std::string_view key, unsigned obj_code) const;
 

src/server/search/index_join.h

@@ -40,6 +40,10 @@ using Key = std::pair<ShardId, search::DocId>;
 using Entry = std::pair<Key, Vector<JoinableValue> /*fields values of this key*/>;
 using EntriesPerIndex = absl::Span<const Vector<Entry> /*one index can store several keys*/>;
 
+// TODO: comments
+using OwnedJoinableValue = std::variant<double, std::string>;
+using OwnedEntry = std::pair<Key, Vector<OwnedJoinableValue>>;
+
 // Stores data for single join expression,
 // e.g. index1.field1 = index2.field2:
 // field - "field1", foreign_index - "index2", foreign_field - "field2"