fix: resolve distributed aggregation, sort merge, join columns, and subquery across shards

Bug 1 - Aggregation missing COUNT column: constant folding was replacing
aggregate functions (COUNT, SUM, AVG, etc.) with NULL when they had no
column references (e.g., COUNT(*)). Fixed by treating aggregate functions
as non-foldable in the constant folding rule. Also added output_exprs to
MERGE_AGGREGATE for proper column naming.

Bug 2 - Sort+Limit returning wrong results: MySQL DECIMAL values were
stored as TAG_STRING instead of TAG_DECIMAL, causing string comparison
("90000" > "110000") instead of numeric. Fixed by returning TAG_DECIMAL
from the MySQL executor and adding TAG_DECIMAL handling to Value::to_double()
and Value::to_int64().

Bug 3 - Cross-shard JOIN columns NULL: collect_tables() recursed into both
sides of SET_OP (UNION ALL) nodes, duplicating table entries and inflating
column offsets so the PROJECT resolver addressed beyond the actual row width.
Fixed by only collecting tables from the left side of SET_OP nodes.

Bug 4 - Distributed subquery empty: SubqueryExecutor built and executed
subquery plans without going through the distributed planner or setting
up the remote executor. Fixed by adding a distribute callback to
PlanExecutor that Session wires up, and propagating remote_executor to
inner subquery executors. Also prevented filter-with-subquery pushdown
to remote shards.

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -116,7 +116,10 @@ class DistributedPlanner {
 
             case PlanNodeType::FILTER: {
                 // Check if child is a SCAN -- push filter to remote
-                if (node->left && node->left->type == PlanNodeType::SCAN) {
+                // (but not if the filter contains a subquery, which must be
+                // evaluated locally after distributed subquery execution)
+                if (node->left && node->left->type == PlanNodeType::SCAN &&
+                    !has_subquery(node->filter.expr)) {
                     return distribute_scan(node->left, node->filter.expr,
                                            nullptr, nullptr, nullptr, false);
                 }
@@ -140,8 +143,6 @@ class DistributedPlanner {
                     // Extract aggregate info from the PROJECT select list
                     push_agg_exprs_from_project(node, agg_child);
                     PlanNode* dist_agg = distribute_aggregate(agg_child);
-                    // If it was distributed (MERGE_AGGREGATE), no need for PROJECT wrapper
-                    // because the merge already produces the right columns
                     if (dist_agg && dist_agg->type == PlanNodeType::MERGE_AGGREGATE) {
                         // Re-add FILTER (HAVING) if present
                         if (node->left && node->left->type == PlanNodeType::FILTER) {
@@ -583,6 +584,18 @@ class DistributedPlanner {
             arena_.allocate(merge_ops.size()));
         std::memcpy(merge->merge_aggregate.merge_ops, merge_ops.data(), merge_ops.size());
 
+        // Store original output expressions for column naming:
+        // group_by expressions + aggregate expressions
+        uint16_t out_count = agg_node->aggregate.group_count + agg_node->aggregate.agg_count;
+        auto** out_exprs = static_cast<const sql_parser::AstNode**>(
+            arena_.allocate(sizeof(sql_parser::AstNode*) * out_count));
+        for (uint16_t i = 0; i < agg_node->aggregate.group_count; ++i)
+            out_exprs[i] = agg_node->aggregate.group_by[i];
+        for (uint16_t i = 0; i < agg_node->aggregate.agg_count; ++i)
+            out_exprs[agg_node->aggregate.group_count + i] = agg_node->aggregate.agg_exprs[i];
+        merge->merge_aggregate.output_exprs = out_exprs;
+        merge->merge_aggregate.output_expr_count = out_count;
+
         // Set left to first child for compatibility with tree walkers
         if (!children.empty()) merge->left = children[0];
 

include/sql_engine/plan_executor.h

@@ -48,6 +48,7 @@
 #include "sql_engine/dml_result.h"
 #include "sql_engine/mutable_data_source.h"
 #include "sql_engine/catalog_resolver.h"
+#include "sql_parser/emitter.h"
 
 #include <unordered_map>
 #include <string>
@@ -81,6 +82,11 @@ class PlanExecutor {
     // Access the subquery executor (for operators that need it)
     SubqueryExecutor<D>* subquery_executor() { return &subquery_exec_; }
 
+    // Set a distribute callback for subquery plans (used by Session to inject
+    // distributed planning into subquery execution).
+    using DistributeFn = std::function<PlanNode*(PlanNode*)>;
+    void set_distribute_fn(DistributeFn fn) { distribute_fn_ = std::move(fn); }
+
     ResultSet execute(PlanNode* plan) {
         if (!plan) return {};
 
@@ -141,15 +147,22 @@ class PlanExecutor {
     std::unordered_map<std::string, MutableDataSource*> mutable_sources_;
     std::vector<std::unique_ptr<Operator>> operators_;
     RemoteExecutor* remote_executor_ = nullptr;
+    DistributeFn distribute_fn_;
     SubqueryExecutor<D> subquery_exec_;
     sql_parser::Arena subquery_plan_arena_{65536, 1048576};
     std::function<Value(sql_parser::StringRef)> outer_resolver_;
 
     void setup_subquery_executor() {
-        // Build plan callback: uses PlanBuilder with our catalog and arena
+        // Build plan callback: uses PlanBuilder with our catalog and arena.
+        // If a distribute function is set (for distributed/sharded execution),
+        // apply it to the plan so subqueries also go through the distributed planner.
         subquery_exec_.set_build_plan([this](const sql_parser::AstNode* ast) -> PlanNode* {
             PlanBuilder<D> builder(catalog_, arena_);
-            return builder.build(ast);
+            PlanNode* plan = builder.build(ast);
+            if (plan && distribute_fn_) {
+                plan = distribute_fn_(plan);
+            }
+            return plan;
         });
         // Execute plan callback: create a fresh executor for the subquery
         // to avoid interfering with the outer operator tree.
@@ -169,6 +182,8 @@ class PlanExecutor {
                     inner_exec.add_data_source(kv.first.c_str(), kv.second);
                 }
             }
+            if (remote_executor_)
+                inner_exec.set_remote_executor(remote_executor_);
             return inner_exec.execute(plan);
         });
         // Correlated execution: pass outer resolver as fallback.
@@ -192,6 +207,8 @@ class PlanExecutor {
                         inner_exec.add_data_source(kv.first.c_str(), kv.second);
                     }
                 }
+                if (remote_executor_)
+                    inner_exec.set_remote_executor(remote_executor_);
                 inner_exec.set_outer_resolver(outer_resolve);
                 return inner_exec.execute(plan);
             });
@@ -503,6 +520,13 @@ class PlanExecutor {
                 collect_tables(node->merge_sort.children[i], tables);
             return;
         }
+        // For SET_OP (UNION ALL), only collect from left side to avoid
+        // duplicating table entries when the same table appears on both
+        // sides (e.g., sharded UNION ALL of RemoteScans for the same table).
+        if (node->type == PlanNodeType::SET_OP) {
+            collect_tables(node->left, tables);
+            return;
+        }
         collect_tables(node->left, tables);
         collect_tables(node->right, tables);
     }
@@ -536,8 +560,10 @@ class PlanExecutor {
                 if (node->remote_scan.table) return node->remote_scan.table->column_count;
                 return 0;
             case PlanNodeType::MERGE_AGGREGATE:
-                // group keys + output agg columns
-                return count_columns(node->left);
+                // Use stored output expression count if available
+                if (node->merge_aggregate.output_expr_count > 0)
+                    return node->merge_aggregate.output_expr_count;
+                return node->merge_aggregate.group_key_count + node->merge_aggregate.merge_op_count;
             case PlanNodeType::MERGE_SORT:
                 return count_columns(node->left);
             case PlanNodeType::INSERT_PLAN:
@@ -970,6 +996,31 @@ class PlanExecutor {
                     build_column_names(plan->derived_scan.inner_plan, rs);
                 }
                 break;
+            case PlanNodeType::MERGE_AGGREGATE: {
+                // Use the stored output expressions for column naming
+                if (plan->merge_aggregate.output_exprs &&
+                    plan->merge_aggregate.output_expr_count > 0) {
+                    for (uint16_t i = 0; i < plan->merge_aggregate.output_expr_count; ++i) {
+                        const sql_parser::AstNode* expr = plan->merge_aggregate.output_exprs[i];
+                        if (expr) {
+                            // Use the emitter to produce a readable name
+                            sql_parser::Emitter<D> emitter(arena_);
+                            emitter.emit(expr);
+                            sql_parser::StringRef name = emitter.result();
+                            if (name.ptr && name.len > 0) {
+                                rs.column_names.emplace_back(name.ptr, name.len);
+                            } else {
+                                rs.column_names.push_back("?column?");
+                            }
+                        } else {
+                            rs.column_names.push_back("?column?");
+                        }
+                    }
+                } else {
+                    build_column_names(plan->left, rs);
+                }
+                break;
+            }
             default:
                 // For wrapping operators, recurse to child
                 build_column_names(plan->left, rs);

include/sql_engine/plan_node.h

@@ -113,6 +113,9 @@ struct PlanNode {
             uint8_t* merge_ops;       // parallel to agg columns
             uint16_t merge_op_count;
             uint16_t group_key_count; // number of leading group-by columns
+            // Original output column expressions (for column naming)
+            const sql_parser::AstNode** output_exprs;
+            uint16_t output_expr_count;
         } merge_aggregate;
 
         struct {

include/sql_engine/rules/constant_folding.h

@@ -23,7 +23,17 @@ namespace rules {
 
 namespace detail_cf {
 
-// Check if an expression has any column references
+// Check if a function name is an aggregate function
+inline bool is_aggregate_function(const sql_parser::AstNode* expr) {
+    if (!expr || expr->type != sql_parser::NodeType::NODE_FUNCTION_CALL) return false;
+    sql_parser::StringRef name = expr->value();
+    return name.equals_ci("COUNT", 5) || name.equals_ci("SUM", 3) ||
+           name.equals_ci("AVG", 3) || name.equals_ci("MIN", 3) ||
+           name.equals_ci("MAX", 3);
+}
+
+// Check if an expression has any column references or aggregate functions
+// (aggregate functions depend on row data and must not be folded)
 inline bool has_column_ref(const sql_parser::AstNode* expr) {
     if (!expr) return false;
     switch (expr->type) {
@@ -33,6 +43,10 @@ inline bool has_column_ref(const sql_parser::AstNode* expr) {
         case sql_parser::NodeType::NODE_IDENTIFIER:
             // Identifiers in expression context are column references
             return true;
+        case sql_parser::NodeType::NODE_FUNCTION_CALL:
+            // Aggregate functions depend on row data, not foldable
+            if (is_aggregate_function(expr)) return true;
+            break;
         default:
             break;
     }

include/sql_engine/session.h

@@ -203,6 +203,16 @@ class Session {
             executor.add_mutable_data_source(kv.first.c_str(), kv.second);
         if (remote_executor_)
             executor.set_remote_executor(remote_executor_);
+        // If sharding is configured, provide a distribute callback so that
+        // subqueries also go through the distributed planner.
+        if (shard_map_ && remote_executor_) {
+            executor.set_distribute_fn(
+                [this](PlanNode* plan) -> PlanNode* {
+                    DistributedPlanner<D> dp(*shard_map_, catalog_, parser_.arena(),
+                                             remote_executor_, &functions_);
+                    return dp.distribute(plan);
+                });
+        }
     }
 };
 

include/sql_engine/value.h

@@ -4,6 +4,7 @@
 #include "sql_engine/types.h"
 #include "sql_parser/common.h"  // StringRef
 #include <cstdint>
+#include <cstdlib>
 #include <type_traits>
 
 namespace sql_engine {
@@ -55,6 +56,15 @@ struct Value {
             case TAG_INT64:  return static_cast<double>(int_val);
             case TAG_UINT64: return static_cast<double>(uint_val);
             case TAG_DOUBLE: return double_val;
+            case TAG_DECIMAL:
+                if (str_val.ptr && str_val.len > 0) {
+                    char buf[128];
+                    uint32_t n = str_val.len < 127 ? str_val.len : 127;
+                    for (uint32_t i = 0; i < n; ++i) buf[i] = str_val.ptr[i];
+                    buf[n] = '\0';
+                    return std::strtod(buf, nullptr);
+                }
+                return 0.0;
             default:         return 0.0;
         }
     }
@@ -66,6 +76,7 @@ struct Value {
             case TAG_INT64:  return int_val;
             case TAG_UINT64: return static_cast<int64_t>(uint_val);
             case TAG_DOUBLE: return static_cast<int64_t>(double_val);
+            case TAG_DECIMAL: return static_cast<int64_t>(to_double());
             default:         return 0;
         }
     }

src/sql_engine/mysql_remote_executor.cpp

@@ -165,9 +165,9 @@ Value MySQLRemoteExecutor::mysql_field_to_value(
 
         case MYSQL_TYPE_DECIMAL:
         case MYSQL_TYPE_NEWDECIMAL: {
-            // Store as string in arena
+            // Store as decimal (numeric string) in arena
             sql_parser::StringRef s = arena_.allocate_string(data, static_cast<uint32_t>(length));
-            return value_string(s);
+            return value_decimal(s);
         }
 
         case MYSQL_TYPE_DATE: {

tests/test_mysql_executor.cpp

@@ -112,8 +112,8 @@ TEST_F(MySQLExecutorTest, DecimalType) {
     sql_parser::StringRef sql{q, static_cast<uint32_t>(strlen(q))};
     auto rs = exec_->execute("test_mysql", sql);
     ASSERT_EQ(rs.row_count(), 1u);
-    // DECIMAL comes as string
-    EXPECT_EQ(rs.rows[0].get(0).tag, sql_engine::Value::TAG_STRING);
+    // DECIMAL comes as numeric decimal type
+    EXPECT_EQ(rs.rows[0].get(0).tag, sql_engine::Value::TAG_DECIMAL);
 }
 
 TEST_F(MySQLExecutorTest, NullHandling) {