feat(query_engine): add join operator

.commitlintrc.js

@@ -37,7 +37,7 @@ const Configuration = {
     "scope-enum": [
       RuleConfigSeverity.Error,
       "always",
-      [...trackedRootDirectories, ...siloSubdirectories, "silo"],
+      [...trackedRootDirectories, ...siloSubdirectories, "silo", "main", "deps"],
     ],
   },
 };

documentation/query_documentation.md

@@ -285,6 +285,84 @@ default.filter(pango_lineage = 'B.1.1.7').phyloSubtree('usherTree')
 {"subtreeNewick": "((key_83:0.00027051)NODE_0000077:3.291e-05,(...)NODE_0000079:1e-06)NODE_0000076;", "missingNodeCount": 0}
 ```
 
+### `join(left, right, on [, type := ...])`
+
+Combines two pipelines on an equality condition (equi-join). `join` can be called as a standalone function or with piped syntax.
+
+The `on` argument is an equality between a left column and a right column. Multiple key pairs are combined with `&&`; a row pair joins only when all key pairs are equal:
+
+```
+join(
+  default.project({primaryKey, country}),
+  default.map({pk := primaryKey, ctry := country}).project({pk, ctry}),
+  primaryKey = pk
+)
+```
+
+Or equivalently using piped syntax:
+
+```
+default.project({primaryKey, country})
+  .join(default.map({pk := primaryKey, ctry := country}).project({pk, ctry}), primaryKey = pk)
+```
+
+Named arguments are also supported:
+
+```
+join(left := <pipeline1>, right := <pipeline2>, on := primaryKey = pk)
+```
+
+Multiple equalities:
+
+```
+join(<left>, <right>, primaryKey = pk && country = ctry)
+```
+
+**Join type:** the optional `type` named argument selects the join kind (default `inner`):
+
+| `type`      | Kept rows                                     | Output columns |
+| ----------- | --------------------------------------------- | -------------- |
+| `inner`     | matching pairs only (default)                 | left + right   |
+| `left`      | all left rows; right null-filled if unmatched | left + right   |
+| `right`     | all right rows; left null-filled if unmatched | left + right   |
+| `full`      | all rows from both sides, null-filled         | left + right   |
+| `leftSemi`  | left rows that have a match                   | left only      |
+| `rightSemi` | right rows that have a match                  | right only     |
+| `leftAnti`  | left rows without a match                     | left only      |
+| `rightAnti` | right rows without a match                    | right only     |
+
+```
+join(<left>, <right>, primaryKey = pk, type := left)
+```
+
+**Output schema:** the left input's columns followed by the right input's columns (semi/anti joins output only the kept side's columns).
+
+**Restrictions:**
+
+- The two inputs must have disjoint column names. A name occurring on both sides would make column references ambiguous and produce a result with duplicate column names, so such a query is rejected — rename one side (e.g. via `map()`) or drop the duplicates (e.g. via `project()`) before joining.
+
+```
+join(
+  default.project({primaryKey, country}),
+  default.project({primaryKey, country}),
+  primaryKey = primaryKey
+)
+```
+
+  fails because `primaryKey` and `country` are present in both inputs; give the two sides distinct names instead (`primaryKey = pk`).
+
+- `filter()` cannot be applied to the output of a `join()`. A predicate above a join is not pushed into a join input, because doing so is not semantics-preserving in general (for example, pushing into the null-supplying side of an outer join would keep null-extended rows that the predicate should have removed). Apply the filter to one of the join inputs instead:
+
+```
+join(
+  default.filter(country='CH').project({primaryKey, country}),
+  default.map({pk := primaryKey}).project({pk}),
+  primaryKey = pk
+)
+```
+
+**Output:** the joined rows. The order of rows is not guaranteed; use `orderBy(...)` for a deterministic order.
+
 ### `unionAll(left, right)`
 
 Concatenates the output of two pipelines. `unionAll` can be called as a standalone function or with piped syntax:

src/silo/query_engine/operator_visitor.h

@@ -9,6 +9,7 @@
 #include "silo/query_engine/operators/fetch_node.h"
 #include "silo/query_engine/operators/filter_node.h"
 #include "silo/query_engine/operators/insertions_node.h"
+#include "silo/query_engine/operators/join_node.h"
 #include "silo/query_engine/operators/map_node.h"
 #include "silo/query_engine/operators/most_recent_common_ancestor_node.h"
 #include "silo/query_engine/operators/mutations_node.h"
@@ -79,6 +80,8 @@ decltype(auto) visit(QueryNode& node, Func&& func) {
          return std::forward<Func>(func)(static_cast<CountFilterNode&>(node));
       case NodeKind::UNION_ALL:
          return std::forward<Func>(func)(static_cast<UnionAllNode&>(node));
+      case NodeKind::JOIN:
+         return std::forward<Func>(func)(static_cast<JoinNode&>(node));
       case NodeKind::SCHEMA:
          return std::forward<Func>(func)(static_cast<SchemaNode&>(node));
       case NodeKind::BITMAP_AGGREGATION:

src/silo/query_engine/operators/join_node.cpp

@@ -0,0 +1,108 @@
+#include "silo/query_engine/operators/join_node.h"
+
+#include <string_view>
+
+#include <arrow/acero/exec_plan.h>
+#include <arrow/acero/options.h>
+#include <arrow/type.h>
+#include <nlohmann/json.hpp>
+
+#include "silo/common/panic.h"
+
+namespace silo::query_engine::operators {
+
+using arrow::acero::JoinType;
+
+JoinNode::JoinNode(
+   QueryNodePtr left,
+   QueryNodePtr right,
+   std::vector<schema::ColumnIdentifier> left_keys,
+   std::vector<schema::ColumnIdentifier> right_keys,
+   JoinType join_type
+)
+    : left(std::move(left)),
+      right(std::move(right)),
+      left_keys(std::move(left_keys)),
+      right_keys(std::move(right_keys)),
+      join_type(join_type) {
+   // Equi-join key vectors are paired positionally (left_keys[i] == right_keys[i]), so they
+   // must have the same length
+   SILO_ASSERT_EQ(this->left_keys.size(), this->right_keys.size());
+}
+
+std::string_view joinTypeToString(JoinType join_type) {
+   switch (join_type) {
+      case JoinType::INNER:
+         return "inner";
+      case JoinType::LEFT_OUTER:
+         return "left";
+      case JoinType::RIGHT_OUTER:
+         return "right";
+      case JoinType::FULL_OUTER:
+         return "full";
+      case JoinType::LEFT_SEMI:
+         return "leftSemi";
+      case JoinType::RIGHT_SEMI:
+         return "rightSemi";
+      case JoinType::LEFT_ANTI:
+         return "leftAnti";
+      case JoinType::RIGHT_ANTI:
+         return "rightAnti";
+   }
+   SILO_UNREACHABLE();
+}
+
+std::vector<schema::ColumnIdentifier> JoinNode::getOutputSchema() const {
+   // Semi/anti joins act as a filter on one input and emit only that side's columns.
+   if (join_type == JoinType::LEFT_SEMI || join_type == JoinType::LEFT_ANTI) {
+      return left->getOutputSchema();
+   }
+   if (join_type == JoinType::RIGHT_SEMI || join_type == JoinType::RIGHT_ANTI) {
+      return right->getOutputSchema();
+   }
+   auto output = left->getOutputSchema();
+   auto right_schema = right->getOutputSchema();
+   output.insert(output.end(), right_schema.begin(), right_schema.end());
+   return output;
+}
+
+arrow::Result<arrow::acero::ExecNode*> JoinNode::addToExecPlan(
+   arrow::acero::ExecPlan& plan,
+   const std::map<schema::TableName, std::shared_ptr<storage::Table>>& tables,
+   const config::QueryOptions& query_options
+) const {
+   ARROW_ASSIGN_OR_RAISE(auto* left_node, left->addToExecPlan(plan, tables, query_options));
+   ARROW_ASSIGN_OR_RAISE(auto* right_node, right->addToExecPlan(plan, tables, query_options));
+
+   std::vector<arrow::FieldRef> left_key_refs;
+   left_key_refs.reserve(left_keys.size());
+   for (const auto& key : left_keys) {
+      left_key_refs.emplace_back(key.name);
+   }
+   std::vector<arrow::FieldRef> right_key_refs;
+   right_key_refs.reserve(right_keys.size());
+   for (const auto& key : right_keys) {
+      right_key_refs.emplace_back(key.name);
+   }
+
+   const arrow::acero::HashJoinNodeOptions options{
+      join_type, std::move(left_key_refs), std::move(right_key_refs)
+   };
+   return arrow::acero::MakeExecNode("hashjoin", &plan, {left_node, right_node}, options);
+}
+
+nlohmann::json JoinNode::toJson() const {
+   nlohmann::json keys = nlohmann::json::array();
+   for (size_t i = 0; i < left_keys.size(); ++i) {
+      keys.push_back({{"left", left_keys[i].name}, {"right", right_keys[i].name}});
+   }
+   return {
+      {"type", nodeKindToString(kind())},
+      {"joinType", joinTypeToString(join_type)},
+      {"on", std::move(keys)},
+      {"left", left->toJson()},
+      {"right", right->toJson()},
+   };
+}
+
+}  // namespace silo::query_engine::operators

src/silo/query_engine/operators/join_node.h

@@ -0,0 +1,58 @@
+#pragma once
+
+#include <map>
+#include <memory>
+#include <vector>
+
+#include <arrow/acero/options.h>
+#include <arrow/result.h>
+#include <nlohmann/json_fwd.hpp>
+
+#include "silo/query_engine/operators/query_node.h"
+#include "silo/schema/database_schema.h"
+#include "silo/storage/table.h"
+
+namespace silo::query_engine::operators {
+
+/// Joins two child pipelines on an equality condition (equi-join), backed by
+/// Arrow's hash-join. `left_keys[i]` is compared for equality against
+/// `right_keys[i]`; all key pairs must match for a row pair to join.
+///
+/// The output schema is the left child's columns followed by the right child's
+/// columns, except for semi/anti joins which output only the columns of the
+/// side they keep. Overlapping column names between the two sides are passed
+/// through as-is (mirroring SQL `SELECT *`); rename or project them downstream
+/// if distinct names are required.
+class JoinNode final : public QueryNode {
+  public:
+   QueryNodePtr left;
+   QueryNodePtr right;
+   std::vector<schema::ColumnIdentifier> left_keys;
+   std::vector<schema::ColumnIdentifier> right_keys;
+   arrow::acero::JoinType join_type;
+
+   JoinNode(
+      QueryNodePtr left,
+      QueryNodePtr right,
+      std::vector<schema::ColumnIdentifier> left_keys,
+      std::vector<schema::ColumnIdentifier> right_keys,
+      arrow::acero::JoinType join_type
+   );
+
+   [[nodiscard]] std::vector<schema::ColumnIdentifier> getOutputSchema() const override;
+
+   [[nodiscard]] arrow::Result<arrow::acero::ExecNode*> addToExecPlan(
+      arrow::acero::ExecPlan& plan,
+      const std::map<schema::TableName, std::shared_ptr<storage::Table>>& tables,
+      const config::QueryOptions& query_options
+   ) const override;
+
+   [[nodiscard]] NodeKind kind() const override { return NodeKind::JOIN; }
+
+   [[nodiscard]] nlohmann::json toJson() const override;
+};
+
+/// Human-readable name for a join type, used in query-plan JSON and diagnostics.
+[[nodiscard]] std::string_view joinTypeToString(arrow::acero::JoinType join_type);
+
+}  // namespace silo::query_engine::operators