Merge pull request #57 from ProxySQL/fix/shard-planner-correctness

fix: stop silent wrong answers in distributed planner

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -178,6 +178,18 @@ class DistributedTransactionManager : public TransactionManager {
         return executor_.execute_dml(backend_name, sql);
     }
 
+    ResultSet route_query(const char* backend_name,
+                          sql_parser::StringRef sql) override {
+        if (!active_) return executor_.execute(backend_name, sql);
+        auto it = sessions_.find(backend_name);
+        if (it != sessions_.end() && it->second) {
+            return it->second->execute(sql);
+        }
+        return executor_.execute(backend_name, sql);
+    }
+
+    bool route_query_supported() const override { return true; }
+
     bool commit() override {
         if (!active_) return false;
         if (participants_.empty()) {

include/sql_engine/distributed_txn.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_engine/operators/aggregate_op.h

@@ -25,6 +25,7 @@
 #include "sql_parser/common.h"
 #include "sql_parser/arena.h"
 #include <cstring>
+#include <cstdlib>
 #include <vector>
 
 namespace sql_engine {
@@ -116,6 +117,55 @@ class PlanBuilder {
         return false;
     }
 
+    static const sql_parser::AstNode* select_item_expr(const sql_parser::AstNode* item) {
+        return item ? item->first_child : nullptr;
+    }
+
+    static sql_parser::StringRef select_item_alias(const sql_parser::AstNode* item) {
+        if (!item) return {};
+        for (const sql_parser::AstNode* c = item->first_child; c; c = c->next_sibling) {
+            if (c->type == sql_parser::NodeType::NODE_ALIAS) return c->value();
+        }
+        return {};
+    }
+
+    static const sql_parser::AstNode* resolve_order_key(
+            const sql_parser::AstNode* key, const sql_parser::AstNode* select_items) {
+        if (!key || !select_items) return key;
+        uint16_t n = count_children(select_items);
+        if (n == 0) return key;
+
+        if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = key->value();
+            if (!sv.ptr || sv.len == 0) return key;
+            int64_t pos = std::strtoll(sv.ptr, nullptr, 10);
+            if (pos < 1 || pos > static_cast<int64_t>(n)) return key;
+            uint16_t idx = 0;
+            for (const sql_parser::AstNode* item = select_items->first_child; item;
+                 item = item->next_sibling, ++idx) {
+                if (idx + 1 == static_cast<uint16_t>(pos)) {
+                    const sql_parser::AstNode* expr = select_item_expr(item);
+                    return expr ? expr : key;
+                }
+            }
+            return key;
+        }
+
+        if (key->type == sql_parser::NodeType::NODE_COLUMN_REF ||
+            key->type == sql_parser::NodeType::NODE_IDENTIFIER) {
+            sql_parser::StringRef name = key->value();
+            for (const sql_parser::AstNode* item = select_items->first_child; item;
+                 item = item->next_sibling) {
+                sql_parser::StringRef alias = select_item_alias(item);
+                if (alias.ptr && alias.equals_ci(name.ptr, name.len)) {
+                    const sql_parser::AstNode* expr = select_item_expr(item);
+                    return expr ? expr : key;
+                }
+            }
+        }
+        return key;
+    }
+
     // Check if an expression (or any descendant) contains an aggregate function call.
     // Does NOT recurse into subqueries -- aggregates inside subqueries belong
     // to the subquery's own aggregation, not the outer query.
@@ -279,10 +329,12 @@ class PlanBuilder {
                 arena_.allocate(sizeof(sql_parser::AstNode*) * cnt));
             auto* dirs = static_cast<uint8_t*>(arena_.allocate(cnt));
 
+            const sql_parser::AstNode* select_items =
+                find_child(select_ast, sql_parser::NodeType::NODE_SELECT_ITEM_LIST);
+
             uint16_t idx = 0;
             for (const sql_parser::AstNode* item = order_by->first_child; item; item = item->next_sibling) {
-                // First child is the key expression
-                keys[idx] = item->first_child;
+                keys[idx] = resolve_order_key(item->first_child, select_items);
                 // Check for DESC direction (second child with "DESC" value)
                 dirs[idx] = 0; // ASC by default
                 const sql_parser::AstNode* dir_node = find_child(item, sql_parser::NodeType::NODE_IDENTIFIER);

include/sql_engine/plan_builder.h

@@ -56,6 +56,8 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <cstdlib>
+#include <cstring>
 
 namespace sql_engine {
 
@@ -846,15 +848,128 @@ class PlanExecutor {
         return ptr;
     }
 
+    static bool same_agg_call(const sql_parser::AstNode* a, const sql_parser::AstNode* b) {
+        if (!a || !b) return false;
+        if (a->type != sql_parser::NodeType::NODE_FUNCTION_CALL ||
+            b->type != sql_parser::NodeType::NODE_FUNCTION_CALL) return false;
+        return a->value().equals_ci(b->value().ptr, b->value().len);
+    }
+
+    const sql_parser::AstNode* rewrite_having_expr(const sql_parser::AstNode* expr,
+                                                    PlanNode* agg_node) {
+        if (!expr || !agg_node) return expr;
+        uint16_t group_count = 0;
+        uint16_t agg_count = 0;
+        const sql_parser::AstNode** agg_exprs = nullptr;
+        const sql_parser::AstNode** group_by = nullptr;
+        if (agg_node->type == PlanNodeType::AGGREGATE) {
+            group_count = agg_node->aggregate.group_count;
+            agg_count = agg_node->aggregate.agg_count;
+            agg_exprs = agg_node->aggregate.agg_exprs;
+            group_by = agg_node->aggregate.group_by;
+        } else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE) {
+            group_count = agg_node->merge_aggregate.group_key_count;
+            if (agg_node->merge_aggregate.output_exprs &&
+                agg_node->merge_aggregate.output_expr_count > group_count) {
+                agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
+                agg_count = static_cast<uint16_t>(
+                    agg_node->merge_aggregate.output_expr_count - group_count);
+                group_by = agg_node->merge_aggregate.output_exprs;
+            }
+        } else {
+            return expr;
+        }
+
+        if (expr->type == sql_parser::NodeType::NODE_FUNCTION_CALL && agg_exprs) {
+            for (uint16_t i = 0; i < agg_count; ++i) {
+                if (same_agg_call(expr, agg_exprs[i])) {
+                    sql_parser::StringRef name = expr->value();
+                    return sql_parser::make_node(
+                        arena_, sql_parser::NodeType::NODE_IDENTIFIER, name);
+                }
+            }
+        }
+
+        bool changed = false;
+        sql_parser::AstNode* clone = sql_parser::make_node(
+            arena_, expr->type, expr->value(), expr->flags);
+        for (const sql_parser::AstNode* c = expr->first_child; c; c = c->next_sibling) {
+            const sql_parser::AstNode* rw = rewrite_having_expr(c, agg_node);
+            if (rw != c) changed = true;
+            if (rw) clone->add_child(const_cast<sql_parser::AstNode*>(rw));
+        }
+        (void)group_count;
+        (void)group_by;
+        return changed ? clone : expr;
+    }
+
+    const TableInfo* make_agg_output_table(PlanNode* agg_node) {
+        if (!agg_node) return nullptr;
+        uint16_t group_count = 0;
+        uint16_t agg_count = 0;
+        const sql_parser::AstNode** group_by = nullptr;
+        const sql_parser::AstNode** agg_exprs = nullptr;
+        if (agg_node->type == PlanNodeType::AGGREGATE) {
+            group_count = agg_node->aggregate.group_count;
+            agg_count = agg_node->aggregate.agg_count;
+            group_by = agg_node->aggregate.group_by;
+            agg_exprs = agg_node->aggregate.agg_exprs;
+        } else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE &&
+                   agg_node->merge_aggregate.output_exprs) {
+            group_count = agg_node->merge_aggregate.group_key_count;
+            group_by = agg_node->merge_aggregate.output_exprs;
+            if (agg_node->merge_aggregate.output_expr_count > group_count) {
+                agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
+                agg_count = static_cast<uint16_t>(
+                    agg_node->merge_aggregate.output_expr_count - group_count);
+            }
+        } else {
+            return nullptr;
+        }
+
+        uint16_t n = static_cast<uint16_t>(group_count + agg_count);
+        if (n == 0) return nullptr;
+        auto* cols = static_cast<ColumnInfo*>(arena_.allocate(sizeof(ColumnInfo) * n));
+        if (!cols) return nullptr;
+        for (uint16_t i = 0; i < group_count; ++i) {
+            cols[i].ordinal = i;
+            cols[i].nullable = true;
+            cols[i].type = SqlType::make_int();
+            cols[i].name = (group_by && group_by[i]) ? group_by[i]->value()
+                                                     : sql_parser::StringRef{};
+        }
+        for (uint16_t i = 0; i < agg_count; ++i) {
+            cols[group_count + i].ordinal = static_cast<uint16_t>(group_count + i);
+            cols[group_count + i].nullable = true;
+            cols[group_count + i].type = SqlType::make_int();
+            cols[group_count + i].name = (agg_exprs && agg_exprs[i])
+                ? agg_exprs[i]->value() : sql_parser::StringRef{};
+        }
+        auto* ti = static_cast<TableInfo*>(arena_.allocate(sizeof(TableInfo)));
+        if (!ti) return nullptr;
+        std::memset(ti, 0, sizeof(TableInfo));
+        ti->columns = cols;
+        ti->column_count = n;
+        return ti;
+    }
+
     Operator* build_filter(PlanNode* node) {
         Operator* child = build_operator(node->left);
         if (!child && node->left) return nullptr;
 
         std::vector<const TableInfo*> tables;
-        collect_tables(node->left, tables);
+        const sql_parser::AstNode* expr = node->filter.expr;
+        if (node->left && (node->left->type == PlanNodeType::AGGREGATE ||
+                           node->left->type == PlanNodeType::MERGE_AGGREGATE)) {
+            expr = rewrite_having_expr(expr, node->left);
+            const TableInfo* synth = make_agg_output_table(node->left);
+            if (synth) tables.push_back(synth);
+        } else {
+            collect_tables(node->left, tables);
+        }
 
         auto op = std::make_unique<FilterOperator<D>>(
-            child, node->filter.expr, catalog_, tables, functions_, arena_,
+            child, expr, catalog_, tables, functions_, arena_,
             &subquery_exec_, outer_resolver_);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));
@@ -1170,12 +1285,21 @@ class PlanExecutor {
 
     uint16_t resolve_column_index(const sql_parser::AstNode* key, const TableInfo* table) {
         if (!key || !table) return 0;
+        if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = key->value();
+            if (!sv.ptr || sv.len == 0) return 0;
+            int64_t n = std::strtoll(sv.ptr, nullptr, 10);
+            if (n < 1) return 0;
+            if (n > static_cast<int64_t>(table->column_count)) {
+                return static_cast<uint16_t>(table->column_count - 1);
+            }
+            return static_cast<uint16_t>(n - 1);
+        }
         sql_parser::StringRef col_name;
         if (key->type == sql_parser::NodeType::NODE_COLUMN_REF ||
             key->type == sql_parser::NodeType::NODE_IDENTIFIER) {
             col_name = key->value();
         } else if (key->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
-            // table.column -- get the column part
             const sql_parser::AstNode* c = key->first_child;
             if (c && c->next_sibling) col_name = c->next_sibling->value();
             else if (c) col_name = c->value();

include/sql_engine/plan_executor.h

@@ -101,7 +101,7 @@ struct PlanNode {
         struct {
             const char* backend_name;
             const char* remote_sql;
-            uint16_t remote_sql_len;
+            uint32_t remote_sql_len;
             const TableInfo* table;       // expected result schema (for SELECT *)
             // Optional projection expressions used to derive result column
             // names when the remote SQL is not a passthrough SELECT *. When

include/sql_engine/plan_node.h

@@ -94,6 +94,28 @@ class RemoteQueryBuilder {
         return sb.finish();
     }
 
+    sql_parser::StringRef build_select_join(
+            const TableInfo* left,
+            const TableInfo* right,
+            const sql_parser::AstNode* on_expr,
+            const sql_parser::AstNode* where_expr)
+    {
+        sql_parser::StringBuilder sb(arena_, 512);
+        sb.append("SELECT * FROM ");
+        if (left) sb.append(left->table_name.ptr, left->table_name.len);
+        sb.append(" JOIN ");
+        if (right) sb.append(right->table_name.ptr, right->table_name.len);
+        if (on_expr) {
+            sb.append(" ON ");
+            emit_expr(on_expr, sb);
+        }
+        if (where_expr) {
+            sb.append(" WHERE ");
+            emit_expr(where_expr, sb);
+        }
+        return sb.finish();
+    }
+
     // Build an INSERT statement string.
     sql_parser::StringRef build_insert(
             const TableInfo* table,