fix: stop silent wrong answers in distributed planner

COUNT(DISTINCT) and unknown aggs no longer merge as SUM_OF_COUNTS.
They gather rows and aggregate locally, with DISTINCT honored in
AggregateOperator. GROUP BY keeps HAVING and ORDER BY. WINDOW is
distributed as gather-then-window so ORDER BY cannot drop it.

Adds mock coverage plus a live 2-shard INSERT-then-point-SELECT
gtest (skips without 13306/13307) and a sqlengine sharded INSERT check.

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -141,28 +141,38 @@ class DistributedPlanner {
             }
 
             case PlanNodeType::PROJECT: {
-                // Check for PROJECT -> [SORT ->] [FILTER ->] AGGREGATE pattern
-                // For aggregate queries, we want to handle the whole thing in distribute_aggregate
+                // PROJECT -> [SORT ->] [FILTER(HAVING) ->] AGGREGATE
+                PlanNode* sort_node = nullptr;
+                PlanNode* having_node = nullptr;
                 PlanNode* agg_child = node->left;
-                if (agg_child && agg_child->type == PlanNodeType::SORT)
+                if (agg_child && agg_child->type == PlanNodeType::SORT) {
+                    sort_node = agg_child;
                     agg_child = agg_child->left;
-                if (agg_child && agg_child->type == PlanNodeType::FILTER)
+                }
+                if (agg_child && agg_child->type == PlanNodeType::FILTER) {
+                    having_node = agg_child;
                     agg_child = agg_child->left;
+                }
                 if (agg_child && agg_child->type == PlanNodeType::AGGREGATE) {
-                    // Extract aggregate info from the PROJECT select list
                     push_agg_exprs_from_project(node, agg_child);
                     PlanNode* dist_agg = distribute_aggregate(agg_child);
-                    if (dist_agg && dist_agg->type == PlanNodeType::MERGE_AGGREGATE) {
-                        // Re-add FILTER (HAVING) if present
-                        if (node->left && node->left->type == PlanNodeType::FILTER) {
+                    if (dist_agg && (dist_agg->type == PlanNodeType::MERGE_AGGREGATE ||
+                                     dist_agg->type == PlanNodeType::AGGREGATE)) {
+                        PlanNode* top = dist_agg;
+                        if (having_node) {
                             PlanNode* having = make_plan_node(arena_, PlanNodeType::FILTER);
-                            having->filter.expr = node->left->filter.expr;
-                            having->left = dist_agg;
-                            return having;
+                            having->filter.expr = having_node->filter.expr;
+                            having->left = top;
+                            top = having;
+                        }
+                        if (sort_node) {
+                            PlanNode* sort = make_plan_node(arena_, PlanNodeType::SORT);
+                            sort->sort = sort_node->sort;
+                            sort->left = top;
+                            top = sort;
                         }
-                        return dist_agg;
+                        return top;
                     }
-                    // For unsharded, the remote already computes everything
                     return dist_agg;
                 }
 
@@ -189,6 +199,13 @@ class DistributedPlanner {
             case PlanNodeType::JOIN:
                 return distribute_join(node);
 
+            case PlanNodeType::WINDOW: {
+                PlanNode* result = make_plan_node(arena_, PlanNodeType::WINDOW);
+                result->window = node->window;
+                result->left = distribute_node(node->left);
+                return result;
+            }
+
             case PlanNodeType::SET_OP: {
                 PlanNode* result = make_plan_node(arena_, PlanNodeType::SET_OP);
                 result->set_op = node->set_op;
@@ -261,6 +278,24 @@ class DistributedPlanner {
         return ctx;
     }
 
+    static bool contains_type(const PlanNode* node, PlanNodeType type) {
+        if (!node) return false;
+        if (node->type == type) return true;
+        if (contains_type(node->left, type)) return true;
+        if (contains_type(node->right, type)) return true;
+        if (node->type == PlanNodeType::MERGE_AGGREGATE) {
+            for (uint16_t i = 0; i < node->merge_aggregate.child_count; ++i) {
+                if (contains_type(node->merge_aggregate.children[i], type)) return true;
+            }
+        }
+        if (node->type == PlanNodeType::MERGE_SORT) {
+            for (uint16_t i = 0; i < node->merge_sort.child_count; ++i) {
+                if (contains_type(node->merge_sort.children[i], type)) return true;
+            }
+        }
+        return false;
+    }
+
     // Case 1 & 2: Distribute a scan (possibly with filter pushed down)
     PlanNode* distribute_scan(PlanNode* scan_node,
                               const sql_parser::AstNode* where_expr,
@@ -548,10 +583,16 @@ class DistributedPlanner {
 
         const TableInfo* table = ctx.scan->scan.table;
         if (!shards_.has_table(table->table_name) || !shards_.is_sharded(table->table_name)) {
-            // Unsharded -- push the whole thing to remote
             return make_unsharded_aggregate(agg_node, ctx, table);
         }
 
+        if (!all_aggregates_two_phase(agg_node)) {
+            PlanNode* result = make_plan_node(arena_, PlanNodeType::AGGREGATE);
+            result->aggregate = agg_node->aggregate;
+            result->left = distribute_node(agg_node->left);
+            return result;
+        }
+
         // Sharded aggregate: each shard computes partial aggregates.
         // Build remote project expressions: group-by cols + partial agg expressions
         const auto& shard_list = shards_.get_shards(table->table_name);
@@ -661,6 +702,24 @@ class DistributedPlanner {
         return make_remote_scan_with_outputs(backend, sql, table, projs);
     }
 
+    static bool is_two_phase_aggregate(const sql_parser::AstNode* expr) {
+        if (!expr || expr->type != sql_parser::NodeType::NODE_FUNCTION_CALL) return false;
+        if (expr->flags & sql_parser::FLAG_FUNC_DISTINCT) 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);
+    }
+
+    bool all_aggregates_two_phase(const PlanNode* agg_node) const {
+        if (!agg_node) return false;
+        if (agg_node->aggregate.agg_count == 0) return true;
+        for (uint16_t i = 0; i < agg_node->aggregate.agg_count; ++i) {
+            if (!is_two_phase_aggregate(agg_node->aggregate.agg_exprs[i])) return false;
+        }
+        return true;
+    }
+
     void decompose_aggregate(const sql_parser::AstNode* expr,
                               std::vector<const sql_parser::AstNode*>& projs,
                               std::vector<uint8_t>& merge_ops) {
@@ -673,11 +732,9 @@ class DistributedPlanner {
         sql_parser::StringRef name = expr->value();
 
         if (name.equals_ci("COUNT", 5)) {
-            // Remote: COUNT(*) or COUNT(col), Local: SUM of counts
             projs.push_back(expr);
             merge_ops.push_back(static_cast<uint8_t>(MergeOp::SUM_OF_COUNTS));
         } else if (name.equals_ci("SUM", 3)) {
-            // Remote: SUM(col), Local: SUM of sums
             projs.push_back(expr);
             merge_ops.push_back(static_cast<uint8_t>(MergeOp::SUM_OF_SUMS));
         } else if (name.equals_ci("AVG", 3)) {
@@ -726,7 +783,16 @@ class DistributedPlanner {
 
     // Case 4: Distributed sort + limit
     PlanNode* distribute_sort(PlanNode* sort_node) {
-        // Check if the child is a scan (possibly through filter) on a sharded table
+        if (contains_type(sort_node->left, PlanNodeType::WINDOW) ||
+            contains_type(sort_node->left, PlanNodeType::DERIVED_SCAN) ||
+            contains_type(sort_node->left, PlanNodeType::AGGREGATE) ||
+            contains_type(sort_node->left, PlanNodeType::MERGE_AGGREGATE)) {
+            PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT);
+            result->sort = sort_node->sort;
+            result->left = distribute_node(sort_node->left);
+            return result;
+        }
+
         ScanContext ctx = extract_scan_context(sort_node->left);
         if (!ctx.scan || !ctx.scan->scan.table) {
             PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT);
@@ -795,6 +861,15 @@ class DistributedPlanner {
         // Check if child is Sort on sharded table
         if (limit_node->left && limit_node->left->type == PlanNodeType::SORT) {
             PlanNode* sort_node = limit_node->left;
+            if (contains_type(sort_node->left, PlanNodeType::WINDOW) ||
+                contains_type(sort_node->left, PlanNodeType::DERIVED_SCAN) ||
+                contains_type(sort_node->left, PlanNodeType::AGGREGATE) ||
+                contains_type(sort_node->left, PlanNodeType::MERGE_AGGREGATE)) {
+                PlanNode* result = make_plan_node(arena_, PlanNodeType::LIMIT);
+                result->limit = limit_node->limit;
+                result->left = distribute_node(limit_node->left);
+                return result;
+            }
             ScanContext ctx = extract_scan_context(sort_node->left);
             if (ctx.scan && ctx.scan->scan.table) {
                 const TableInfo* table = ctx.scan->scan.table;
@@ -839,7 +914,16 @@ class DistributedPlanner {
             }
         }
 
-        // Check if child is scan on sharded/unsharded table (limit without sort)
+        if (contains_type(limit_node->left, PlanNodeType::WINDOW) ||
+            contains_type(limit_node->left, PlanNodeType::DERIVED_SCAN) ||
+            contains_type(limit_node->left, PlanNodeType::AGGREGATE) ||
+            contains_type(limit_node->left, PlanNodeType::MERGE_AGGREGATE)) {
+            PlanNode* result = make_plan_node(arena_, PlanNodeType::LIMIT);
+            result->limit = limit_node->limit;
+            result->left = distribute_node(limit_node->left);
+            return result;
+        }
+
         ScanContext ctx = extract_scan_context(limit_node->left);
         if (ctx.scan && ctx.scan->scan.table) {
             const TableInfo* table = ctx.scan->scan.table;

include/sql_engine/operators/aggregate_op.h

@@ -12,6 +12,8 @@
 #include <string>
 #include <cstring>
 #include <cmath>
+#include <unordered_set>
+#include "sql_parser/common.h"
 
 namespace sql_engine {
 
@@ -131,6 +133,8 @@ class AggregateOperator : public Operator {
         Value max_val{};
         bool has_value = false;
         bool count_star = false; // COUNT(*)
+        bool distinct = false;
+        std::unordered_set<std::string> seen;
     };
 
     struct GroupState {
@@ -186,9 +190,9 @@ class AggregateOperator : public Operator {
 
         if (expr->type == sql_parser::NodeType::NODE_FUNCTION_CALL) {
             sql_parser::StringRef name = expr->value();
+            state.distinct = (expr->flags & sql_parser::FLAG_FUNC_DISTINCT) != 0;
             if (name.equals_ci("COUNT", 5)) {
                 state.type = AggType::COUNT;
-                // Check for COUNT(*)
                 const sql_parser::AstNode* arg = expr->first_child;
                 if (arg && arg->type == sql_parser::NodeType::NODE_ASTERISK) {
                     state.count_star = true;
@@ -203,25 +207,32 @@ class AggregateOperator : public Operator {
         state.type = AggType::EXPR;
     }
 
+    static bool note_distinct(AggState& state, const sql_parser::AstNode* expr,
+                             const Value& v) {
+        bool distinct = state.distinct ||
+            (expr && (expr->flags & sql_parser::FLAG_FUNC_DISTINCT));
+        if (!distinct) return true;
+        return state.seen.insert(value_to_string(v)).second;
+    }
+
     void update_agg(AggState& state, const sql_parser::AstNode* expr,
                     const std::function<Value(sql_parser::StringRef)>& resolver) {
         switch (state.type) {
             case AggType::COUNT: {
                 if (state.count_star) {
                     state.count++;
                 } else {
-                    // COUNT(expr) - count non-null values
                     const sql_parser::AstNode* arg = expr->first_child;
                     Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
-                    if (!v.is_null()) state.count++;
+                    if (!v.is_null() && note_distinct(state, expr, v)) state.count++;
                 }
                 break;
             }
             case AggType::SUM:
             case AggType::AVG: {
                 const sql_parser::AstNode* arg = expr->first_child;
                 Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
-                if (!v.is_null()) {
+                if (!v.is_null() && note_distinct(state, expr, v)) {
                     state.sum += v.to_double();
                     state.count++;
                     state.has_value = true;
@@ -231,7 +242,7 @@ class AggregateOperator : public Operator {
             case AggType::MIN: {
                 const sql_parser::AstNode* arg = expr->first_child;
                 Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
-                if (!v.is_null()) {
+                if (!v.is_null() && note_distinct(state, expr, v)) {
                     if (!state.has_value || compare_values(v, state.min_val) < 0) {
                         state.min_val = v;
                         state.has_value = true;
@@ -242,7 +253,7 @@ class AggregateOperator : public Operator {
             case AggType::MAX: {
                 const sql_parser::AstNode* arg = expr->first_child;
                 Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
-                if (!v.is_null()) {
+                if (!v.is_null() && note_distinct(state, expr, v)) {
                     if (!state.has_value || compare_values(v, state.max_val) > 0) {
                         state.max_val = v;
                         state.has_value = true;

include/sql_parser/common.h

@@ -66,6 +66,10 @@ static constexpr uint16_t FLAG_SET_OP_ALL = 0x01;
 // which matters for SHOW search_path / SHOW <var> canonical re-emission.
 static constexpr uint16_t FLAG_IDENT_DELIMITED = 0x01;
 
+// -- Flags for NODE_FUNCTION_CALL --
+// Set when the call was written as FN(DISTINCT ...).
+static constexpr uint16_t FLAG_FUNC_DISTINCT = 0x01;
+
 // -- Statement type (always set, even for PARTIAL/ERROR) --
 
 enum class StmtType : uint8_t {

include/sql_parser/emitter.h

@@ -1116,6 +1116,7 @@ class Emitter {
     void emit_function_call(const AstNode* node) {
         emit_value(node);
         sb_.append_char('(');
+        if (node->flags & FLAG_FUNC_DISTINCT) sb_.append("DISTINCT ");
         bool first = true;
         for (const AstNode* arg = node->first_child; arg; arg = arg->next_sibling) {
             if (!first) sb_.append(", ");

include/sql_parser/expression_parser.h

@@ -331,6 +331,10 @@ class ExpressionParser {
             // argument list. Model it as a function call so consumers can
             // reject or handle the expression without leaving valid input
             // unconsumed.
+            if (tok_.peek().type == TokenType::TK_DISTINCT) {
+                func->flags |= FLAG_FUNC_DISTINCT;
+                tok_.skip();
+            }
             if (name_token.text.equals_ci("CAST", 4)) {
                 AstNode* arg = parse();
                 if (!arg || tok_.peek().type != TokenType::TK_AS) return func;

scripts/test_sqlengine.sh

@@ -290,6 +290,13 @@ test_sharded() {
     # Total = 890000.
     out=$(run_sharded "SELECT SUM(salary) FROM users")
     assert_contains "sharded: SUM(salary) all users = 890000" "${out}" "890000"
+
+    # Engine INSERT then point-SELECT must agree (RANGE routing).
+    out=$(run_sharded "INSERT INTO users (id, name, age, dept, salary) VALUES (11, 'Zed', 40, 'Test', 1)")
+    assert_contains "sharded: INSERT id=11" "${out}" "Query OK, 1 row"
+    out=$(run_sharded "SELECT name FROM users WHERE id = 11")
+    assert_contains "sharded: point SELECT after INSERT (Zed)" "${out}" "Zed"
+    run_sharded "DELETE FROM users WHERE id = 11" >/dev/null
 }
 
 # ----------------------------------------------------------------------