From e742dfdc82b9ac64d1de06aeb303ade01e50044f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 02:52:54 +0700 Subject: [PATCH 1/2] fix: distribute derived tables and stop silent sort/SQL bugs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrite DERIVED_SCAN inner plans so FROM (SELECT ...) hits remote shards. Store remote_sql_len as uint32_t so statements longer than 64KB are not truncated. Resolve ORDER BY position/alias in the plan builder, and only MERGE_SORT when every key is a table column — expressions gather and sort locally instead of comparing column 0. --- include/sql_engine/distributed_planner.h | 70 +++++++++++++++-- include/sql_engine/plan_builder.h | 56 +++++++++++++- include/sql_engine/plan_executor.h | 12 ++- include/sql_engine/plan_node.h | 2 +- tests/test_distributed_planner.cpp | 97 ++++++++++++++++++++++++ tests/test_plan_executor.cpp | 10 +++ 6 files changed, 236 insertions(+), 11 deletions(-) diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 8cc0940..010d6df 100644 --- a/include/sql_engine/distributed_planner.h +++ b/include/sql_engine/distributed_planner.h @@ -206,6 +206,13 @@ class DistributedPlanner { return result; } + case PlanNodeType::DERIVED_SCAN: { + PlanNode* result = make_plan_node(arena_, PlanNodeType::DERIVED_SCAN); + result->derived_scan = node->derived_scan; + result->derived_scan.inner_plan = distribute(node->derived_scan.inner_plan); + return result; + } + case PlanNodeType::SET_OP: { PlanNode* result = make_plan_node(arena_, PlanNodeType::SET_OP); result->set_op = node->set_op; @@ -293,6 +300,9 @@ class DistributedPlanner { if (contains_type(node->merge_sort.children[i], type)) return true; } } + if (node->type == PlanNodeType::DERIVED_SCAN) { + return contains_type(node->derived_scan.inner_plan, type); + } return false; } @@ -545,7 +555,7 @@ class DistributedPlanner { std::memcpy(bn, backend, blen + 1); node->remote_scan.backend_name = bn; node->remote_scan.remote_sql = sql.ptr; - node->remote_scan.remote_sql_len = static_cast(sql.len); + node->remote_scan.remote_sql_len = sql.len; node->remote_scan.table = table; // Caller is responsible for setting output_exprs when the remote SQL // is not a passthrough SELECT *. make_plan_node() already zero-fills @@ -782,15 +792,53 @@ class DistributedPlanner { } // Case 4: Distributed sort + limit + PlanNode* local_sort(PlanNode* sort_node) { + PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT); + result->sort = sort_node->sort; + result->left = distribute_node(sort_node->left); + return result; + } + + int sort_key_table_ordinal(const sql_parser::AstNode* key, const TableInfo* table) const { + if (!key || !table) return -1; + if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) { + sql_parser::StringRef sv = key->value(); + if (!sv.ptr || sv.len == 0) return -1; + int64_t n = std::strtoll(sv.ptr, nullptr, 10); + if (n < 1 || n > static_cast(table->column_count)) return -1; + return static_cast(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) { + 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(); + } else { + return -1; + } + if (!col_name.ptr) return -1; + const ColumnInfo* col = catalog_.get_column(table, col_name); + if (!col) return -1; + return static_cast(col->ordinal); + } + + bool all_sort_keys_are_table_columns(const PlanNode* sort_node, const TableInfo* table) const { + if (!sort_node || !table) return false; + for (uint16_t i = 0; i < sort_node->sort.count; ++i) { + if (sort_key_table_ordinal(sort_node->sort.keys[i], table) < 0) return false; + } + return true; + } + PlanNode* distribute_sort(PlanNode* sort_node) { 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; + return local_sort(sort_node); } ScanContext ctx = extract_scan_context(sort_node->left); @@ -809,6 +857,10 @@ class DistributedPlanner { return result; } + if (!all_sort_keys_are_table_columns(sort_node, table)) { + return local_sort(sort_node); + } + if (!shards_.is_sharded(table->table_name)) { // Unsharded -- push sort to remote sql_parser::StringRef sql = qb_.build_select( @@ -875,8 +927,12 @@ class DistributedPlanner { const TableInfo* table = ctx.scan->scan.table; if (shards_.has_table(table->table_name) && shards_.is_sharded(table->table_name)) { - // Case 4: Sharded sort + limit - // Each shard: ORDER BY + LIMIT, MergeSort, then outer Limit + if (!all_sort_keys_are_table_columns(sort_node, table)) { + PlanNode* result = make_plan_node(arena_, PlanNodeType::LIMIT); + result->limit = limit_node->limit; + result->left = distribute_node(limit_node->left); + return result; + } int64_t remote_limit = limit_node->limit.count + limit_node->limit.offset; PlanNode* merge = make_sharded_merge_sort( diff --git a/include/sql_engine/plan_builder.h b/include/sql_engine/plan_builder.h index c31f653..700da0d 100644 --- a/include/sql_engine/plan_builder.h +++ b/include/sql_engine/plan_builder.h @@ -25,6 +25,7 @@ #include "sql_parser/common.h" #include "sql_parser/arena.h" #include +#include #include 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(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(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(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); diff --git a/include/sql_engine/plan_executor.h b/include/sql_engine/plan_executor.h index b38afb8..c92466b 100644 --- a/include/sql_engine/plan_executor.h +++ b/include/sql_engine/plan_executor.h @@ -56,6 +56,7 @@ #include #include #include +#include namespace sql_engine { @@ -1170,12 +1171,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(table->column_count)) { + return static_cast(table->column_count - 1); + } + return static_cast(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(); diff --git a/include/sql_engine/plan_node.h b/include/sql_engine/plan_node.h index 2a880c4..24d7412 100644 --- a/include/sql_engine/plan_node.h +++ b/include/sql_engine/plan_node.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 diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index 2daa5c7..7ed2607 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -390,6 +390,10 @@ class DistributedPlannerTest : public ::testing::Test { find_nodes(node->merge_sort.children[i], type, out); return; } + if (node->type == PlanNodeType::DERIVED_SCAN) { + find_nodes(node->derived_scan.inner_plan, type, out); + return; + } find_nodes(node->left, type, out); find_nodes(node->right, type, out); } @@ -964,3 +968,96 @@ TEST_F(DistributedPlannerTest, WindowGatherCorrectness) { EXPECT_EQ(dist_rs.row_count(), 15u); EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs)); } + +TEST_F(DistributedPlannerTest, DerivedScanIsDistributed) { + Parser parser; + const char* sql = "SELECT name FROM (SELECT name, age FROM users WHERE age > 20) AS t"; + auto pr = parser.parse(sql, std::strlen(sql)); + ASSERT_EQ(pr.status, ParseResult::OK); + + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + ASSERT_NE(plan, nullptr); + + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector scans, remotes; + find_nodes(dist, PlanNodeType::SCAN, scans); + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + EXPECT_TRUE(scans.empty()) << "inner SCAN must be rewritten"; + EXPECT_FALSE(remotes.empty()); +} + +TEST_F(DistributedPlannerTest, DerivedScanCorrectness) { + const char* sql = "SELECT name FROM (SELECT name, age FROM users WHERE age > 20) AS t"; + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_GT(local_rs.row_count(), 0u); + EXPECT_EQ(local_rs.row_count(), dist_rs.row_count()); + EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs)); +} + +TEST_F(DistributedPlannerTest, OrderByExpressionDoesNotMergeOnColumnZero) { + Parser parser; + const char* sql = "SELECT name, age FROM users ORDER BY age + 1"; + auto pr = parser.parse(sql, std::strlen(sql)); + ASSERT_EQ(pr.status, ParseResult::OK); + + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + ASSERT_NE(plan, nullptr); + + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector merges; + find_nodes(dist, PlanNodeType::MERGE_SORT, merges); + EXPECT_TRUE(merges.empty()) << "expression ORDER BY must not MERGE_SORT on col 0"; + + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_EQ(local_rs.row_count(), dist_rs.row_count()); + EXPECT_TRUE(compare_results_ordered(local_rs, dist_rs)); +} + +TEST_F(DistributedPlannerTest, OrderByAliasAndPositionCorrectness) { + const char* sql = "SELECT name AS n, age AS a FROM users ORDER BY a DESC"; + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_EQ(local_rs.row_count(), 15u); + EXPECT_TRUE(compare_results_ordered(local_rs, dist_rs)); + + const char* sql2 = "SELECT name, age FROM users ORDER BY 2 DESC"; + auto local2 = execute_local(sql2); + auto dist2 = execute_distributed(sql2); + EXPECT_TRUE(compare_results_ordered(local2, dist2)); +} + +TEST_F(DistributedPlannerTest, RemoteSqlLenIsNotTruncated) { + std::string name(70000, 'x'); + std::string sql = "SELECT * FROM users WHERE name = '" + name + "'"; + Parser parser; + auto pr = parser.parse(sql.c_str(), sql.size()); + ASSERT_EQ(pr.status, ParseResult::OK); + + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + ASSERT_NE(plan, nullptr); + + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + ASSERT_FALSE(remotes.empty()); + for (auto* rs : remotes) { + EXPECT_GT(rs->remote_scan.remote_sql_len, 65535u); + std::string remote(rs->remote_scan.remote_sql, rs->remote_scan.remote_sql_len); + EXPECT_NE(remote.find(name), std::string::npos); + EXPECT_EQ(rs->remote_scan.remote_sql_len, remote.size()); + } +} diff --git a/tests/test_plan_executor.cpp b/tests/test_plan_executor.cpp index 32e5206..06316d9 100644 --- a/tests/test_plan_executor.cpp +++ b/tests/test_plan_executor.cpp @@ -147,6 +147,16 @@ TEST_F(PlanExecutorTest, SelectDistinctDept) { EXPECT_EQ(depts.size(), 2u); } +TEST_F(PlanExecutorTest, OrderByAliasAndPosition) { + auto by_alias = run_query("SELECT name AS n, age AS a FROM users ORDER BY a DESC"); + ASSERT_EQ(by_alias.row_count(), 5u); + EXPECT_EQ(by_alias.rows[0].get(1).int_val, 35); + + auto by_pos = run_query("SELECT name, age FROM users ORDER BY 2 DESC"); + ASSERT_EQ(by_pos.row_count(), 5u); + EXPECT_EQ(by_pos.rows[0].get(1).int_val, 35); +} + TEST_F(PlanExecutorTest, CountDistinctDept) { parser.reset(); const char* sql = "SELECT COUNT(DISTINCT dept) FROM users"; From 886f741fa3aed4b1b8e56964c91ed46ee5036da6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 03:10:10 +0700 Subject: [PATCH 2/2] fix: close remaining sharding correctness holes Refuse multi-table UPDATE/DELETE when any table is sharded or backends differ. Evaluate HAVING against aggregate output so COUNT(*) filters work locally and distributed. Push equi-joins of same-layout sharded tables to each shard. sqlengine and mysql_server use 2PC when backends are configured; SELECTs inside an open distributed txn hit pinned sessions. Optional --txn-log on sqlengine. --- include/sql_engine/distributed_planner.h | 149 +++++++++++++++++++--- include/sql_engine/distributed_txn.h | 12 ++ include/sql_engine/plan_executor.h | 118 ++++++++++++++++- include/sql_engine/remote_query_builder.h | 22 ++++ include/sql_engine/session.h | 46 ++++++- include/sql_engine/shard_map.h | 28 ++++ include/sql_engine/transaction_manager.h | 8 ++ tests/test_distributed_dml.cpp | 14 ++ tests/test_distributed_planner.cpp | 38 ++++++ tests/test_plan_executor.cpp | 7 + tools/mysql_server.cpp | 6 +- tools/sqlengine.cpp | 29 ++++- 12 files changed, 448 insertions(+), 29 deletions(-) diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 010d6df..77976e5 100644 --- a/include/sql_engine/distributed_planner.h +++ b/include/sql_engine/distributed_planner.h @@ -999,12 +999,72 @@ class DistributedPlanner { return result; } - // Case 5: Cross-backend join + bool join_on_shard_keys(const sql_parser::AstNode* cond, + sql_parser::StringRef left_key, + sql_parser::StringRef right_key) const { + if (!cond || cond->type != sql_parser::NodeType::NODE_BINARY_OP) return false; + sql_parser::StringRef op = cond->value(); + if (op.len != 1 || op.ptr[0] != '=') return false; + const sql_parser::AstNode* l = cond->first_child; + const sql_parser::AstNode* r = l ? l->next_sibling : nullptr; + if (!l || !r) return false; + return (is_shard_key_ref(l, left_key) && is_shard_key_ref(r, right_key)) || + (is_shard_key_ref(l, right_key) && is_shard_key_ref(r, left_key)); + } + + PlanNode* distribute_colocated_join(PlanNode* join_node, + const TableInfo* left_table, + const TableInfo* right_table) { + ScanContext lctx = extract_scan_context(join_node->left); + ScanContext rctx = extract_scan_context(join_node->right); + const sql_parser::AstNode* where_expr = nullptr; + if (lctx.where_expr && rctx.where_expr) { + sql_parser::AstNode* and_node = sql_parser::make_node( + arena_, sql_parser::NodeType::NODE_BINARY_OP, + sql_parser::StringRef{"AND", 3}); + and_node->add_child(const_cast(lctx.where_expr)); + and_node->add_child(const_cast(rctx.where_expr)); + where_expr = and_node; + } else if (lctx.where_expr) { + where_expr = lctx.where_expr; + } else { + where_expr = rctx.where_expr; + } + + const auto& shard_list = shards_.get_shards(left_table->table_name); + PlanNode* current = nullptr; + for (const auto& shard : shard_list) { + sql_parser::StringRef sql = qb_.build_select_join( + left_table, right_table, join_node->join.condition, where_expr); + PlanNode* rs = make_remote_scan(shard.backend_name.c_str(), sql, left_table); + if (!current) { + current = rs; + } else { + PlanNode* union_node = make_plan_node(arena_, PlanNodeType::SET_OP); + union_node->set_op.op = SET_OP_UNION; + union_node->set_op.all = true; + union_node->left = current; + union_node->right = rs; + current = union_node; + } + } + return current ? current : join_node; + } + PlanNode* distribute_join(PlanNode* join_node) { - // Get tables from each side const TableInfo* left_table = find_table(join_node->left); const TableInfo* right_table = find_table(join_node->right); + if (left_table && right_table && + shards_.is_sharded(left_table->table_name) && + shards_.is_sharded(right_table->table_name) && + shards_.same_routing(left_table->table_name, right_table->table_name) && + join_on_shard_keys(join_node->join.condition, + shards_.get_shard_key(left_table->table_name), + shards_.get_shard_key(right_table->table_name))) { + return distribute_colocated_join(join_node, left_table, right_table); + } + PlanNode* left_dist = nullptr; PlanNode* right_dist = nullptr; @@ -1172,18 +1232,13 @@ class DistributedPlanner { PlanNode* distribute_update(PlanNode* plan) { const auto& up = plan->update_plan; const TableInfo* table = up.table; - if (!table || !shards_.has_table(table->table_name)) return plan; - // Multi-table UPDATE: emit full SQL from AST, route to primary table's backend if (up.original_ast) { - sql_parser::StringRef sql = qb_.build_update_from_ast(up.original_ast); - if (!shards_.is_sharded(table->table_name)) { - return make_remote_scan(shards_.get_backend(table->table_name), sql, table); - } - const auto& shard_list = shards_.get_shards(table->table_name); - return scatter_dml_to_shards(table, shard_list, [&]() { return sql; }); + return distribute_multi_table_dml(up.original_ast, table, true); } + if (!table || !shards_.has_table(table->table_name)) return plan; + // Check for cross-shard subqueries in WHERE and rewrite const sql_parser::AstNode* where_expr = up.where_expr; if (where_expr && has_subquery(where_expr) && remote_executor_) { @@ -1220,18 +1275,13 @@ class DistributedPlanner { PlanNode* distribute_delete(PlanNode* plan) { const auto& dp = plan->delete_plan; const TableInfo* table = dp.table; - if (!table || !shards_.has_table(table->table_name)) return plan; - // Multi-table DELETE: emit full SQL from AST, route to primary table's backend if (dp.original_ast) { - sql_parser::StringRef sql = qb_.build_delete_from_ast(dp.original_ast); - if (!shards_.is_sharded(table->table_name)) { - return make_remote_scan(shards_.get_backend(table->table_name), sql, table); - } - const auto& shard_list = shards_.get_shards(table->table_name); - return scatter_dml_to_shards(table, shard_list, [&]() { return sql; }); + return distribute_multi_table_dml(dp.original_ast, table, false); } + if (!table || !shards_.has_table(table->table_name)) return plan; + // Check for cross-shard subqueries in WHERE and rewrite const sql_parser::AstNode* where_expr = dp.where_expr; if (where_expr && has_subquery(where_expr) && remote_executor_) { @@ -1318,7 +1368,68 @@ class DistributedPlanner { return false; } - // Scatter DML SQL to all shards, combining results via UNION ALL + void collect_ast_table_names(const sql_parser::AstNode* n, + std::vector& out) const { + if (!n) return; + if (n->type == sql_parser::NodeType::NODE_TABLE_REF && n->first_child) { + const sql_parser::AstNode* name = n->first_child; + if (name->type == sql_parser::NodeType::NODE_IDENTIFIER) { + out.push_back(name->value()); + } else if (name->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) { + const sql_parser::AstNode* schema = name->first_child; + const sql_parser::AstNode* table = schema ? schema->next_sibling : nullptr; + if (table) out.push_back(table->value()); + else if (schema) out.push_back(schema->value()); + } + } + for (const sql_parser::AstNode* c = n->first_child; c; c = c->next_sibling) { + collect_ast_table_names(c, out); + } + } + + PlanNode* distribute_multi_table_dml(const sql_parser::AstNode* ast, + const TableInfo* primary, + bool is_update) { + std::vector names; + collect_ast_table_names(ast, names); + const char* backend = nullptr; + bool saw_mapped = false; + for (sql_parser::StringRef name : names) { + if (!shards_.has_table(name)) continue; + saw_mapped = true; + if (shards_.is_sharded(name)) { + return fail_dml(is_update + ? "multi-table UPDATE is not supported on sharded tables" + : "multi-table DELETE is not supported on sharded tables"); + } + const char* b = shards_.get_backend(name); + if (backend && b && std::strcmp(backend, b) != 0) { + return fail_dml(is_update + ? "multi-table UPDATE spans multiple backends" + : "multi-table DELETE spans multiple backends"); + } + if (b) backend = b; + } + if (!backend && primary && shards_.has_table(primary->table_name)) { + if (shards_.is_sharded(primary->table_name)) { + return fail_dml(is_update + ? "multi-table UPDATE is not supported on sharded tables" + : "multi-table DELETE is not supported on sharded tables"); + } + backend = shards_.get_backend(primary->table_name); + saw_mapped = true; + } + if (!backend || !saw_mapped) { + return fail_dml(is_update + ? "multi-table UPDATE is not supported on sharded tables" + : "multi-table DELETE is not supported on sharded tables"); + } + sql_parser::StringRef sql = is_update + ? qb_.build_update_from_ast(ast) + : qb_.build_delete_from_ast(ast); + return make_remote_scan(backend, sql, primary); + } + PlanNode* scatter_dml_to_shards(const TableInfo* table, const std::vector& shard_list, std::function build_sql) { diff --git a/include/sql_engine/distributed_txn.h b/include/sql_engine/distributed_txn.h index aae1e82..9236b68 100644 --- a/include/sql_engine/distributed_txn.h +++ b/include/sql_engine/distributed_txn.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()) { diff --git a/include/sql_engine/plan_executor.h b/include/sql_engine/plan_executor.h index c92466b..b6395b1 100644 --- a/include/sql_engine/plan_executor.h +++ b/include/sql_engine/plan_executor.h @@ -57,6 +57,7 @@ #include #include #include +#include namespace sql_engine { @@ -847,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( + 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(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( + agg_node->merge_aggregate.output_expr_count - group_count); + } + } else { + return nullptr; + } + + uint16_t n = static_cast(group_count + agg_count); + if (n == 0) return nullptr; + auto* cols = static_cast(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(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(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 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>( - 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)); diff --git a/include/sql_engine/remote_query_builder.h b/include/sql_engine/remote_query_builder.h index 5610618..ae99ad9 100644 --- a/include/sql_engine/remote_query_builder.h +++ b/include/sql_engine/remote_query_builder.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, diff --git a/include/sql_engine/session.h b/include/sql_engine/session.h index 4013861..a2dd90d 100644 --- a/include/sql_engine/session.h +++ b/include/sql_engine/session.h @@ -13,6 +13,7 @@ #include "sql_engine/result_set.h" #include "sql_engine/dml_result.h" #include "sql_engine/mutable_data_source.h" +#include "sql_engine/remote_executor.h" #include "sql_parser/parser.h" #include "sql_parser/common.h" @@ -25,6 +26,44 @@ namespace sql_engine { +class TxnRoutingExecutor : public RemoteExecutor { +public: + void bind(RemoteExecutor* inner, TransactionManager* txn) { + inner_ = inner; + txn_ = txn; + } + + ResultSet execute(const char* backend_name, sql_parser::StringRef sql) override { + if (txn_ && txn_->in_transaction() && txn_->is_distributed() && + txn_->route_query_supported()) { + return txn_->route_query(backend_name, sql); + } + return inner_ ? inner_->execute(backend_name, sql) : ResultSet{}; + } + + DmlResult execute_dml(const char* backend_name, sql_parser::StringRef sql) override { + if (txn_ && txn_->in_transaction() && txn_->is_distributed()) { + return txn_->route_dml(backend_name, sql); + } + if (inner_) return inner_->execute_dml(backend_name, sql); + DmlResult r; + r.error_message = "no remote executor"; + return r; + } + + bool allows_unpinned_distributed_2pc() const override { + return inner_ && inner_->allows_unpinned_distributed_2pc(); + } + + std::unique_ptr checkout_session(const char* backend_name) override { + return inner_ ? inner_->checkout_session(backend_name) : nullptr; + } + +private: + RemoteExecutor* inner_ = nullptr; + TransactionManager* txn_ = nullptr; +}; + // Session is the high-level API that ties together parsing, planning, // optimization, execution, and transaction management. // @@ -306,6 +345,7 @@ class Session { FunctionRegistry functions_; Optimizer optimizer_; RemoteExecutor* remote_executor_ = nullptr; + TxnRoutingExecutor routing_exec_; const ShardMap* shard_map_ = nullptr; bool parallel_open_enabled_ = false; std::unordered_map sources_; @@ -377,8 +417,10 @@ class Session { executor.add_data_source(kv.first.c_str(), kv.second); for (auto& kv : mutable_sources_) executor.add_mutable_data_source(kv.first.c_str(), kv.second); - if (remote_executor_) - executor.set_remote_executor(remote_executor_); + if (remote_executor_) { + routing_exec_.bind(remote_executor_, &txn_mgr_); + executor.set_remote_executor(&routing_exec_); + } if (parallel_open_enabled_) { executor.set_parallel_open(true); if (pool_) diff --git a/include/sql_engine/shard_map.h b/include/sql_engine/shard_map.h index bff566b..0404c99 100644 --- a/include/sql_engine/shard_map.h +++ b/include/sql_engine/shard_map.h @@ -155,6 +155,34 @@ class ShardMap { return 0; } + bool same_routing(sql_parser::StringRef a, sql_parser::StringRef b) const { + const TableShardConfig* ca = lookup(a); + const TableShardConfig* cb = lookup(b); + if (!ca || !cb) return false; + if (ca->strategy != cb->strategy) return false; + if (ca->shards.size() != cb->shards.size() || ca->shards.empty()) return false; + for (size_t i = 0; i < ca->shards.size(); ++i) { + if (ca->shards[i].backend_name != cb->shards[i].backend_name) return false; + } + if (ca->strategy == RoutingStrategy::RANGE) { + if (ca->ranges.size() != cb->ranges.size()) return false; + for (size_t i = 0; i < ca->ranges.size(); ++i) { + if (ca->ranges[i].upper_inclusive != cb->ranges[i].upper_inclusive || + ca->ranges[i].shard_index != cb->ranges[i].shard_index) return false; + } + } + if (ca->strategy == RoutingStrategy::LIST) { + if (ca->list.size() != cb->list.size()) return false; + for (size_t i = 0; i < ca->list.size(); ++i) { + if (ca->list[i].is_int != cb->list[i].is_int || + ca->list[i].int_val != cb->list[i].int_val || + ca->list[i].str_val != cb->list[i].str_val || + ca->list[i].shard_index != cb->list[i].shard_index) return false; + } + } + return true; + } + // Get the single backend for an unsharded table. const char* get_backend(sql_parser::StringRef table_name) const { const TableShardConfig* cfg = lookup(table_name); diff --git a/include/sql_engine/transaction_manager.h b/include/sql_engine/transaction_manager.h index 7dd0605..67e2f3f 100644 --- a/include/sql_engine/transaction_manager.h +++ b/include/sql_engine/transaction_manager.h @@ -2,6 +2,7 @@ #define SQL_ENGINE_TRANSACTION_MANAGER_H #include "sql_engine/dml_result.h" +#include "sql_engine/result_set.h" #include "sql_parser/common.h" namespace sql_engine { @@ -40,6 +41,13 @@ class TransactionManager { r.error_message = "route_dml not supported by this transaction manager"; return r; } + + virtual ResultSet route_query(const char* /*backend_name*/, + sql_parser::StringRef /*sql*/) { + return {}; + } + + virtual bool route_query_supported() const { return false; } }; } // namespace sql_engine diff --git a/tests/test_distributed_dml.cpp b/tests/test_distributed_dml.cpp index fbc9805..d9d32d6 100644 --- a/tests/test_distributed_dml.cpp +++ b/tests/test_distributed_dml.cpp @@ -770,3 +770,17 @@ TEST_F(DistributedDmlTest, InsertThenPointSelectList) { EXPECT_EQ(execute_distributed_select("SELECT name FROM users WHERE id = 7").row_count(), 1u); EXPECT_EQ(execute_distributed_select("SELECT name FROM users WHERE id = 20").row_count(), 1u); } + +TEST_F(DistributedDmlTest, MultiTableUpdateOnShardedFails) { + auto result = execute_distributed_dml( + "UPDATE users u JOIN orders o ON u.id = o.user_id SET u.age = 30"); + EXPECT_FALSE(result.success); + EXPECT_NE(result.error_message.find("sharded"), std::string::npos); +} + +TEST_F(DistributedDmlTest, MultiTableDeleteOnShardedFails) { + auto result = execute_distributed_dml( + "DELETE u FROM users u JOIN orders o ON u.id = o.user_id"); + EXPECT_FALSE(result.success); + EXPECT_NE(result.error_message.find("sharded"), std::string::npos); +} diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index 7ed2607..8a51eaa 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -1061,3 +1061,41 @@ TEST_F(DistributedPlannerTest, RemoteSqlLenIsNotTruncated) { EXPECT_EQ(rs->remote_scan.remote_sql_len, remote.size()); } } + +TEST_F(DistributedPlannerTest, HavingCountCorrectness) { + const char* sql = "SELECT dept, COUNT(*) FROM users GROUP BY dept HAVING COUNT(*) > 4"; + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_EQ(local_rs.row_count(), 2u); + EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs)); +} + +TEST_F(DistributedPlannerTest, ColocatedJoinPushedToShards) { + shard_map.add_table(TableShardConfig{ + "orders", "user_id", + {ShardInfo{"shard_1"}, ShardInfo{"shard_2"}, ShardInfo{"shard_3"}} + }); + + Parser parser; + const char* sql = "SELECT * FROM users JOIN orders ON users.id = orders.user_id"; + auto pr = parser.parse(sql, std::strlen(sql)); + ASSERT_EQ(pr.status, ParseResult::OK); + + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + ASSERT_NE(plan, nullptr); + + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector joins, remotes; + find_nodes(dist, PlanNodeType::JOIN, joins); + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + EXPECT_TRUE(joins.empty()) << "co-located join should not stay local"; + ASSERT_FALSE(remotes.empty()); + for (auto* rs : remotes) { + std::string remote(rs->remote_scan.remote_sql, rs->remote_scan.remote_sql_len); + EXPECT_NE(remote.find("JOIN"), std::string::npos) << remote; + } +} diff --git a/tests/test_plan_executor.cpp b/tests/test_plan_executor.cpp index 06316d9..cb10fdf 100644 --- a/tests/test_plan_executor.cpp +++ b/tests/test_plan_executor.cpp @@ -177,6 +177,13 @@ TEST_F(PlanExecutorTest, CountDistinctDept) { EXPECT_EQ(rs.rows[0].get(0).int_val, 2); } +TEST_F(PlanExecutorTest, HavingCountFilter) { + auto rs = run_query("SELECT dept, COUNT(*) FROM users GROUP BY dept HAVING COUNT(*) > 2"); + ASSERT_EQ(rs.row_count(), 1u); + EXPECT_EQ(std::string(rs.rows[0].get(0).str_val.ptr, rs.rows[0].get(0).str_val.len), + "Engineering"); +} + // SELECT name FROM users WHERE name LIKE 'A%' → LIKE filter TEST_F(PlanExecutorTest, SelectWithLike) { auto rs = run_query("SELECT name FROM users WHERE name LIKE 'A%'"); diff --git a/tools/mysql_server.cpp b/tools/mysql_server.cpp index bcf968d..3c82c1f 100644 --- a/tools/mysql_server.cpp +++ b/tools/mysql_server.cpp @@ -49,6 +49,7 @@ #include "sql_engine/in_memory_catalog.h" #include "sql_engine/data_source.h" #include "sql_engine/local_txn.h" +#include "sql_engine/distributed_txn.h" #include "sql_engine/multi_remote_executor.h" #include "sql_engine/thread_safe_executor.h" #include "sql_engine/shard_map.h" @@ -582,14 +583,13 @@ static void handle_connection(int client_fd, uint32_t conn_id, const ServerConte } // Set up per-connection session - Arena txn_arena{65536, 1048576}; - LocalTransactionManager txn_mgr(txn_arena); - ThreadSafeMultiRemoteExecutor remote_exec; for (auto& bc : ctx.backends) { remote_exec.add_backend(bc); } + DistributedTransactionManager txn_mgr( + remote_exec, DistributedTransactionManager::BackendDialect::MYSQL); Session session(ctx.catalog, txn_mgr); session.set_remote_executor(&remote_exec); session.set_parallel_open(true); // thread-safe executor enables parallel shard I/O diff --git a/tools/sqlengine.cpp b/tools/sqlengine.cpp index 5522162..ea6adc1 100644 --- a/tools/sqlengine.cpp +++ b/tools/sqlengine.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include "sql_parser/parser.h" #include "sql_parser/common.h" @@ -25,6 +26,8 @@ #include "sql_engine/in_memory_catalog.h" #include "sql_engine/data_source.h" #include "sql_engine/local_txn.h" +#include "sql_engine/distributed_txn.h" +#include "sql_engine/durable_txn_log.h" #include "sql_engine/multi_remote_executor.h" #include "sql_engine/thread_safe_executor.h" #include "sql_engine/shard_map.h" @@ -226,6 +229,7 @@ static void print_usage(const char* prog) { << "Options:\n" << " --backend URL Add a backend (mysql://... or pgsql://...)\n" << " --shard SPEC Add shard config (table:key:shard1,shard2)\n" + << " --txn-log PATH Durable 2PC WAL (backend mode only)\n" << " --help Show this help\n" << "\n" << "In-memory mode (no --backend): evaluates expressions locally.\n" @@ -239,6 +243,7 @@ static void print_usage(const char* prog) { int main(int argc, char* argv[]) { std::vector backends; std::vector shards; + std::string txn_log_path; // Parse command-line args for (int i = 1; i < argc; ++i) { @@ -254,6 +259,9 @@ int main(int argc, char* argv[]) { return 1; } backends.push_back(std::move(pb.config)); + } else if (arg == "--txn-log" && i + 1 < argc) { + ++i; + txn_log_path = argv[i]; } else if (arg == "--shard" && i + 1 < argc) { ++i; auto ps = parse_shard_spec(argv[i]); @@ -274,7 +282,10 @@ int main(int argc, char* argv[]) { // Set up arena for transaction manager Arena txn_arena{65536, 1048576}; - LocalTransactionManager txn_mgr(txn_arena); + LocalTransactionManager local_txn(txn_arena); + std::unique_ptr dtxn; + DurableTransactionLog txn_log; + TransactionManager* txn_mgr = &local_txn; // Set up shard map ShardMap shard_map; @@ -344,8 +355,20 @@ int main(int argc, char* argv[]) { } } - // Create session - Session session(catalog, txn_mgr); + if (remote_exec) { + dtxn.reset(new DistributedTransactionManager( + *remote_exec, DistributedTransactionManager::BackendDialect::MYSQL)); + if (!txn_log_path.empty()) { + if (!txn_log.open(txn_log_path)) { + std::cerr << "Error: cannot open txn log " << txn_log_path << std::endl; + return 1; + } + dtxn->set_durable_log(&txn_log); + } + txn_mgr = dtxn.get(); + } + + Session session(catalog, *txn_mgr); if (remote_exec) { session.set_remote_executor(remote_exec); session.set_parallel_open(true); // thread-safe executor enables parallel shard I/O