perf: add parallel remote execution, shard-key routing, hash join operator (#26, #27, #28)

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -276,9 +276,167 @@ class DistributedPlanner {
         }
 
         // Case 2: Sharded -- N RemoteScans + UNION ALL
-        const auto& shard_list = shards_.get_shards(table->table_name);
+        // Optimization (#27): if WHERE contains shard_key = <literal> or
+        // shard_key IN (<literals>), route to only the relevant shard(s).
+        const auto& full_shard_list = shards_.get_shards(table->table_name);
+        std::vector<ShardInfo> pruned = prune_shards(table, where_expr, full_shard_list);
         return make_sharded_union(table, where_expr, nullptr, 0, nullptr, 0,
-                                  nullptr, nullptr, 0, -1, false, shard_list);
+                                  nullptr, nullptr, 0, -1, false, pruned);
+    }
+
+    // Shard pruning (#27): analyze WHERE for shard_key = <literal> or
+    // shard_key IN (<literal_list>). Returns a subset of shards when pruning
+    // is possible, otherwise returns the full shard list.
+    std::vector<ShardInfo> prune_shards(const TableInfo* table,
+                                         const sql_parser::AstNode* where_expr,
+                                         const std::vector<ShardInfo>& all_shards) {
+        if (!where_expr || all_shards.empty()) return all_shards;
+
+        sql_parser::StringRef shard_key = shards_.get_shard_key(table->table_name);
+        if (!shard_key.ptr || shard_key.len == 0) return all_shards;
+
+        // Try to extract shard key literal values from WHERE expression
+        std::vector<size_t> target_indices;
+        extract_shard_targets(where_expr, shard_key, table->table_name,
+                              all_shards.size(), target_indices);
+
+        if (target_indices.empty()) return all_shards;
+
+        // Deduplicate and collect matching shards
+        std::vector<bool> included(all_shards.size(), false);
+        for (size_t idx : target_indices) {
+            if (idx < all_shards.size()) included[idx] = true;
+        }
+        std::vector<ShardInfo> result;
+        for (size_t i = 0; i < all_shards.size(); ++i) {
+            if (included[i]) result.push_back(all_shards[i]);
+        }
+        return result.empty() ? all_shards : result;
+    }
+
+    // Walk a WHERE expression looking for shard_key = <literal> or
+    // shard_key IN (<literal>, ...). Populates target_indices with
+    // the shard index for each matched literal.
+    void extract_shard_targets(const sql_parser::AstNode* expr,
+                                sql_parser::StringRef shard_key,
+                                sql_parser::StringRef table_name,
+                                size_t num_shards,
+                                std::vector<size_t>& target_indices) {
+        if (!expr) return;
+
+        // Check for shard_key = <literal>
+        if (expr->type == sql_parser::NodeType::NODE_BINARY_OP) {
+            sql_parser::StringRef op = expr->value();
+            if (op.len == 1 && op.ptr[0] == '=') {
+                const sql_parser::AstNode* left_node = expr->first_child;
+                const sql_parser::AstNode* right_node = left_node ? left_node->next_sibling : nullptr;
+                if (left_node && right_node) {
+                    // Check if one side is the shard key column and the other is a literal
+                    const sql_parser::AstNode* col_node = nullptr;
+                    const sql_parser::AstNode* lit_node = nullptr;
+                    if (is_shard_key_ref(left_node, shard_key) && is_literal(right_node)) {
+                        col_node = left_node; lit_node = right_node;
+                    } else if (is_shard_key_ref(right_node, shard_key) && is_literal(left_node)) {
+                        col_node = right_node; lit_node = left_node;
+                    }
+                    if (col_node && lit_node) {
+                        size_t idx = literal_to_shard_index(lit_node, table_name, num_shards);
+                        target_indices.push_back(idx);
+                        return;
+                    }
+                }
+            }
+            // Recurse into AND branches
+            if (op.len == 3 &&
+                (op.ptr[0] == 'A' || op.ptr[0] == 'a') &&
+                (op.ptr[1] == 'N' || op.ptr[1] == 'n') &&
+                (op.ptr[2] == 'D' || op.ptr[2] == 'd')) {
+                const sql_parser::AstNode* left_node = expr->first_child;
+                const sql_parser::AstNode* right_node = left_node ? left_node->next_sibling : nullptr;
+                // For AND, either branch matching is sufficient (both must be true,
+                // so if one constrains the shard key, we can prune).
+                std::vector<size_t> left_targets, right_targets;
+                extract_shard_targets(left_node, shard_key, table_name, num_shards, left_targets);
+                extract_shard_targets(right_node, shard_key, table_name, num_shards, right_targets);
+                // Use whichever branch found shard targets (prefer the more selective one)
+                if (!left_targets.empty() && !right_targets.empty()) {
+                    // Intersect: both constraints must hold
+                    std::vector<bool> lset(num_shards, false), rset(num_shards, false);
+                    for (auto i : left_targets) if (i < num_shards) lset[i] = true;
+                    for (auto i : right_targets) if (i < num_shards) rset[i] = true;
+                    for (size_t i = 0; i < num_shards; ++i) {
+                        if (lset[i] && rset[i]) target_indices.push_back(i);
+                    }
+                } else if (!left_targets.empty()) {
+                    target_indices.insert(target_indices.end(), left_targets.begin(), left_targets.end());
+                } else if (!right_targets.empty()) {
+                    target_indices.insert(target_indices.end(), right_targets.begin(), right_targets.end());
+                }
+                return;
+            }
+        }
+
+        // Check for shard_key IN (literal_list)
+        if (expr->type == sql_parser::NodeType::NODE_IN_LIST) {
+            const sql_parser::AstNode* col_expr = expr->first_child;
+            if (col_expr && is_shard_key_ref(col_expr, shard_key)) {
+                for (const sql_parser::AstNode* item = col_expr->next_sibling; item; item = item->next_sibling) {
+                    if (is_literal(item)) {
+                        target_indices.push_back(literal_to_shard_index(item, table_name, num_shards));
+                    } else {
+                        // Non-literal in IN list -- can't prune
+                        target_indices.clear();
+                        return;
+                    }
+                }
+            }
+        }
+    }
+
+    bool is_shard_key_ref(const sql_parser::AstNode* node, sql_parser::StringRef shard_key) const {
+        if (!node) return false;
+        if (node->type == sql_parser::NodeType::NODE_COLUMN_REF ||
+            node->type == sql_parser::NodeType::NODE_IDENTIFIER) {
+            return node->value().equals_ci(shard_key.ptr, shard_key.len);
+        }
+        if (node->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
+            // table.column -- check the column part
+            const sql_parser::AstNode* c = node->first_child;
+            if (c && c->next_sibling) {
+                return c->next_sibling->value().equals_ci(shard_key.ptr, shard_key.len);
+            }
+        }
+        return false;
+    }
+
+    static bool is_literal(const sql_parser::AstNode* node) {
+        if (!node) return false;
+        return node->type == sql_parser::NodeType::NODE_LITERAL_INT ||
+               node->type == sql_parser::NodeType::NODE_LITERAL_FLOAT ||
+               node->type == sql_parser::NodeType::NODE_LITERAL_STRING;
+    }
+
+    size_t literal_to_shard_index(const sql_parser::AstNode* lit,
+                                    sql_parser::StringRef table_name,
+                                    size_t num_shards) const {
+        if (!lit || num_shards == 0) return 0;
+        if (lit->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = lit->value();
+            int64_t val = 0;
+            if (sv.ptr && sv.len > 0) val = std::strtoll(sv.ptr, nullptr, 10);
+            return shards_.shard_index_for_int(table_name, val);
+        }
+        if (lit->type == sql_parser::NodeType::NODE_LITERAL_STRING) {
+            sql_parser::StringRef sv = lit->value();
+            return shards_.shard_index_for_string(table_name, sv.ptr, sv.len);
+        }
+        if (lit->type == sql_parser::NodeType::NODE_LITERAL_FLOAT) {
+            sql_parser::StringRef sv = lit->value();
+            double dv = sv.ptr ? std::strtod(sv.ptr, nullptr) : 0.0;
+            int64_t iv = static_cast<int64_t>(dv);
+            return shards_.shard_index_for_int(table_name, iv);
+        }
+        return 0;
     }
 
     // Build N RemoteScans with UNION ALL

include/sql_engine/operators/hash_join_op.h

@@ -0,0 +1,186 @@
+#ifndef SQL_ENGINE_OPERATORS_HASH_JOIN_OP_H
+#define SQL_ENGINE_OPERATORS_HASH_JOIN_OP_H
+
+#include "sql_engine/operator.h"
+#include "sql_engine/expression_eval.h"
+#include "sql_engine/catalog.h"
+#include "sql_engine/plan_node.h"
+#include "sql_parser/arena.h"
+#include <functional>
+#include <vector>
+#include <unordered_map>
+
+namespace sql_engine {
+
+// HashJoinOperator — builds a hash table on the right (build) side keyed by
+// the equi-join column, then probes it with each row from the left (probe)
+// side.  O(n + m) for equi-joins, vs O(n * m) for nested-loop.
+//
+// Supports INNER and LEFT equi-joins.  Falls back to NestedLoopJoinOperator
+// for non-equi or CROSS joins (selection is done in PlanExecutor::build_join).
+template <sql_parser::Dialect D>
+class HashJoinOperator : public Operator {
+public:
+    HashJoinOperator(Operator* left, Operator* right,
+                     uint8_t join_type,
+                     uint16_t left_join_col,   // ordinal of join key in left row
+                     uint16_t right_join_col,  // ordinal of join key in right row
+                     uint16_t left_cols, uint16_t right_cols,
+                     sql_parser::Arena& arena)
+        : left_(left), right_(right), join_type_(join_type),
+          left_join_col_(left_join_col), right_join_col_(right_join_col),
+          left_cols_(left_cols), right_cols_(right_cols),
+          arena_(arena) {}
+
+    void open() override {
+        // Build phase: consume right side into hash table keyed by join column
+        hash_table_.clear();
+        right_->open();
+        Row r{};
+        while (right_->next(r)) {
+            uint64_t h = hash_value(r.get(right_join_col_));
+            hash_table_[h].push_back(copy_row(r));
+        }
+        right_->close();
+
+        // Probe side
+        left_->open();
+        probe_exhausted_ = false;
+        has_probe_row_ = false;
+        match_idx_ = 0;
+        current_matches_ = nullptr;
+        probe_matched_ = false;
+    }
+
+    bool next(Row& out) override {
+        uint16_t total_cols = left_cols_ + right_cols_;
+
+        while (true) {
+            // Try to emit from current matches
+            if (has_probe_row_ && current_matches_) {
+                while (match_idx_ < current_matches_->size()) {
+                    const Row& rr = (*current_matches_)[match_idx_];
+                    match_idx_++;
+                    // Verify actual equality (not just hash match)
+                    if (values_equal(probe_row_.get(left_join_col_), rr.get(right_join_col_))) {
+                        probe_matched_ = true;
+                        out = combine_rows(probe_row_, rr, total_cols);
+                        return true;
+                    }
+                }
+            }
+
+            // Done with current probe row's matches
+            if (has_probe_row_ && join_type_ == JOIN_LEFT && !probe_matched_) {
+                // LEFT JOIN: emit probe row with NULLs for right side
+                out = make_row(arena_, total_cols);
+                for (uint16_t i = 0; i < left_cols_ && i < probe_row_.column_count; ++i)
+                    out.set(i, probe_row_.get(i));
+                for (uint16_t i = 0; i < right_cols_; ++i)
+                    out.set(left_cols_ + i, value_null());
+                has_probe_row_ = false;
+                return true;
+            }
+
+            // Advance to next probe row
+            if (probe_exhausted_) return false;
+            if (!left_->next(probe_row_)) {
+                probe_exhausted_ = true;
+                return false;
+            }
+
+            has_probe_row_ = true;
+            probe_matched_ = false;
+            match_idx_ = 0;
+
+            uint64_t h = hash_value(probe_row_.get(left_join_col_));
+            auto it = hash_table_.find(h);
+            current_matches_ = (it != hash_table_.end()) ? &it->second : nullptr;
+        }
+    }
+
+    void close() override {
+        left_->close();
+        hash_table_.clear();
+    }
+
+private:
+    Operator* left_;
+    Operator* right_;
+    uint8_t join_type_;
+    uint16_t left_join_col_;
+    uint16_t right_join_col_;
+    uint16_t left_cols_;
+    uint16_t right_cols_;
+    sql_parser::Arena& arena_;
+
+    std::unordered_map<uint64_t, std::vector<Row>> hash_table_;
+    Row probe_row_{};
+    bool has_probe_row_ = false;
+    bool probe_exhausted_ = false;
+    bool probe_matched_ = false;
+    size_t match_idx_ = 0;
+    const std::vector<Row>* current_matches_ = nullptr;
+
+    Row copy_row(const Row& src) {
+        Row dst = make_row(arena_, src.column_count);
+        for (uint16_t i = 0; i < src.column_count; ++i)
+            dst.set(i, src.get(i));
+        return dst;
+    }
+
+    Row combine_rows(const Row& left, const Row& right, uint16_t total_cols) {
+        Row out = make_row(arena_, total_cols);
+        for (uint16_t i = 0; i < left_cols_ && i < left.column_count; ++i)
+            out.set(i, left.get(i));
+        for (uint16_t i = 0; i < right_cols_ && i < right.column_count; ++i)
+            out.set(left_cols_ + i, right.get(i));
+        return out;
+    }
+
+    static uint64_t hash_value(const Value& v) {
+        if (v.is_null()) return 0;
+        switch (v.tag) {
+            case Value::TAG_BOOL:   return std::hash<bool>{}(v.bool_val);
+            case Value::TAG_INT64:  return std::hash<int64_t>{}(v.int_val);
+            case Value::TAG_UINT64: return std::hash<uint64_t>{}(v.uint_val);
+            case Value::TAG_DOUBLE: return std::hash<double>{}(v.double_val);
+            case Value::TAG_STRING: {
+                // FNV-1a hash
+                uint64_t h = 14695981039346656037ULL;
+                for (uint32_t i = 0; i < v.str_val.len; ++i) {
+                    h ^= static_cast<uint64_t>(static_cast<unsigned char>(v.str_val.ptr[i]));
+                    h *= 1099511628211ULL;
+                }
+                return h;
+            }
+            default: return 0;
+        }
+    }
+
+    static bool values_equal(const Value& a, const Value& b) {
+        if (a.is_null() || b.is_null()) return false;
+        if (a.tag == b.tag) {
+            switch (a.tag) {
+                case Value::TAG_BOOL:   return a.bool_val == b.bool_val;
+                case Value::TAG_INT64:  return a.int_val == b.int_val;
+                case Value::TAG_UINT64: return a.uint_val == b.uint_val;
+                case Value::TAG_DOUBLE: return a.double_val == b.double_val;
+                case Value::TAG_STRING:
+                    return a.str_val.len == b.str_val.len &&
+                           (a.str_val.len == 0 ||
+                            std::memcmp(a.str_val.ptr, b.str_val.ptr, a.str_val.len) == 0);
+                default: return false;
+            }
+        }
+        // Cross-type numeric comparison
+        if (a.is_numeric() && b.is_numeric()) {
+            return a.to_double() == b.to_double();
+        }
+        return false;
+    }
+};
+
+} // namespace sql_engine
+
+#endif // SQL_ENGINE_OPERATORS_HASH_JOIN_OP_H