From 5da29e23a7fa9671b43c06272cf04bc3ea617bde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 02:42:34 +0700 Subject: [PATCH 1/3] 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. --- include/sql_engine/distributed_planner.h | 120 +++++++++++++++++--- include/sql_engine/operators/aggregate_op.h | 23 +++- include/sql_parser/common.h | 4 + include/sql_parser/emitter.h | 1 + include/sql_parser/expression_parser.h | 4 + scripts/test_sqlengine.sh | 7 ++ tests/test_distributed_planner.cpp | 115 +++++++++++++++++++ tests/test_distributed_real.cpp | 93 +++++++++++++++ tests/test_expression.cpp | 9 ++ tests/test_plan_executor.cpp | 20 ++++ 10 files changed, 372 insertions(+), 24 deletions(-) diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 161a97a..8cc0940 100644 --- a/include/sql_engine/distributed_planner.h +++ b/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& projs, std::vector& 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(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(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; diff --git a/include/sql_engine/operators/aggregate_op.h b/include/sql_engine/operators/aggregate_op.h index f2b3a79..20039e0 100644 --- a/include/sql_engine/operators/aggregate_op.h +++ b/include/sql_engine/operators/aggregate_op.h @@ -12,6 +12,8 @@ #include #include #include +#include +#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 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,6 +207,14 @@ 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& resolver) { switch (state.type) { @@ -210,10 +222,9 @@ class AggregateOperator : public Operator { 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(arg, resolver, functions_, arena_); - if (!v.is_null()) state.count++; + if (!v.is_null() && note_distinct(state, expr, v)) state.count++; } break; } @@ -221,7 +232,7 @@ class AggregateOperator : public Operator { case AggType::AVG: { const sql_parser::AstNode* arg = expr->first_child; Value v = evaluate_expression(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(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(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; diff --git a/include/sql_parser/common.h b/include/sql_parser/common.h index 7cc651d..753805f 100644 --- a/include/sql_parser/common.h +++ b/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 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 { diff --git a/include/sql_parser/emitter.h b/include/sql_parser/emitter.h index 39f6596..41049f0 100644 --- a/include/sql_parser/emitter.h +++ b/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(", "); diff --git a/include/sql_parser/expression_parser.h b/include/sql_parser/expression_parser.h index 51b15c8..050d93e 100644 --- a/include/sql_parser/expression_parser.h +++ b/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; diff --git a/scripts/test_sqlengine.sh b/scripts/test_sqlengine.sh index 5ee9ddc..0c0e37c 100755 --- a/scripts/test_sqlengine.sh +++ b/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 } # ---------------------------------------------------------------------- diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index 0fdb483..2daa5c7 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -849,3 +849,118 @@ TEST_F(DistributedPlannerTest, ShardMap_IndexForInt) { EXPECT_EQ(idx1, idx2); EXPECT_LT(idx1, 3u); } + +TEST_F(DistributedPlannerTest, CountDistinctIsNotTwoPhaseMerge) { + Parser parser; + const char* sql = "SELECT COUNT(DISTINCT dept) FROM users"; + auto pr = parser.parse(sql, std::strlen(sql)); + ASSERT_EQ(pr.status, ParseResult::OK); + + const AstNode* items = nullptr; + for (const AstNode* c = pr.ast->first_child; c; c = c->next_sibling) { + if (c->type == NodeType::NODE_SELECT_ITEM_LIST) items = c; + } + ASSERT_NE(items, nullptr); + ASSERT_NE(items->first_child, nullptr); + const AstNode* count_expr = items->first_child->first_child; + ASSERT_NE(count_expr, nullptr); + EXPECT_EQ(count_expr->type, NodeType::NODE_FUNCTION_CALL); + EXPECT_NE(static_cast(count_expr->flags & FLAG_FUNC_DISTINCT), 0u); + + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + ASSERT_NE(plan, nullptr); + if (plan->type == PlanNodeType::PROJECT && plan->project.count > 0) { + EXPECT_NE(static_cast(plan->project.exprs[0]->flags & FLAG_FUNC_DISTINCT), 0u); + } + + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector merges; + find_nodes(dist, PlanNodeType::MERGE_AGGREGATE, merges); + EXPECT_TRUE(merges.empty()) << "COUNT(DISTINCT) must not use SUM_OF_COUNTS merge"; +} + +TEST_F(DistributedPlannerTest, CountDistinctCorrectness) { + Parser parser; + const char* sql = "SELECT COUNT(DISTINCT dept) FROM users"; + 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); + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + ASSERT_EQ(dist->type, PlanNodeType::AGGREGATE); + ASSERT_EQ(dist->aggregate.agg_count, 1u); + EXPECT_NE(static_cast(dist->aggregate.agg_exprs[0]->flags & FLAG_FUNC_DISTINCT), 0u); + + auto local_rs = execute_local("SELECT COUNT(DISTINCT dept) FROM users"); + auto dist_rs = execute_distributed("SELECT COUNT(DISTINCT dept) FROM users"); + EXPECT_EQ(local_rs.row_count(), 1u); + EXPECT_EQ(dist_rs.row_count(), 1u); + EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs)); + ASSERT_GE(dist_rs.row_count(), 1u); + EXPECT_EQ(dist_rs.rows[0].get(0).int_val, 3); +} + +TEST_F(DistributedPlannerTest, GroupByHavingKeepsFilter) { + Parser parser; + const char* sql = "SELECT dept, COUNT(*) FROM users GROUP BY dept HAVING COUNT(*) > 4 ORDER BY dept"; + 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 filters, sorts, merges; + find_nodes(dist, PlanNodeType::FILTER, filters); + find_nodes(dist, PlanNodeType::SORT, sorts); + find_nodes(dist, PlanNodeType::MERGE_AGGREGATE, merges); + EXPECT_FALSE(filters.empty()) << "HAVING filter must be kept"; + EXPECT_FALSE(sorts.empty()) << "ORDER BY must be kept"; + EXPECT_FALSE(merges.empty()); +} + +TEST_F(DistributedPlannerTest, GroupByOrderByCorrectness) { + const char* sql = "SELECT dept, COUNT(*) FROM users GROUP BY dept ORDER BY dept"; + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_EQ(local_rs.row_count(), 3u); + EXPECT_TRUE(compare_results_ordered(local_rs, dist_rs)); +} + +TEST_F(DistributedPlannerTest, WindowNotDroppedByOrderBy) { + Parser parser; + const char* sql = "SELECT name, ROW_NUMBER() OVER (ORDER BY age) AS rn FROM users ORDER BY name"; + 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 windows; + find_nodes(dist, PlanNodeType::WINDOW, windows); + EXPECT_FALSE(windows.empty()) << "WINDOW node must survive ORDER BY distribution"; +} + +TEST_F(DistributedPlannerTest, WindowGatherCorrectness) { + const char* sql = "SELECT name, ROW_NUMBER() OVER (ORDER BY age) AS rn FROM users"; + auto local_rs = execute_local(sql); + auto dist_rs = execute_distributed(sql); + EXPECT_EQ(local_rs.row_count(), 15u); + EXPECT_EQ(dist_rs.row_count(), 15u); + EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs)); +} diff --git a/tests/test_distributed_real.cpp b/tests/test_distributed_real.cpp index 24a0388..ef904f0 100644 --- a/tests/test_distributed_real.cpp +++ b/tests/test_distributed_real.cpp @@ -15,6 +15,8 @@ #include "sql_engine/shard_map.h" #include "sql_engine/in_memory_catalog.h" #include "sql_engine/function_registry.h" +#include "sql_engine/session.h" +#include "sql_engine/local_txn.h" #include "sql_parser/parser.h" #include @@ -272,4 +274,95 @@ TEST_F(MultiExecutorTest, UnknownBackendReturnsEmpty) { EXPECT_TRUE(rs.empty()); } +static bool mysql_port_available(uint16_t port) { + MYSQL* conn = mysql_init(nullptr); + if (!conn) return false; + unsigned int timeout = 2; + mysql_options(conn, MYSQL_OPT_CONNECT_TIMEOUT, &timeout); + bool ok = mysql_real_connect(conn, "127.0.0.1", "root", "test", + "testdb", port, nullptr, 0) != nullptr; + mysql_close(conn); + return ok; +} + +class LiveShardedWriteTest : public ::testing::Test { +protected: + void SetUp() override { + exec_ = std::make_unique(); + BackendConfig s1; + s1.name = "shard1"; + s1.host = "127.0.0.1"; + s1.port = 13306; + s1.user = "root"; + s1.password = "test"; + s1.database = "testdb"; + s1.dialect = Dialect::MySQL; + BackendConfig s2 = s1; + s2.name = "shard2"; + s2.port = 13307; + exec_->add_backend(s1); + exec_->add_backend(s2); + + catalog_.add_table("", "shard_write_t", { + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(64), true}, + }); + + TableShardConfig cfg; + cfg.table_name = "shard_write_t"; + cfg.shard_key = "id"; + cfg.shards = {{"shard1"}, {"shard2"}}; + cfg.strategy = RoutingStrategy::RANGE; + cfg.ranges = {{5, 0}, {100000, 1}}; + shard_map_.add_table(cfg); + } + + void TearDown() override { + if (exec_) exec_->disconnect_all(); + } + + std::unique_ptr exec_; + InMemoryCatalog catalog_; + ShardMap shard_map_; +}; + +TEST_F(LiveShardedWriteTest, InsertThenPointSelect) { + if (!mysql_port_available(13306) || !mysql_port_available(13307)) { + GTEST_SKIP() << "Need MySQL on 13306 and 13307 (start_sharding_demo.sh)"; + } + + const char* ddl = + "CREATE TABLE IF NOT EXISTS shard_write_t (id INT PRIMARY KEY, name VARCHAR(64))"; + StringRef ddl_ref{ddl, static_cast(std::strlen(ddl))}; + exec_->execute_dml("shard1", ddl_ref); + exec_->execute_dml("shard2", ddl_ref); + const char* wipe = "DELETE FROM shard_write_t"; + StringRef wipe_ref{wipe, static_cast(std::strlen(wipe))}; + exec_->execute_dml("shard1", wipe_ref); + exec_->execute_dml("shard2", wipe_ref); + + Arena txn_arena{65536, 1048576}; + LocalTransactionManager txn(txn_arena); + Session session(catalog_, txn); + session.set_remote_executor(exec_.get()); + session.set_shard_map(&shard_map_); + + auto ins1 = session.execute_statement( + "INSERT INTO shard_write_t (id, name) VALUES (3, 'low')"); + auto ins2 = session.execute_statement( + "INSERT INTO shard_write_t (id, name) VALUES (9, 'high')"); + EXPECT_TRUE(ins1.success) << ins1.error_message; + EXPECT_TRUE(ins2.success) << ins2.error_message; + + auto low = session.execute_query("SELECT name FROM shard_write_t WHERE id = 3"); + auto high = session.execute_query("SELECT name FROM shard_write_t WHERE id = 9"); + ASSERT_EQ(low.row_count(), 1u); + ASSERT_EQ(high.row_count(), 1u); + EXPECT_EQ(std::string(low.rows[0].get(0).str_val.ptr, low.rows[0].get(0).str_val.len), "low"); + EXPECT_EQ(std::string(high.rows[0].get(0).str_val.ptr, high.rows[0].get(0).str_val.len), "high"); + + session.execute_statement("DELETE FROM shard_write_t WHERE id = 3"); + session.execute_statement("DELETE FROM shard_write_t WHERE id = 9"); +} + } // namespace diff --git a/tests/test_expression.cpp b/tests/test_expression.cpp index ef98bbb..485c4a3 100644 --- a/tests/test_expression.cpp +++ b/tests/test_expression.cpp @@ -301,6 +301,15 @@ TEST_F(ExpressionTest, FunctionCall) { EXPECT_EQ(node->first_child->type, NodeType::NODE_ASTERISK); } +TEST_F(ExpressionTest, FunctionCallDistinct) { + AstNode* node = parse_expr("COUNT(DISTINCT dept)"); + ASSERT_NE(node, nullptr); + EXPECT_EQ(node->type, NodeType::NODE_FUNCTION_CALL); + EXPECT_NE(static_cast(node->flags & FLAG_FUNC_DISTINCT), 0u); + ASSERT_NE(node->first_child, nullptr); + EXPECT_EQ(node->first_child->type, NodeType::NODE_COLUMN_REF); +} + TEST_F(ExpressionTest, FunctionCallMultiArg) { AstNode* node = parse_expr("COALESCE(a, b, 0)"); ASSERT_NE(node, nullptr); diff --git a/tests/test_plan_executor.cpp b/tests/test_plan_executor.cpp index 2ca5e0f..32e5206 100644 --- a/tests/test_plan_executor.cpp +++ b/tests/test_plan_executor.cpp @@ -147,6 +147,26 @@ TEST_F(PlanExecutorTest, SelectDistinctDept) { EXPECT_EQ(depts.size(), 2u); } +TEST_F(PlanExecutorTest, CountDistinctDept) { + parser.reset(); + const char* sql = "SELECT COUNT(DISTINCT dept) FROM users"; + auto r = parser.parse(sql, std::strlen(sql)); + ASSERT_EQ(r.status, ParseResult::OK); + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(r.ast); + ASSERT_NE(plan, nullptr); + ASSERT_EQ(plan->type, PlanNodeType::PROJECT); + ASSERT_GE(plan->project.count, 1u); + EXPECT_EQ(plan->project.exprs[0]->type, NodeType::NODE_FUNCTION_CALL); + EXPECT_NE(static_cast(plan->project.exprs[0]->flags & FLAG_FUNC_DISTINCT), 0u); + + PlanExecutor executor(functions, catalog, parser.arena()); + executor.add_data_source("users", users_source); + auto rs = executor.execute(plan); + ASSERT_EQ(rs.row_count(), 1u); + EXPECT_EQ(rs.rows[0].get(0).int_val, 2); +} + // 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%'"); 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 2/3] 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 3/3] 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