Merge pull request #61 from ProxySQL/feat/shard-parallel-composite-move

feat: parallel scatter, composite keys, LIST prune, plan-cache redistribute

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -17,86 +17,82 @@ class SetOpOperator : public Operator {
 public:
     SetOpOperator(Operator* left, Operator* right, uint8_t op, bool all,
                   bool parallel_open = false, ThreadPool* pool = nullptr)
-        : left_(left), right_(right), op_(op), all_(all),
+        : op_(op), all_(all), parallel_open_(parallel_open), pool_(pool) {
+        children_.push_back(left);
+        children_.push_back(right);
+    }
+
+    explicit SetOpOperator(std::vector<Operator*> children,
+                           bool parallel_open = false, ThreadPool* pool = nullptr)
+        : children_(std::move(children)), op_(SET_OP_UNION), all_(true),
           parallel_open_(parallel_open), pool_(pool) {}
 
     void open() override {
-        if (parallel_open_ && pool_) {
-            // Thread-pool parallel open: ~1-2us dispatch vs ~200us for std::async
-            auto fl = pool_->submit([this]{ left_->open(); });
-            auto fr = pool_->submit([this]{ right_->open(); });
-            fl.get();
-            fr.get();
-        } else if (parallel_open_) {
-            // Fallback: std::async when no pool available
-            auto fl = std::async(std::launch::async, [this]{ left_->open(); });
-            auto fr = std::async(std::launch::async, [this]{ right_->open(); });
-            fl.get();
-            fr.get();
+        if (children_.empty()) return;
+        if (parallel_open_ && children_.size() > 1) {
+            std::vector<std::future<void>> futures;
+            futures.reserve(children_.size());
+            for (size_t i = 0; i < children_.size(); ++i) {
+                auto launcher = [this, i]{ children_[i]->open(); };
+                if (pool_) {
+                    futures.push_back(pool_->submit(std::move(launcher)));
+                } else {
+                    futures.push_back(std::async(std::launch::async, std::move(launcher)));
+                }
+            }
+            for (auto& f : futures) f.get();
         } else {
-            left_->open();
-            right_->open();
+            for (auto* c : children_) c->open();
         }
-        reading_left_ = true;
+        child_idx_ = 0;
         seen_.clear();
         expected_col_count_ = -1;
 
-        if (op_ == SET_OP_INTERSECT || op_ == SET_OP_EXCEPT) {
-            // Materialize right side into a set
+        if ((op_ == SET_OP_INTERSECT || op_ == SET_OP_EXCEPT) && children_.size() >= 2) {
             right_set_.clear();
             Row r{};
-            while (right_->next(r)) {
+            while (children_[1]->next(r)) {
                 check_col_count(r);
                 check_operator_row_limit(right_set_.size(), kDefaultMaxOperatorRows, "SetOpOperator");
                 right_set_.insert(row_key(r));
             }
-            right_->close();
+            children_[1]->close();
+            right_closed_ = true;
         }
     }
 
     bool next(Row& out) override {
+        if (children_.empty()) return false;
+
         if (op_ == SET_OP_UNION && !all_) {
-            // UNION (deduplicated)
-            while (true) {
-                bool got = false;
-                if (reading_left_) {
-                    got = left_->next(out);
-                    if (!got) { reading_left_ = false; }
-                }
-                if (!reading_left_) {
-                    got = right_->next(out);
-                    if (!got) return false;
+            while (child_idx_ < children_.size()) {
+                if (!children_[child_idx_]->next(out)) {
+                    ++child_idx_;
+                    continue;
                 }
-                if (got) {
-                    check_col_count(out);
-                    std::string key = row_key(out);
-                    if (seen_.find(key) == seen_.end()) {
-                        check_operator_row_limit(seen_.size(), kDefaultMaxOperatorRows, "SetOpOperator");
-                    }
-                    if (seen_.insert(key).second) return true;
+                check_col_count(out);
+                std::string key = row_key(out);
+                if (seen_.find(key) == seen_.end()) {
+                    check_operator_row_limit(seen_.size(), kDefaultMaxOperatorRows, "SetOpOperator");
                 }
+                if (seen_.insert(key).second) return true;
             }
+            return false;
         }
 
         if (op_ == SET_OP_UNION && all_) {
-            // UNION ALL: yield left then right
-            if (reading_left_) {
-                if (left_->next(out)) {
+            while (child_idx_ < children_.size()) {
+                if (children_[child_idx_]->next(out)) {
                     check_col_count(out);
                     return true;
                 }
-                reading_left_ = false;
-            }
-            if (right_->next(out)) {
-                check_col_count(out);
-                return true;
+                ++child_idx_;
             }
             return false;
         }
 
         if (op_ == SET_OP_INTERSECT) {
-            // Yield left rows that also appear in right
-            while (left_->next(out)) {
+            while (children_[0]->next(out)) {
                 check_col_count(out);
                 std::string key = row_key(out);
                 if (right_set_.count(key)) {
@@ -108,8 +104,7 @@ class SetOpOperator : public Operator {
         }
 
         if (op_ == SET_OP_EXCEPT) {
-            // Yield left rows that don't appear in right
-            while (left_->next(out)) {
+            while (children_[0]->next(out)) {
                 check_col_count(out);
                 std::string key = row_key(out);
                 if (!right_set_.count(key)) {
@@ -124,20 +119,22 @@ class SetOpOperator : public Operator {
     }
 
     void close() override {
-        left_->close();
-        right_->close();
+        for (size_t i = 0; i < children_.size(); ++i) {
+            if (right_closed_ && i == 1) continue;
+            children_[i]->close();
+        }
         seen_.clear();
         right_set_.clear();
     }
 
 private:
-    Operator* left_;
-    Operator* right_;
+    std::vector<Operator*> children_;
     uint8_t op_;
     bool all_;
     bool parallel_open_;
     ThreadPool* pool_ = nullptr;
-    bool reading_left_ = true;
+    size_t child_idx_ = 0;
+    bool right_closed_ = false;
     std::unordered_set<std::string> seen_;
     std::unordered_set<std::string> right_set_;
     // Column count established by the first row we see. Used to detect

include/sql_engine/operators/set_op_op.h

@@ -1171,11 +1171,29 @@ class PlanExecutor {
     }
 
     Operator* build_set_op(PlanNode* node) {
+        std::vector<PlanNode*> remote_leaves;
+        if (node->set_op.op == SET_OP_UNION && node->set_op.all &&
+            collect_union_all_remotes(node, remote_leaves) &&
+            remote_leaves.size() > 2) {
+            std::vector<Operator*> children;
+            children.reserve(remote_leaves.size());
+            for (PlanNode* leaf : remote_leaves) {
+                Operator* child = build_remote_scan(leaf);
+                if (!child) return nullptr;
+                children.push_back(child);
+            }
+            bool parallel = parallel_open_enabled_ && children.size() > 1;
+            auto op = std::make_unique<SetOpOperator>(
+                std::move(children), parallel, parallel ? pool_ : nullptr);
+            Operator* ptr = op.get();
+            operators_.push_back(std::move(op));
+            return ptr;
+        }
+
         Operator* left = build_operator(node->left);
         Operator* right = build_operator(node->right);
         if (!left || !right) return nullptr;
 
-        // Enable parallel open when both children are remote scans and executor is thread-safe
         bool parallel = parallel_open_enabled_ &&
                          (node->left && node->left->type == PlanNodeType::REMOTE_SCAN &&
                           node->right && node->right->type == PlanNodeType::REMOTE_SCAN);
@@ -1187,6 +1205,21 @@ class PlanExecutor {
         return ptr;
     }
 
+    static bool collect_union_all_remotes(const PlanNode* node,
+                                          std::vector<PlanNode*>& out) {
+        if (!node) return false;
+        if (node->type == PlanNodeType::REMOTE_SCAN) {
+            out.push_back(const_cast<PlanNode*>(node));
+            return true;
+        }
+        if (node->type == PlanNodeType::SET_OP &&
+            node->set_op.op == SET_OP_UNION && node->set_op.all) {
+            return collect_union_all_remotes(node->left, out) &&
+                   collect_union_all_remotes(node->right, out);
+        }
+        return false;
+    }
+
     Operator* build_remote_scan(PlanNode* node) {
         if (!remote_executor_) return nullptr;
         sql_parser::StringRef sql{node->remote_scan.remote_sql,

include/sql_engine/plan_executor.h

@@ -126,15 +126,16 @@ class Session {
         std::string sql_key(sql, len);
         auto cache_it = plan_cache_.find(sql_key);
         if (cache_it != plan_cache_.end()) {
-            // Cache hit: move this entry to the front of the LRU list.
             plan_cache_order_.splice(plan_cache_order_.begin(),
                                      plan_cache_order_,
                                      cache_it->second);
             exec_arena_.reset();
             auto& entry = *cache_it->second;
+            PlanNode* plan = maybe_distribute(entry.plan, exec_arena_);
+            if (!plan) return {};
             PlanExecutor<D> executor(functions_, catalog_, exec_arena_);
             wire_executor(executor);
-            return executor.execute(entry.plan);
+            return executor.execute(plan);
         }
 
         // Cache miss: full parse -> plan -> optimize -> distribute pipeline
@@ -164,25 +165,21 @@ class Session {
 
         plan = optimizer_.optimize(plan, cached_parser->arena());
 
-        // Distribute across shards if shard map is configured
+        ResultSet rs;
         if (shard_map_ && remote_executor_) {
-            DistributedPlanner<D> dplanner(*shard_map_, catalog_, cached_parser->arena(), remote_executor_, &functions_);
-            plan = dplanner.distribute(plan);
+            exec_arena_.reset();
+            PlanNode* dist = maybe_distribute(plan, exec_arena_);
+            if (!dist) return {};
+            PlanExecutor<D> executor(functions_, catalog_, exec_arena_);
+            wire_executor(executor);
+            rs = executor.execute(dist);
+        } else {
+            PlanExecutor<D> executor(functions_, catalog_, cached_parser->arena());
+            wire_executor(executor);
+            rs = executor.execute(plan);
         }
 
-        // Execute first using the parser's arena. preprocess_aggregates (called
-        // inside execute()) may allocate into the arena to modify the plan in-place.
-        // Those allocations must persist in the parser arena (not exec_arena_) so
-        // the cached plan remains valid across calls.
-        PlanExecutor<D> executor(functions_, catalog_, cached_parser->arena());
-        wire_executor(executor);
-        ResultSet rs = executor.execute(plan);
-
-        // Cache the plan, enforcing the LRU bound. The parser arena is kept
-        // alive via unique_ptr stored in the CachedPlan so all plan/AST
-        // string pointers remain valid for subsequent cache hits.
         insert_into_plan_cache(std::move(sql_key), std::move(cached_parser), plan);
-
         return rs;
     }
 
@@ -378,6 +375,15 @@ class Session {
     std::unordered_map<std::string, CacheIter> plan_cache_;
     size_t plan_cache_max_size_ = 1024;
 
+    PlanNode* maybe_distribute(PlanNode* plan, sql_parser::Arena& arena) {
+        if (!plan || !shard_map_ || !remote_executor_) return plan;
+        DistributedPlanner<D> dplanner(*shard_map_, catalog_, arena,
+                                       remote_executor_, &functions_);
+        PlanNode* dist = dplanner.distribute(plan);
+        if (dplanner.last_error()) return nullptr;
+        return dist;
+    }
+
     void insert_into_plan_cache(std::string key,
                                 std::unique_ptr<sql_parser::Parser<D>> parser,
                                 PlanNode* plan) {