From 9e264dae834ee4bbb59ba4153e0d01c1a396121d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 03:38:09 +0700 Subject: [PATCH 1/3] feat: prune OR/range, pool PostgreSQL, 2PC stress, honest docs Prune shard_key OR-of-equalities and RANGE inequalities/BETWEEN. Placeholders still scatter. ThreadSafeMultiRemoteExecutor pools PostgreSQL as well as MySQL. engine_stress_test uses DistributedTransactionManager. README matches current Session/ShardMap and tool 2PC behavior. --- AGENTS.md | 2 +- README.md | 30 +-- include/sql_engine/distributed_planner.h | 80 ++++++++ include/sql_engine/pg_connection_pool.h | 108 ++++++++++ include/sql_engine/shard_map.h | 27 +++ include/sql_engine/thread_safe_executor.h | 233 +++++++++++++++++++++- tests/test_distributed_planner.cpp | 66 ++++++ tests/test_pgsql_executor.cpp | 10 + tests/test_shard_map.cpp | 18 ++ tools/engine_stress_test.cpp | 6 +- 10 files changed, 558 insertions(+), 22 deletions(-) create mode 100644 include/sql_engine/pg_connection_pool.h diff --git a/AGENTS.md b/AGENTS.md index 3d1e56a..be0e703 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,7 +7,7 @@ Trust the `Makefile` over prose. Extension recipes live in `CLAUDE.md`. `docs/su - Parser: header-only templates in `include/sql_parser/` except `src/sql_parser/{arena,parser}.cpp` - Engine: headers in `include/sql_engine/` (`operators/`, `functions/`, `rules/`); compiled files are the explicit `ENGINE_SRCS` list - High-level API: `Session` (`include/sql_engine/session.h`) — parse → plan → optimize → distribute → execute -- Production remote path: `ThreadSafeMultiRemoteExecutor`, not the single-connection executors +- Production remote path: `ThreadSafeMultiRemoteExecutor` (pooled MySQL **and** PostgreSQL), not the single-connection executors - All shard routing (SELECT prune and DML) goes through `ShardMap`. Do not add a private hash in the planner. - Backend URL / shard-spec parsing: `tool_config_parser` — do not add another copy in tools - Do not edit `third_party/` diff --git a/README.md b/README.md index dfab227..4dfc03f 100644 --- a/README.md +++ b/README.md @@ -93,7 +93,7 @@ echo "SELECT 1 + 2, UPPER('hello'), COALESCE(NULL, 42)" | ./sqlengine # Against a MySQL backend ./sqlengine --backend "mysql://root:pass@127.0.0.1:3306/mydb?name=primary" -# Sharded across two backends +# Sharded across two backends (2PC is on; optional --txn-log PATH) ./sqlengine \ --backend "mysql://root:pass@host1:3306/db?name=shard1" \ --backend "mysql://root:pass@host2:3306/db?name=shard2" \ @@ -171,20 +171,20 @@ ResultSet rs = executor.execute(plan); #include "sql_engine/session.h" #include "sql_engine/thread_safe_executor.h" #include "sql_engine/shard_map.h" -#include "sql_engine/local_txn.h" +#include "sql_engine/distributed_txn.h" -// Backends (connection-pooled, thread-safe) +// Backends (connection-pooled, thread-safe; MySQL or PostgreSQL) ThreadSafeMultiRemoteExecutor executor; executor.add_backend({.name = "shard1", .host = "h1", .port = 3306, ...}); executor.add_backend({.name = "shard2", .host = "h2", .port = 3306, ...}); // Sharding policy: "users" is sharded on "id" across shard1, shard2 ShardMap shards; -shards.add_sharded_table("users", "id", {"shard1", "shard2"}); +shards.add_table({"users", "id", {{"shard1"}, {"shard2"}}}); -// Catalog, transactions, session +// Catalog + 2PC (required for atomic multi-shard DML) InMemoryCatalog catalog; /* ... add_table(...) ... */ -LocalTransactionManager txn; +DistributedTransactionManager txn(executor); Session session(catalog, txn); session.set_remote_executor(&executor); session.set_shard_map(&shards); @@ -354,11 +354,11 @@ auto report = recovery.recover(); ### Distributed execution -- **Shard routing** — shard-key lookups go to one backend; scatter queries go to all -- **Distributed aggregation** — per-shard partial aggregates + coordinator merge (COUNT+SUM+MIN+MAX + AVG from SUM/COUNT) -- **Distributed sort** — per-shard sort + coordinator merge -- **Cross-shard joins** — hash-join coordinator; materialized subquery cache -- **Cross-shard DML** — scatter INSERT/UPDATE/DELETE when no shard key; single-shard when key present +- **Shard routing** — equality / `IN` / `OR` of equalities prune via `ShardMap`; RANGE also prunes `<`/`>`/`BETWEEN`. Placeholders scatter. +- **Distributed aggregation** — per-shard partial aggregates + coordinator merge (`COUNT`/`SUM`/`MIN`/`MAX`/`AVG`). `COUNT(DISTINCT)` gathers then aggregates locally. +- **Distributed sort** — per-shard sort + coordinator merge when keys are table columns +- **Joins** — co-located same-key joins push down; otherwise gather both sides and join locally +- **Cross-shard DML** — routed by `ShardMap`; missing/non-literal shard key and multi-table DML on shards fail closed - **Cross-shard INSERT ... SELECT** — source materialized, rows routed by destination shard key ### Transactions @@ -372,10 +372,10 @@ auto report = recovery.recover(); ### Backends & connectivity -- **MySQL** — libmysqlclient with pooled and single-connection paths, UTF-8, configurable timeouts -- **PostgreSQL** — libpq with statement_timeout, UTC-normalized TIMESTAMPTZ handling +- **MySQL** — libmysqlclient with pooled (`ThreadSafeMultiRemoteExecutor`) and single-connection paths +- **PostgreSQL** — libpq pooled on the same executor, plus a single-connection path; `statement_timeout` and UTC TIMESTAMPTZ - **SSL/TLS** — `ssl_mode`, `ssl_ca`, `ssl_cert`, `ssl_key` configurable per backend for both dialects -- **Connection pool** — thread-safe with health checks, reconnection, RAII `ConnectionGuard` +- **Connection pool** — thread-safe per dialect, RAII checkout, poison-on-error - **MySQL wire-protocol server** — `mysql_server` speaks the MySQL protocol; backends are ParserSQL engines ### Thread-safety @@ -388,7 +388,7 @@ auto report = recovery.recover(); | Tool | Build | Purpose | |---|---|---| -| `sqlengine` | `make build-sqlengine` | Interactive SQL CLI; stdin, one-shot, or REPL; optional backends and sharding | +| `sqlengine` | `make build-sqlengine` | Interactive SQL CLI; 2PC when `--backend` is set; optional `--txn-log` | | `mysql_server` | `make mysql-server` | MySQL wire-protocol server fronted by the ParserSQL engine | | `corpus_test` | `make build-corpus-test` | Read SQL from stdin/files, parse each, report OK/PARTIAL/ERROR | | `engine_stress_test` | `make engine-stress` | Direct-API engine stress test | diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 77976e5..251109e 100644 --- a/include/sql_engine/distributed_planner.h +++ b/include/sql_engine/distributed_planner.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -401,6 +402,62 @@ class DistributedPlanner { } } } + if (is_compare_op(op) && + shards_.routing_strategy(table_name) == RoutingStrategy::RANGE) { + const sql_parser::AstNode* left_node = expr->first_child; + const sql_parser::AstNode* right_node = left_node ? left_node->next_sibling : nullptr; + if (left_node && right_node) { + const sql_parser::AstNode* col = nullptr; + const sql_parser::AstNode* lit = nullptr; + bool key_on_left = false; + if (is_shard_key_ref(left_node, shard_key) && is_literal(right_node)) { + col = left_node; lit = right_node; key_on_left = true; + } else if (is_shard_key_ref(right_node, shard_key) && is_literal(left_node)) { + col = right_node; lit = left_node; key_on_left = false; + } + if (col && lit) { + int64_t v = literal_to_int(lit); + int64_t lo = INT64_MIN, hi = INT64_MAX; + char c0 = op.ptr[0]; + bool has_eq = op.len == 2 && op.ptr[1] == '='; + if (c0 == '<' && key_on_left) { + hi = has_eq ? v : (v == INT64_MIN ? INT64_MIN : v - 1); + } else if (c0 == '>' && key_on_left) { + lo = has_eq ? v : (v == INT64_MAX ? INT64_MAX : v + 1); + } else if (c0 == '<' && !key_on_left) { + lo = has_eq ? v : (v == INT64_MAX ? INT64_MAX : v + 1); + } else if (c0 == '>' && !key_on_left) { + hi = has_eq ? v : (v == INT64_MIN ? INT64_MIN : v - 1); + } + shards_.collect_int_range_shards(table_name, lo, hi, target_indices); + return; + } + } + } + if (op.len == 2 && + (op.ptr[0] == 'O' || op.ptr[0] == 'o') && + (op.ptr[1] == 'R' || op.ptr[1] == 'r')) { + const sql_parser::AstNode* left_node = expr->first_child; + const sql_parser::AstNode* right_node = left_node ? left_node->next_sibling : nullptr; + std::vector left_targets, right_targets; + extract_shard_targets(left_node, shard_key, table_name, num_shards, left_targets); + extract_shard_targets(right_node, shard_key, table_name, num_shards, right_targets); + if (left_targets.empty() || right_targets.empty()) return; + std::vector seen(num_shards, false); + for (auto i : left_targets) { + if (i < num_shards && !seen[i]) { + seen[i] = true; + target_indices.push_back(i); + } + } + for (auto i : right_targets) { + if (i < num_shards && !seen[i]) { + seen[i] = true; + target_indices.push_back(i); + } + } + return; + } // Recurse into AND branches if (op.len == 3 && (op.ptr[0] == 'A' || op.ptr[0] == 'a') && @@ -446,6 +503,29 @@ class DistributedPlanner { } } } + + if (expr->type == sql_parser::NodeType::NODE_BETWEEN && + shards_.routing_strategy(table_name) == RoutingStrategy::RANGE) { + const sql_parser::AstNode* col = expr->first_child; + const sql_parser::AstNode* lo = col ? col->next_sibling : nullptr; + const sql_parser::AstNode* hi = lo ? lo->next_sibling : nullptr; + if (col && is_shard_key_ref(col, shard_key) && is_literal(lo) && is_literal(hi)) { + shards_.collect_int_range_shards(table_name, + literal_to_int(lo), literal_to_int(hi), + target_indices); + } + } + } + + static bool is_compare_op(sql_parser::StringRef op) { + if (op.len == 1) return op.ptr[0] == '<' || op.ptr[0] == '>'; + if (op.len == 2) return (op.ptr[0] == '<' || op.ptr[0] == '>') && op.ptr[1] == '='; + return false; + } + + static int64_t literal_to_int(const sql_parser::AstNode* lit) { + if (!lit || !lit->value().ptr) return 0; + return std::strtoll(lit->value().ptr, nullptr, 10); } bool is_shard_key_ref(const sql_parser::AstNode* node, sql_parser::StringRef shard_key) const { diff --git a/include/sql_engine/pg_connection_pool.h b/include/sql_engine/pg_connection_pool.h new file mode 100644 index 0000000..22288e5 --- /dev/null +++ b/include/sql_engine/pg_connection_pool.h @@ -0,0 +1,108 @@ +#ifndef SQL_ENGINE_PG_CONNECTION_POOL_H +#define SQL_ENGINE_PG_CONNECTION_POOL_H + +#include "sql_engine/backend_config.h" +#include +#include +#include +#include +#include +#include +#include + +#ifndef SQL_ENGINE_PG_STATEMENT_TIMEOUT_MS +#define SQL_ENGINE_PG_STATEMENT_TIMEOUT_MS 30000 +#endif + +namespace sql_engine { + +class PgConnectionPool { +public: + PgConnectionPool() = default; + + ~PgConnectionPool() { + for (auto& kv : backends_) { + auto& be = *kv.second; + std::lock_guard lk(be.mu); + for (PGconn* c : be.idle) { + if (c) PQfinish(c); + } + } + } + + void add_backend(const BackendConfig& config) { + auto be = std::make_unique(); + be->config = config; + backends_[config.name] = std::move(be); + } + + bool has_backend(const std::string& name) const { + return backends_.find(name) != backends_.end(); + } + + PGconn* checkout(const std::string& backend) { + Backend& be = get_backend(backend); + { + std::lock_guard lk(be.mu); + if (!be.idle.empty()) { + PGconn* c = be.idle.back(); + be.idle.pop_back(); + if (c && PQstatus(c) == CONNECTION_OK) return c; + if (c) PQfinish(c); + } + } + return create_connection(be); + } + + void checkin(const std::string& backend, PGconn* conn) { + if (!conn) return; + Backend& be = get_backend(backend); + std::lock_guard lk(be.mu); + be.idle.push_back(conn); + } + +private: + struct Backend { + BackendConfig config; + std::mutex mu; + std::vector idle; + }; + + std::unordered_map> backends_; + + Backend& get_backend(const std::string& name) { + auto it = backends_.find(name); + if (it == backends_.end()) { + throw std::runtime_error("PgConnectionPool: unknown backend: " + name); + } + return *it->second; + } + + static PGconn* create_connection(Backend& be) { + const BackendConfig& cfg = be.config; + std::string conninfo = "host=" + cfg.host + + " port=" + std::to_string(cfg.port) + + " user=" + cfg.user + + " password=" + cfg.password + + " dbname=" + cfg.database + + " connect_timeout=5" + + " options='-c statement_timeout=" + + std::to_string(SQL_ENGINE_PG_STATEMENT_TIMEOUT_MS) + "'"; + if (!cfg.ssl_mode.empty()) conninfo += " sslmode=" + cfg.ssl_mode; + if (!cfg.ssl_ca.empty()) conninfo += " sslrootcert=" + cfg.ssl_ca; + if (!cfg.ssl_cert.empty()) conninfo += " sslcert=" + cfg.ssl_cert; + if (!cfg.ssl_key.empty()) conninfo += " sslkey=" + cfg.ssl_key; + + PGconn* c = PQconnectdb(conninfo.c_str()); + if (PQstatus(c) != CONNECTION_OK) { + std::string err = PQerrorMessage(c); + PQfinish(c); + throw std::runtime_error("PgConnectionPool connect failed for " + cfg.name + ": " + err); + } + return c; + } +}; + +} // namespace sql_engine + +#endif diff --git a/include/sql_engine/shard_map.h b/include/sql_engine/shard_map.h index 0404c99..cbb6061 100644 --- a/include/sql_engine/shard_map.h +++ b/include/sql_engine/shard_map.h @@ -3,6 +3,7 @@ #include "sql_parser/common.h" #include +#include #include #include #include @@ -155,6 +156,32 @@ class ShardMap { return 0; } + RoutingStrategy routing_strategy(sql_parser::StringRef table_name) const { + const TableShardConfig* cfg = lookup(table_name); + return cfg ? cfg->strategy : RoutingStrategy::HASH; + } + + // RANGE only. Inclusive [lo, hi]. HASH/LIST yield no indices (caller scatters). + void collect_int_range_shards(sql_parser::StringRef table_name, + int64_t lo, int64_t hi, + std::vector& out) const { + const TableShardConfig* cfg = lookup(table_name); + if (!cfg || cfg->strategy != RoutingStrategy::RANGE || cfg->ranges.empty()) + return; + if (lo > hi) return; + size_t n = cfg->shards.size(); + const auto& ranges = cfg->ranges; + for (size_t i = 0; i < ranges.size(); ++i) { + int64_t seg_hi = (i + 1 == ranges.size()) + ? INT64_MAX : ranges[i].upper_inclusive; + int64_t seg_lo = (i == 0) ? INT64_MIN + : (cfg->ranges[i - 1].upper_inclusive == INT64_MAX + ? INT64_MAX : cfg->ranges[i - 1].upper_inclusive + 1); + if (seg_lo <= hi && seg_hi >= lo) + out.push_back(clamp_index(ranges[i].shard_index, n)); + } + } + bool same_routing(sql_parser::StringRef a, sql_parser::StringRef b) const { const TableShardConfig* ca = lookup(a); const TableShardConfig* cb = lookup(b); diff --git a/include/sql_engine/thread_safe_executor.h b/include/sql_engine/thread_safe_executor.h index a6f92d3..a5b045d 100644 --- a/include/sql_engine/thread_safe_executor.h +++ b/include/sql_engine/thread_safe_executor.h @@ -14,6 +14,7 @@ #include "sql_engine/remote_executor.h" #include "sql_engine/remote_session.h" #include "sql_engine/connection_pool.h" +#include "sql_engine/pg_connection_pool.h" #include "sql_engine/backend_config.h" #include "sql_engine/result_set.h" #include "sql_engine/dml_result.h" @@ -23,6 +24,7 @@ #include "sql_parser/common.h" #include +#include #include #include #include @@ -132,6 +134,89 @@ inline ResultSet mysql_result_to_resultset_impl(MYSQL_RES* res) { return rs; } +#ifndef BOOLOID +#define BOOLOID 16 +#define INT2OID 21 +#define INT4OID 23 +#define INT8OID 20 +#define FLOAT4OID 700 +#define FLOAT8OID 701 +#define NUMERICOID 1700 +#define DATEOID 1082 +#define TIMEOID 1083 +#define TIMESTAMPOID 1114 +#define TIMESTAMPTZOID 1184 +#define BYTEAOID 17 +#define JSONOID 114 +#define JSONBOID 3802 +#define OIDOID 26 +#endif + +inline Value pg_field_to_value_impl( + ResultSet& rs, const char* data, int length, Oid type, bool is_null) +{ + if (is_null) return value_null(); + switch (type) { + case BOOLOID: + return value_bool(data[0] == 't' || data[0] == 'T'); + case INT2OID: + case INT4OID: + case INT8OID: + case OIDOID: + return value_int(std::strtoll(data, nullptr, 10)); + case FLOAT4OID: + case FLOAT8OID: + return value_double(std::strtod(data, nullptr)); + case NUMERICOID: { + sql_parser::StringRef s = rs.own_string(data, static_cast(length)); + return value_string(s); + } + case DATEOID: + return value_date(datetime_parse::parse_date(data)); + case TIMESTAMPOID: + return value_datetime(datetime_parse::parse_datetime(data)); + case TIMESTAMPTZOID: + return value_timestamp(datetime_parse::parse_datetime_tz(data)); + case TIMEOID: + return value_time(datetime_parse::parse_time(data)); + case BYTEAOID: { + sql_parser::StringRef s = rs.own_string(data, static_cast(length)); + return value_bytes(s); + } + case JSONOID: + case JSONBOID: { + sql_parser::StringRef s = rs.own_string(data, static_cast(length)); + return value_json(s); + } + default: { + sql_parser::StringRef s = rs.own_string(data, static_cast(length)); + return value_string(s); + } + } +} + +inline ResultSet pg_result_to_resultset_impl(PGresult* res) { + ResultSet rs; + int num_fields = PQnfields(res); + int num_rows = PQntuples(res); + rs.column_count = static_cast(num_fields); + for (int i = 0; i < num_fields; ++i) { + rs.column_names.emplace_back(PQfname(res, i)); + } + for (int r = 0; r < num_rows; ++r) { + Row& row = rs.add_heap_row(rs.column_count); + for (int c = 0; c < num_fields; ++c) { + bool is_null = PQgetisnull(res, r, c) != 0; + Oid oid = PQftype(res, c); + const char* data = PQgetvalue(res, r, c); + int length = PQgetlength(res, r, c); + row.set(static_cast(c), + pg_field_to_value_impl(rs, data, length, oid, is_null)); + } + } + return rs; +} + } // namespace detail // ---------------------------------------------------------------------------- @@ -280,19 +365,122 @@ class PooledMySQLSession : public RemoteSession { bool poisoned_ = false; }; +class PgConnectionGuard { +public: + PgConnectionGuard(PgConnectionPool& pool, std::string name) + : pool_(pool), name_(std::move(name)), conn_(pool_.checkout(name_)) {} + ~PgConnectionGuard() { + if (!conn_) return; + if (poisoned_) PQfinish(conn_); + else pool_.checkin(name_, conn_); + } + PgConnectionGuard(const PgConnectionGuard&) = delete; + PgConnectionGuard& operator=(const PgConnectionGuard&) = delete; + PGconn* get() const { return conn_; } + void poison() { poisoned_ = true; } +private: + PgConnectionPool& pool_; + std::string name_; + PGconn* conn_; + bool poisoned_ = false; +}; + +class PooledPgSession : public RemoteSession { +public: + PooledPgSession(PgConnectionPool& pool, std::string name) + : pool_(pool), name_(std::move(name)), conn_(pool_.checkout(name_)) {} + ~PooledPgSession() override { + if (!conn_) return; + if (poisoned_) PQfinish(conn_); + else pool_.checkin(name_, conn_); + } + PooledPgSession(const PooledPgSession&) = delete; + PooledPgSession& operator=(const PooledPgSession&) = delete; + + ResultSet execute(sql_parser::StringRef sql) override { + ResultSet rs; + if (!conn_) { poisoned_ = true; return rs; } + std::string q(sql.ptr, sql.len); + PGresult* res = PQexec(conn_, q.c_str()); + if (!res) { poisoned_ = true; return rs; } + ExecStatusType st = PQresultStatus(res); + if (st != PGRES_TUPLES_OK) { + PQclear(res); + return rs; + } + rs = detail::pg_result_to_resultset_impl(res); + PQclear(res); + return rs; + } + + DmlResult execute_dml(sql_parser::StringRef sql) override { + DmlResult result; + if (!conn_) { + poisoned_ = true; + result.error_message = "no connection"; + return result; + } + std::string q(sql.ptr, sql.len); + PGresult* res = PQexec(conn_, q.c_str()); + if (!res) { + poisoned_ = true; + result.error_message = "PQexec returned null"; + return result; + } + ExecStatusType st = PQresultStatus(res); + if (st != PGRES_COMMAND_OK && st != PGRES_TUPLES_OK) { + result.error_message = PQresultErrorMessage(res); + PQclear(res); + return result; + } + const char* tuples = PQcmdTuples(res); + if (tuples && tuples[0]) + result.affected_rows = static_cast(std::strtoull(tuples, nullptr, 10)); + result.success = true; + PQclear(res); + return result; + } + + void poison() override { poisoned_ = true; } + +private: + PgConnectionPool& pool_; + std::string name_; + PGconn* conn_; + bool poisoned_ = false; +}; + class ThreadSafeMultiRemoteExecutor : public RemoteExecutor { public: ThreadSafeMultiRemoteExecutor() = default; ~ThreadSafeMultiRemoteExecutor() override = default; void add_backend(const BackendConfig& config) { - pool_.add_backend(config); - // Track dialect per backend (currently only MySQL is pooled) std::lock_guard lk(mu_); backend_dialects_[config.name] = config.dialect; + if (config.dialect == sql_parser::Dialect::PostgreSQL) + pg_pool_.add_backend(config); + else + pool_.add_backend(config); } ResultSet execute(const char* backend_name, sql_parser::StringRef sql) override { + if (is_pg(backend_name)) { + PgConnectionGuard guard(pg_pool_, std::string(backend_name)); + ResultSet rs; + PGconn* conn = guard.get(); + if (!conn) { guard.poison(); return rs; } + std::string q(sql.ptr, sql.len); + PGresult* res = PQexec(conn, q.c_str()); + if (!res) { guard.poison(); return rs; } + if (PQresultStatus(res) != PGRES_TUPLES_OK) { + PQclear(res); + return rs; + } + rs = detail::pg_result_to_resultset_impl(res); + PQclear(res); + return rs; + } ConnectionGuard guard(pool_, std::string(backend_name)); ResultSet rs; MYSQL* conn = guard.get(); @@ -326,6 +514,35 @@ class ThreadSafeMultiRemoteExecutor : public RemoteExecutor { } DmlResult execute_dml(const char* backend_name, sql_parser::StringRef sql) override { + if (is_pg(backend_name)) { + PgConnectionGuard guard(pg_pool_, std::string(backend_name)); + DmlResult result; + PGconn* conn = guard.get(); + if (!conn) { + guard.poison(); + result.error_message = "failed to acquire connection"; + return result; + } + std::string q(sql.ptr, sql.len); + PGresult* res = PQexec(conn, q.c_str()); + if (!res) { + guard.poison(); + result.error_message = "PQexec returned null"; + return result; + } + ExecStatusType st = PQresultStatus(res); + if (st != PGRES_COMMAND_OK && st != PGRES_TUPLES_OK) { + result.error_message = PQresultErrorMessage(res); + PQclear(res); + return result; + } + const char* tuples = PQcmdTuples(res); + if (tuples && tuples[0]) + result.affected_rows = static_cast(std::strtoull(tuples, nullptr, 10)); + result.success = true; + PQclear(res); + return result; + } ConnectionGuard guard(pool_, std::string(backend_name)); DmlResult result; MYSQL* conn = guard.get(); @@ -363,6 +580,8 @@ class ThreadSafeMultiRemoteExecutor : public RemoteExecutor { // the connection is returned to the pool on destruction (or closed // if the session was poisoned). std::unique_ptr checkout_session(const char* backend_name) override { + if (is_pg(backend_name)) + return std::make_unique(pg_pool_, std::string(backend_name)); return std::make_unique(pool_, std::string(backend_name)); } @@ -372,9 +591,17 @@ class ThreadSafeMultiRemoteExecutor : public RemoteExecutor { private: ConnectionPool pool_; - std::mutex mu_; + PgConnectionPool pg_pool_; + mutable std::mutex mu_; std::unordered_map backend_dialects_; + bool is_pg(const char* backend_name) const { + std::lock_guard lk(mu_); + auto it = backend_dialects_.find(backend_name ? backend_name : ""); + return it != backend_dialects_.end() && + it->second == sql_parser::Dialect::PostgreSQL; + } + // Wrap the free function for the legacy (unpinned) pooled execute path. static ResultSet mysql_result_to_resultset(MYSQL_RES* res) { return detail::mysql_result_to_resultset_impl(res); diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index 8a51eaa..62000b0 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -808,6 +808,72 @@ TEST_F(DistributedPlannerTest, ShardRouting_NoShardKey_AllShards) { << "Non-shard-key filter should query all shards"; } +TEST_F(DistributedPlannerTest, ShardRouting_OrOfEqualitiesPrunesUnion) { + Parser parser; + auto pr = parser.parse("SELECT * FROM users WHERE id = 4 OR id = 4", 42); + 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); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + EXPECT_EQ(remotes.size(), 1u); +} + +TEST_F(DistributedPlannerTest, ShardRouting_OrWithNonKeyScatters) { + Parser parser; + auto pr = parser.parse("SELECT * FROM users WHERE id = 4 OR age > 20", 45); + 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); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + EXPECT_EQ(remotes.size(), 3u); +} + +TEST_F(DistributedPlannerTest, ShardRouting_PlaceholderScatters) { + Parser parser; + auto pr = parser.parse("SELECT * FROM users WHERE id = ?", 32); + 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); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + EXPECT_EQ(remotes.size(), 3u); +} + +TEST_F(DistributedPlannerTest, ShardRouting_RangeInequalityPrunes) { + TableShardConfig cfg; + cfg.table_name = "users"; + cfg.shard_key = "id"; + cfg.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + cfg.strategy = RoutingStrategy::RANGE; + cfg.ranges = {{5, 0}, {10, 1}, {100000, 2}}; + shard_map.add_table(cfg); + + auto count_remotes = [&](const char* sql) { + Parser parser; + auto pr = parser.parse(sql, std::strlen(sql)); + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + return remotes.size(); + }; + + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id <= 5"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id > 10"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id BETWEEN 6 AND 10"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id >= 1 AND id <= 15"), 3u); +} + TEST_F(DistributedPlannerTest, ShardRouting_Correctness) { // The shard routing uses FNV-1a 64-bit hash of the key % num_shards // to pick a shard, but the test fixture uses sequential partitioning diff --git a/tests/test_pgsql_executor.cpp b/tests/test_pgsql_executor.cpp index 777f623..44f18f9 100644 --- a/tests/test_pgsql_executor.cpp +++ b/tests/test_pgsql_executor.cpp @@ -1,5 +1,6 @@ #include #include "sql_engine/pgsql_remote_executor.h" +#include "sql_engine/thread_safe_executor.h" #include "sql_engine/backend_config.h" #include @@ -227,4 +228,13 @@ TEST_F(PgSQLExecutorTest, DateType) { EXPECT_EQ(rs.rows[0].get(0).tag, sql_engine::Value::TAG_DATE); } +TEST(ThreadSafePgTest, PooledSelect) { + SKIP_IF_NO_PGSQL(); + sql_engine::ThreadSafeMultiRemoteExecutor exec; + exec.add_backend(make_pgsql_config("pool_pg")); + sql_parser::StringRef sql{"SELECT 1", 8}; + auto rs = exec.execute("pool_pg", sql); + EXPECT_EQ(rs.row_count(), 1u); +} + } // namespace diff --git a/tests/test_shard_map.cpp b/tests/test_shard_map.cpp index 515de3f..ae7cf64 100644 --- a/tests/test_shard_map.cpp +++ b/tests/test_shard_map.cpp @@ -9,6 +9,7 @@ #include "sql_engine/shard_map.h" #include +#include using namespace sql_engine; using sql_parser::StringRef; @@ -108,6 +109,23 @@ TEST(ShardMapRangeTest, MatchesDemoDataPlacement) { EXPECT_EQ(map.shard_index_for_int(sref("users"), 10), 1u); } +TEST(ShardMapRangeTest, CollectIntRangeShards) { + TableShardConfig cfg = make_two_shards(RoutingStrategy::RANGE); + cfg.ranges = {ShardRange{5, 0}, ShardRange{10, 1}}; + ShardMap map; + map.add_table(cfg); + + std::vector lo; + map.collect_int_range_shards(sref("users"), INT64_MIN, 5, lo); + ASSERT_EQ(lo.size(), 1u); + EXPECT_EQ(lo[0], 0u); + + std::vector hi; + map.collect_int_range_shards(sref("users"), 6, INT64_MAX, hi); + ASSERT_EQ(hi.size(), 1u); + EXPECT_EQ(hi[0], 1u); +} + TEST(ShardMapRangeTest, AboveMaxBoundFallsToLastShard) { TableShardConfig cfg = make_two_shards(RoutingStrategy::RANGE); cfg.ranges = {ShardRange{5, 0}, ShardRange{10, 1}}; diff --git a/tools/engine_stress_test.cpp b/tools/engine_stress_test.cpp index 6d79cf0..492df7c 100644 --- a/tools/engine_stress_test.cpp +++ b/tools/engine_stress_test.cpp @@ -35,6 +35,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" @@ -103,14 +104,13 @@ static void worker_thread( (void)thread_id; // Each thread gets its own arena, txn manager, executor, and session - Arena txn_arena{65536, 1048576}; - LocalTransactionManager txn_mgr(txn_arena); - ThreadSafeMultiRemoteExecutor remote_exec; for (auto& bc : backends) { remote_exec.add_backend(bc); } + DistributedTransactionManager txn_mgr( + remote_exec, DistributedTransactionManager::BackendDialect::MYSQL); Session session(catalog, txn_mgr); session.set_remote_executor(&remote_exec); session.set_parallel_open(true); // thread-safe executor enables parallel shard I/O From 7ade9219dd12b0a9ab98fb44964d1f8e16eb4ce3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 04:36:24 +0700 Subject: [PATCH 2/3] feat: parallel scatter, composite keys, LIST prune, plan-cache redistribute Flatten UNION ALL remotes so N shards open together. Route composite HASH keys, fail unknown tables, and move rows on shard-key UPDATE. LIST misses fail closed and BETWEEN/IN prune; composite colocated joins push when every key part is equated. Cache the logical plan and re-distribute on each hit. --- include/sql_engine/distributed_planner.h | 638 ++++++++++++++++++----- include/sql_engine/operators/set_op_op.h | 105 ++-- include/sql_engine/plan_executor.h | 35 +- include/sql_engine/session.h | 40 +- include/sql_engine/shard_map.h | 161 +++++- src/sql_engine/tool_config_parser.cpp | 12 + tests/test_distributed_dml.cpp | 180 ++++++- tests/test_distributed_planner.cpp | 142 +++++ tests/test_operators.cpp | 20 + tests/test_shard_map.cpp | 62 ++- tests/test_ssl_config.cpp | 16 + 11 files changed, 1198 insertions(+), 213 deletions(-) diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 251109e..1a8491f 100644 --- a/include/sql_engine/distributed_planner.h +++ b/include/sql_engine/distributed_planner.h @@ -12,9 +12,12 @@ #include "sql_engine/result_set.h" #include "sql_engine/plan_builder.h" #include "sql_engine/plan_executor.h" +#include "sql_engine/catalog_resolver.h" #include "sql_parser/arena.h" #include "sql_parser/ast.h" #include "sql_parser/common.h" +#include "sql_parser/string_builder.h" +#include "sql_parser/emitter.h" #include #include #include @@ -22,6 +25,7 @@ #include #include #include +#include namespace sql_engine { @@ -318,7 +322,8 @@ class DistributedPlanner { const TableInfo* table = scan_node->scan.table; if (!table) return scan_node; - if (!shards_.has_table(table->table_name)) return scan_node; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); if (!shards_.is_sharded(table->table_name)) { // Case 1: Unsharded -- single RemoteScan @@ -348,13 +353,18 @@ class DistributedPlanner { const std::vector& all_shards) { if (!where_expr || all_shards.empty()) return all_shards; - sql_parser::StringRef shard_key = shards_.get_shard_key(table->table_name); - if (!shard_key.ptr || shard_key.len == 0) return all_shards; + const auto& keys = shards_.get_shard_keys(table->table_name); + if (keys.empty()) return all_shards; - // Try to extract shard key literal values from WHERE expression std::vector target_indices; - extract_shard_targets(where_expr, shard_key, table->table_name, - all_shards.size(), target_indices); + if (keys.size() > 1) { + extract_composite_targets(where_expr, keys, table->table_name, target_indices); + } else { + sql_parser::StringRef shard_key{keys[0].c_str(), + static_cast(keys[0].size())}; + extract_shard_targets(where_expr, shard_key, table->table_name, + all_shards.size(), target_indices); + } if (target_indices.empty()) return all_shards; @@ -396,8 +406,9 @@ class DistributedPlanner { col_node = right_node; lit_node = left_node; } if (col_node && lit_node) { - size_t idx = literal_to_shard_index(lit_node, table_name, num_shards); - target_indices.push_back(idx); + size_t idx = 0; + if (try_literal_to_shard_index(lit_node, table_name, idx)) + target_indices.push_back(idx); return; } } @@ -494,7 +505,9 @@ class DistributedPlanner { if (col_expr && is_shard_key_ref(col_expr, shard_key)) { for (const sql_parser::AstNode* item = col_expr->next_sibling; item; item = item->next_sibling) { if (is_literal(item)) { - target_indices.push_back(literal_to_shard_index(item, table_name, num_shards)); + size_t idx = 0; + if (try_literal_to_shard_index(item, table_name, idx)) + target_indices.push_back(idx); } else { // Non-literal in IN list -- can't prune target_indices.clear(); @@ -504,15 +517,23 @@ class DistributedPlanner { } } - if (expr->type == sql_parser::NodeType::NODE_BETWEEN && - shards_.routing_strategy(table_name) == RoutingStrategy::RANGE) { - const sql_parser::AstNode* col = expr->first_child; - const sql_parser::AstNode* lo = col ? col->next_sibling : nullptr; - const sql_parser::AstNode* hi = lo ? lo->next_sibling : nullptr; - if (col && is_shard_key_ref(col, shard_key) && is_literal(lo) && is_literal(hi)) { - shards_.collect_int_range_shards(table_name, - literal_to_int(lo), literal_to_int(hi), - target_indices); + if (expr->type == sql_parser::NodeType::NODE_BETWEEN) { + RoutingStrategy strat = shards_.routing_strategy(table_name); + if (strat == RoutingStrategy::RANGE || strat == RoutingStrategy::LIST) { + const sql_parser::AstNode* col = expr->first_child; + const sql_parser::AstNode* lo = col ? col->next_sibling : nullptr; + const sql_parser::AstNode* hi = lo ? lo->next_sibling : nullptr; + if (col && is_shard_key_ref(col, shard_key) && is_literal(lo) && is_literal(hi)) { + if (strat == RoutingStrategy::RANGE) { + shards_.collect_int_range_shards(table_name, + literal_to_int(lo), literal_to_int(hi), + target_indices); + } else { + shards_.collect_int_list_shards(table_name, + literal_to_int(lo), literal_to_int(hi), + target_indices); + } + } } } } @@ -551,27 +572,27 @@ class DistributedPlanner { node->type == sql_parser::NodeType::NODE_LITERAL_STRING; } - size_t literal_to_shard_index(const sql_parser::AstNode* lit, + bool try_literal_to_shard_index(const sql_parser::AstNode* lit, sql_parser::StringRef table_name, - size_t num_shards) const { - if (!lit || num_shards == 0) return 0; + size_t& out) const { + if (!lit) return false; if (lit->type == sql_parser::NodeType::NODE_LITERAL_INT) { sql_parser::StringRef sv = lit->value(); int64_t val = 0; if (sv.ptr && sv.len > 0) val = std::strtoll(sv.ptr, nullptr, 10); - return shards_.shard_index_for_int(table_name, val); + return shards_.try_shard_index_for_int(table_name, val, out); } if (lit->type == sql_parser::NodeType::NODE_LITERAL_STRING) { sql_parser::StringRef sv = lit->value(); - return shards_.shard_index_for_string(table_name, sv.ptr, sv.len); + return shards_.try_shard_index_for_string(table_name, sv.ptr, sv.len, out); } if (lit->type == sql_parser::NodeType::NODE_LITERAL_FLOAT) { sql_parser::StringRef sv = lit->value(); double dv = sv.ptr ? std::strtod(sv.ptr, nullptr) : 0.0; - int64_t iv = static_cast(dv); - return shards_.shard_index_for_int(table_name, iv); + return shards_.try_shard_index_for_int( + table_name, static_cast(dv), out); } - return 0; + return false; } // Build N RemoteScans with UNION ALL @@ -628,6 +649,7 @@ class DistributedPlanner { PlanNode* make_remote_scan(const char* backend, sql_parser::StringRef sql, const TableInfo* table) { + if (!backend) return fail_dml("table not in shard map"); PlanNode* node = make_plan_node(arena_, PlanNodeType::REMOTE_SCAN); // Copy backend name to arena uint32_t blen = static_cast(std::strlen(backend)); @@ -672,7 +694,9 @@ class DistributedPlanner { } const TableInfo* table = ctx.scan->scan.table; - if (!shards_.has_table(table->table_name) || !shards_.is_sharded(table->table_name)) { + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); + if (!shards_.is_sharded(table->table_name)) { return make_unsharded_aggregate(agg_node, ctx, table); } @@ -930,12 +954,8 @@ class DistributedPlanner { } const TableInfo* table = ctx.scan->scan.table; - if (!shards_.has_table(table->table_name)) { - PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT); - result->sort = sort_node->sort; - result->left = distribute_node(sort_node->left); - return result; - } + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); if (!all_sort_keys_are_table_columns(sort_node, table)) { return local_sort(sort_node); @@ -1092,6 +1112,58 @@ class DistributedPlanner { (is_shard_key_ref(l, right_key) && is_shard_key_ref(r, left_key)); } + void collect_eq_pairs(const sql_parser::AstNode* expr, + std::vector>& out) const { + if (!expr || expr->type != sql_parser::NodeType::NODE_BINARY_OP) return; + sql_parser::StringRef op = expr->value(); + if (op.len == 1 && op.ptr[0] == '=') { + const sql_parser::AstNode* l = expr->first_child; + const sql_parser::AstNode* r = l ? l->next_sibling : nullptr; + if (l && r) out.push_back({l, r}); + return; + } + if (op.len == 3 && + (op.ptr[0] == 'A' || op.ptr[0] == 'a') && + (op.ptr[1] == 'N' || op.ptr[1] == 'n') && + (op.ptr[2] == 'D' || op.ptr[2] == 'd')) { + collect_eq_pairs(expr->first_child, out); + if (expr->first_child) + collect_eq_pairs(expr->first_child->next_sibling, out); + } + } + + bool join_covers_composite_keys(const sql_parser::AstNode* cond, + const std::vector& left_keys, + const std::vector& right_keys) const { + if (left_keys.empty() || left_keys.size() != right_keys.size()) return false; + if (left_keys.size() == 1) { + sql_parser::StringRef lk{left_keys[0].c_str(), + static_cast(left_keys[0].size())}; + sql_parser::StringRef rk{right_keys[0].c_str(), + static_cast(right_keys[0].size())}; + return join_on_shard_keys(cond, lk, rk); + } + std::vector> eqs; + collect_eq_pairs(cond, eqs); + for (size_t i = 0; i < left_keys.size(); ++i) { + sql_parser::StringRef lk{left_keys[i].c_str(), + static_cast(left_keys[i].size())}; + sql_parser::StringRef rk{right_keys[i].c_str(), + static_cast(right_keys[i].size())}; + bool found = false; + for (const auto& eq : eqs) { + if ((is_shard_key_ref(eq.first, lk) && is_shard_key_ref(eq.second, rk)) || + (is_shard_key_ref(eq.first, rk) && is_shard_key_ref(eq.second, lk))) { + found = true; + break; + } + } + if (!found) return false; + } + return true; + } + PlanNode* distribute_colocated_join(PlanNode* join_node, const TableInfo* left_table, const TableInfo* right_table) { @@ -1139,9 +1211,9 @@ class DistributedPlanner { 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))) { + join_covers_composite_keys(join_node->join.condition, + shards_.get_shard_keys(left_table->table_name), + shards_.get_shard_keys(right_table->table_name))) { return distribute_colocated_join(join_node, left_table, right_table); } @@ -1198,7 +1270,9 @@ class DistributedPlanner { PlanNode* distribute_insert(PlanNode* plan) { const auto& ip = plan->insert_plan; const TableInfo* table = ip.table; - if (!table || !shards_.has_table(table->table_name)) return plan; + if (!table) return plan; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); // Check for INSERT ... SELECT (select_source stores the SELECT AST) if (ip.select_source && ip.select_source->type == PlanNodeType::DERIVED_SCAN @@ -1217,53 +1291,44 @@ class DistributedPlanner { return make_remote_scan(shards_.get_backend(table->table_name), sql, table); } - // Sharded: group rows by shard key value - sql_parser::StringRef shard_key = shards_.get_shard_key(table->table_name); - if (!shard_key.ptr) return plan; + const auto& key_names = shards_.get_shard_keys(table->table_name); + if (key_names.empty()) return plan; - // Find shard key column ordinal in the column list - int shard_col_idx = -1; - if (ip.columns && ip.column_count > 0) { - for (uint16_t i = 0; i < ip.column_count; ++i) { - if (ip.columns[i] && ip.columns[i]->value().equals_ci(shard_key.ptr, shard_key.len)) { - shard_col_idx = static_cast(i); - break; - } - } - } else if (table) { - // No explicit column list -- match by table column order - for (uint16_t i = 0; i < table->column_count; ++i) { - if (table->columns[i].name.equals_ci(shard_key.ptr, shard_key.len)) { - shard_col_idx = static_cast(i); - break; - } - } + std::vector key_ords(key_names.size(), -1); + for (size_t k = 0; k < key_names.size(); ++k) { + sql_parser::StringRef kn{key_names[k].c_str(), + static_cast(key_names[k].size())}; + key_ords[k] = find_insert_key_ordinal(ip.columns, ip.column_count, table, kn); } - - if (shard_col_idx < 0) { - return fail_dml("cannot route INSERT: shard key column is not present"); + for (int ord : key_ords) { + if (ord < 0) + return fail_dml("cannot route INSERT: shard key column is not present"); } const auto& shard_list = shards_.get_shards(table->table_name); - // Group rows by ShardMap route. Map: shard_index -> list of row indices. std::unordered_map> shard_rows; for (uint16_t ri = 0; ri < ip.row_count; ++ri) { const sql_parser::AstNode* row_ast = ip.value_rows[ri]; if (!row_ast) continue; - const sql_parser::AstNode* expr = row_ast->first_child; - for (int j = 0; j < shard_col_idx && expr; ++j) { - expr = expr->next_sibling; + std::vector parts; + parts.reserve(key_ords.size()); + for (int ord : key_ords) { + const sql_parser::AstNode* expr = row_ast->first_child; + for (int j = 0; j < ord && expr; ++j) expr = expr->next_sibling; + if (!expr) + return fail_dml("cannot route INSERT: missing shard key value"); + Value v = evaluate_shard_key_value(expr); + if (!value_is_routable(v)) + return fail_dml("cannot route INSERT: shard key is not a literal"); + parts.push_back(value_to_part(v)); } - if (!expr) { - return fail_dml("cannot route INSERT: missing shard key value"); - } - size_t shard_idx = 0; - if (!route_value(table->table_name, evaluate_shard_key_value(expr), shard_idx)) { - return fail_dml("cannot route INSERT: shard key is not a literal"); + if (!shards_.try_shard_index_for_parts(table->table_name, parts.data(), + parts.size(), shard_idx)) { + return fail_dml("cannot route INSERT: shard key value is not mapped"); } shard_rows[shard_idx].push_back(ri); } @@ -1317,7 +1382,9 @@ class DistributedPlanner { return distribute_multi_table_dml(up.original_ast, table, true); } - if (!table || !shards_.has_table(table->table_name)) return plan; + if (!table) return plan; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); // Check for cross-shard subqueries in WHERE and rewrite const sql_parser::AstNode* where_expr = up.where_expr; @@ -1332,9 +1399,8 @@ class DistributedPlanner { return make_remote_scan(shards_.get_backend(table->table_name), sql, table); } - sql_parser::StringRef shard_key = shards_.get_shard_key(table->table_name); - if (assigns_shard_key(up.set_columns, up.set_count, shard_key)) { - return fail_dml("cannot UPDATE shard key column"); + if (assigns_any_shard_key(up.set_columns, up.set_count, table->table_name)) { + return distribute_update_move(plan, table, where_expr); } const auto& shard_list = shards_.get_shards(table->table_name); @@ -1360,7 +1426,9 @@ class DistributedPlanner { return distribute_multi_table_dml(dp.original_ast, table, false); } - if (!table || !shards_.has_table(table->table_name)) return plan; + if (!table) return plan; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); // Check for cross-shard subqueries in WHERE and rewrite const sql_parser::AstNode* where_expr = dp.where_expr; @@ -1413,19 +1481,16 @@ class DistributedPlanner { } bool route_value(sql_parser::StringRef table_name, const Value& v, size_t& shard_idx) const { - if (v.tag == Value::TAG_INT64) { - shard_idx = shards_.shard_index_for_int(table_name, v.int_val); - return true; - } + if (!shards_.has_table(table_name)) return false; + if (v.tag == Value::TAG_INT64) + return shards_.try_shard_index_for_int(table_name, v.int_val, shard_idx); if (v.tag == Value::TAG_UINT64) { - shard_idx = shards_.shard_index_for_int( - table_name, static_cast(v.uint_val)); - return true; + return shards_.try_shard_index_for_int( + table_name, static_cast(v.uint_val), shard_idx); } if (v.tag == Value::TAG_STRING && v.str_val.ptr) { - shard_idx = shards_.shard_index_for_string( - table_name, v.str_val.ptr, v.str_val.len); - return true; + return shards_.try_shard_index_for_string( + table_name, v.str_val.ptr, v.str_val.len, shard_idx); } return false; } @@ -1439,6 +1504,341 @@ class DistributedPlanner { return false; } + bool assigns_any_shard_key(const sql_parser::AstNode** set_columns, uint16_t set_count, + sql_parser::StringRef table_name) const { + for (const auto& k : shards_.get_shard_keys(table_name)) { + sql_parser::StringRef kn{k.c_str(), static_cast(k.size())}; + if (assigns_shard_key(set_columns, set_count, kn)) return true; + } + return false; + } + + static bool value_is_routable(const Value& v) { + return v.tag == Value::TAG_INT64 || v.tag == Value::TAG_UINT64 || + (v.tag == Value::TAG_STRING && v.str_val.ptr); + } + + static ShardKeyPart value_to_part(const Value& v) { + ShardKeyPart p; + if (v.tag == Value::TAG_INT64) { + p.is_int = true; + p.int_val = v.int_val; + } else if (v.tag == Value::TAG_UINT64) { + p.is_int = true; + p.int_val = static_cast(v.uint_val); + } else { + p.is_int = false; + p.str = v.str_val.ptr; + p.str_len = v.str_val.len; + } + return p; + } + + static int find_insert_key_ordinal(const sql_parser::AstNode** columns, + uint16_t column_count, + const TableInfo* table, + sql_parser::StringRef key) { + if (columns && column_count > 0) { + for (uint16_t i = 0; i < column_count; ++i) { + if (columns[i] && columns[i]->value().equals_ci(key.ptr, key.len)) + return static_cast(i); + } + return -1; + } + if (!table) return -1; + for (uint16_t i = 0; i < table->column_count; ++i) { + if (table->columns[i].name.equals_ci(key.ptr, key.len)) + return static_cast(i); + } + return -1; + } + + void collect_key_eq_values(const sql_parser::AstNode* expr, + sql_parser::StringRef key, + std::vector& out) { + if (!expr) return; + if (expr->type == sql_parser::NodeType::NODE_BINARY_OP) { + sql_parser::StringRef op = expr->value(); + if (op.len == 1 && op.ptr[0] == '=') { + const sql_parser::AstNode* l = expr->first_child; + const sql_parser::AstNode* r = l ? l->next_sibling : nullptr; + if (l && r) { + const sql_parser::AstNode* lit = nullptr; + if (is_shard_key_ref(l, key) && is_literal(r)) lit = r; + else if (is_shard_key_ref(r, key) && is_literal(l)) lit = l; + if (lit) { + Value v = evaluate_shard_key_value(lit); + if (value_is_routable(v)) out.push_back(v); + } + } + return; + } + if (op.len == 3 && + (op.ptr[0] == 'A' || op.ptr[0] == 'a') && + (op.ptr[1] == 'N' || op.ptr[1] == 'n') && + (op.ptr[2] == 'D' || op.ptr[2] == 'd')) { + collect_key_eq_values(expr->first_child, key, out); + if (expr->first_child) + collect_key_eq_values(expr->first_child->next_sibling, key, out); + } + return; + } + if (expr->type == sql_parser::NodeType::NODE_IN_LIST) { + const sql_parser::AstNode* col = expr->first_child; + if (col && is_shard_key_ref(col, key)) { + for (const sql_parser::AstNode* item = col->next_sibling; item; + item = item->next_sibling) { + if (!is_literal(item)) { + out.clear(); + return; + } + Value v = evaluate_shard_key_value(item); + if (value_is_routable(v)) out.push_back(v); + } + } + } + } + + void extract_composite_targets(const sql_parser::AstNode* where_expr, + const std::vector& keys, + sql_parser::StringRef table_name, + std::vector& target_indices) { + std::vector> dims(keys.size()); + for (size_t k = 0; k < keys.size(); ++k) { + sql_parser::StringRef kn{keys[k].c_str(), + static_cast(keys[k].size())}; + collect_key_eq_values(where_expr, kn, dims[k]); + if (dims[k].empty()) return; + } + std::vector cursor(keys.size(), 0); + for (;;) { + std::vector parts(keys.size()); + for (size_t k = 0; k < keys.size(); ++k) + parts[k] = value_to_part(dims[k][cursor[k]]); + size_t idx = 0; + if (shards_.try_shard_index_for_parts(table_name, parts.data(), + parts.size(), idx)) + target_indices.push_back(idx); + size_t d = keys.size(); + while (d-- > 0) { + if (++cursor[d] < dims[d].size()) break; + cursor[d] = 0; + } + if (d == static_cast(-1)) break; + } + } + + PlanNode* append_remote(PlanNode* current, PlanNode* next) { + if (!next) return current; + if (!current) return next; + 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 = next; + return union_node; + } + + Value copy_value_arena(const Value& v) { + if ((v.tag == Value::TAG_STRING || v.tag == Value::TAG_DECIMAL || + v.tag == Value::TAG_BYTES || v.tag == Value::TAG_JSON) && + v.str_val.ptr && v.str_val.len > 0) { + char* p = static_cast(arena_.allocate(v.str_val.len)); + std::memcpy(p, v.str_val.ptr, v.str_val.len); + Value out = v; + out.str_val = sql_parser::StringRef{p, v.str_val.len}; + return out; + } + return v; + } + + Row copy_row_arena(const Row& src) { + Row dst = make_row(arena_, src.column_count); + for (uint16_t i = 0; i < src.column_count; ++i) + dst.set(i, copy_value_arena(src.get(i))); + return dst; + } + + bool route_row_keys(const TableInfo* table, const Row& row, size_t& idx) const { + const auto& keys = shards_.get_shard_keys(table->table_name); + if (keys.empty()) return false; + std::vector parts; + parts.reserve(keys.size()); + for (const auto& k : keys) { + sql_parser::StringRef kn{k.c_str(), static_cast(k.size())}; + const ColumnInfo* col = catalog_.get_column(table, kn); + if (!col || col->ordinal >= row.column_count) return false; + Value v = row.get(col->ordinal); + if (!value_is_routable(v)) return false; + parts.push_back(value_to_part(v)); + } + return shards_.try_shard_index_for_parts( + table->table_name, parts.data(), parts.size(), idx); + } + + Row apply_update_set(const Row& src, const TableInfo* table, + const sql_parser::AstNode** set_cols, + const sql_parser::AstNode** set_exprs, + uint16_t set_count) { + Row dst = copy_row_arena(src); + auto resolve = make_resolver(catalog_, table, src.values); + for (uint16_t i = 0; i < set_count; ++i) { + if (!set_cols[i]) continue; + const ColumnInfo* col = catalog_.get_column(table, set_cols[i]->value()); + if (!col) continue; + Value nv = value_null(); + if (functions_) { + nv = evaluate_expression(set_exprs[i], resolve, *functions_, arena_); + } else { + nv = evaluate_shard_key_value(set_exprs[i]); + } + dst.set(col->ordinal, copy_value_arena(nv)); + } + return dst; + } + + sql_parser::StringRef build_identity_pred(const TableInfo* table, const Row& row, + sql_parser::StringBuilder& sb) { + for (uint16_t i = 0; i < table->column_count && i < row.column_count; ++i) { + if (i > 0) sb.append(" AND "); + sb.append(table->columns[i].name.ptr, table->columns[i].name.len); + if (row.get(i).is_null()) { + sb.append(" IS NULL", 8); + } else { + sb.append(" = "); + emit_value(row.get(i), sb); + } + } + return sb.finish(); + } + + sql_parser::StringRef build_delete_identity(const TableInfo* table, const Row& row) { + sql_parser::StringBuilder sb(arena_, 256); + sb.append("DELETE FROM "); + sb.append(table->table_name.ptr, table->table_name.len); + sb.append(" WHERE "); + return build_identity_pred(table, row, sb); + } + + sql_parser::StringRef build_update_identity(const TableInfo* table, + const sql_parser::AstNode** set_cols, + const sql_parser::AstNode** set_exprs, + uint16_t set_count, + const Row& old_row) { + sql_parser::StringBuilder sb(arena_, 256); + sb.append("UPDATE "); + sb.append(table->table_name.ptr, table->table_name.len); + sb.append(" SET "); + for (uint16_t i = 0; i < set_count; ++i) { + if (i > 0) sb.append(", "); + if (set_cols[i]) { + sql_parser::StringRef cn = set_cols[i]->value(); + sb.append(cn.ptr, cn.len); + } + sb.append(" = "); + if (set_exprs[i]) { + sql_parser::Emitter emitter(arena_); + emitter.emit(set_exprs[i]); + sql_parser::StringRef ev = emitter.result(); + sb.append(ev.ptr, ev.len); + } + } + sb.append(" WHERE "); + return build_identity_pred(table, old_row, sb); + } + + PlanNode* distribute_update_move(PlanNode* plan, const TableInfo* table, + const sql_parser::AstNode* where_expr) { + if (!remote_executor_) + return fail_dml("cannot UPDATE shard key without a remote executor"); + + const auto& up = plan->update_plan; + const auto& shard_list = shards_.get_shards(table->table_name); + std::vector pruned = prune_shards(table, where_expr, shard_list); + if (pruned.empty()) return plan; + + struct Move { + size_t src = 0; + size_t dst = 0; + Row old_row{}; + Row new_row{}; + }; + std::vector moves; + + for (const auto& shard : pruned) { + size_t src = 0; + for (size_t i = 0; i < shard_list.size(); ++i) { + if (shard_list[i].backend_name == shard.backend_name) { + src = i; + break; + } + } + sql_parser::StringRef sql = qb_.build_select( + table, where_expr, nullptr, 0, nullptr, 0, + nullptr, nullptr, 0, -1, false); + ResultSet rs = remote_executor_->execute(shard.backend_name.c_str(), sql); + for (const auto& row : rs.rows) { + Move m; + m.src = src; + m.old_row = copy_row_arena(row); + m.new_row = apply_update_set(m.old_row, table, up.set_columns, + up.set_exprs, up.set_count); + if (!route_row_keys(table, m.new_row, m.dst)) + return fail_dml("cannot UPDATE shard key: new key is not routable"); + moves.push_back(m); + } + } + + if (moves.empty()) { + sql_parser::StringRef sql = qb_.build_update( + table, up.set_columns, up.set_exprs, up.set_count, where_expr); + return make_remote_scan(pruned[0].backend_name.c_str(), sql, table); + } + + bool any_move = false; + for (const auto& m : moves) { + if (m.src != m.dst) { any_move = true; break; } + } + if (!any_move) { + if (pruned.size() == 1) { + sql_parser::StringRef sql = qb_.build_update( + table, up.set_columns, up.set_exprs, up.set_count, where_expr); + return make_remote_scan(pruned[0].backend_name.c_str(), sql, table); + } + const sql_parser::AstNode* final_where = where_expr; + return scatter_dml_to_shards(table, pruned, [&]() { + return qb_.build_update( + table, up.set_columns, up.set_exprs, up.set_count, final_where); + }); + } + + PlanNode* current = nullptr; + std::unordered_map> inserts; + for (const auto& m : moves) { + if (m.src == m.dst) { + sql_parser::StringRef sql = build_update_identity( + table, up.set_columns, up.set_exprs, up.set_count, m.old_row); + current = append_remote( + current, make_remote_scan(shard_list[m.src].backend_name.c_str(), + sql, table)); + } else { + current = append_remote( + current, + make_remote_scan(shard_list[m.src].backend_name.c_str(), + build_delete_identity(table, m.old_row), table)); + inserts[m.dst].push_back(m.new_row); + } + } + for (auto& kv : inserts) { + current = append_remote( + current, + make_remote_scan(shard_list[kv.first].backend_name.c_str(), + build_insert_from_rows(table, nullptr, 0, kv.second), + table)); + } + return current ? current : plan; + } + bool is_column_ref(const sql_parser::AstNode* node, sql_parser::StringRef col_name) const { if (!node) return false; if (node->type == sql_parser::NodeType::NODE_COLUMN_REF || @@ -1834,7 +2234,8 @@ class DistributedPlanner { } // Determine target shards for each row - if (!shards_.has_table(table->table_name)) return plan; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); if (!shards_.is_sharded(table->table_name)) { // Unsharded: build a single INSERT with all rows @@ -1843,42 +2244,37 @@ class DistributedPlanner { return make_remote_scan(shards_.get_backend(table->table_name), sql, table); } - // Sharded: group rows by shard key - sql_parser::StringRef shard_key = shards_.get_shard_key(table->table_name); - if (!shard_key.ptr) return plan; + const auto& key_names = shards_.get_shard_keys(table->table_name); + if (key_names.empty()) return plan; - // Find shard key column ordinal in the result set - int shard_col_idx = -1; - if (ip.columns && ip.column_count > 0) { - for (uint16_t i = 0; i < ip.column_count; ++i) { - if (ip.columns[i] && ip.columns[i]->value().equals_ci(shard_key.ptr, shard_key.len)) { - shard_col_idx = static_cast(i); - break; - } - } - } else if (table) { - for (uint16_t i = 0; i < table->column_count; ++i) { - if (table->columns[i].name.equals_ci(shard_key.ptr, shard_key.len)) { - shard_col_idx = static_cast(i); - break; - } - } + std::vector key_ords(key_names.size(), -1); + for (size_t k = 0; k < key_names.size(); ++k) { + sql_parser::StringRef kn{key_names[k].c_str(), + static_cast(key_names[k].size())}; + key_ords[k] = find_insert_key_ordinal(ip.columns, ip.column_count, table, kn); } - - if (shard_col_idx < 0) { - return fail_dml("cannot route INSERT ... SELECT: shard key column is not present"); + for (int ord : key_ords) { + if (ord < 0) + return fail_dml("cannot route INSERT ... SELECT: shard key column is not present"); } const auto& shard_list = shards_.get_shards(table->table_name); std::unordered_map> shard_rows; for (size_t ri = 0; ri < rs.rows.size(); ++ri) { - if (shard_col_idx >= rs.rows[ri].column_count) { - return fail_dml("cannot route INSERT ... SELECT: missing shard key value"); + std::vector parts; + parts.reserve(key_ords.size()); + for (int ord : key_ords) { + if (ord >= rs.rows[ri].column_count) + return fail_dml("cannot route INSERT ... SELECT: missing shard key value"); + Value v = rs.rows[ri].get(static_cast(ord)); + if (!value_is_routable(v)) + return fail_dml("cannot route INSERT ... SELECT: shard key is not a literal"); + parts.push_back(value_to_part(v)); } - Value v = rs.rows[ri].get(static_cast(shard_col_idx)); size_t shard_idx = 0; - if (!route_value(table->table_name, v, shard_idx)) { + if (!shards_.try_shard_index_for_parts(table->table_name, parts.data(), + parts.size(), shard_idx)) { return fail_dml("cannot route INSERT ... SELECT: shard key is not a literal"); } shard_rows[shard_idx].push_back(ri); @@ -2008,17 +2404,13 @@ class DistributedPlanner { } const TableInfo* table = ctx.scan->scan.table; - if (!shards_.has_table(table->table_name) || !shards_.is_sharded(table->table_name)) { - if (shards_.has_table(table->table_name) && !shards_.is_sharded(table->table_name)) { - // Unsharded: push DISTINCT to remote - sql_parser::StringRef sql = qb_.build_select( - table, ctx.where_expr, proj_exprs, proj_count, - nullptr, 0, nullptr, nullptr, 0, -1, true); - return make_remote_scan(shards_.get_backend(table->table_name), sql, table); - } - PlanNode* result = make_plan_node(arena_, PlanNodeType::DISTINCT); - result->left = distribute_node(distinct_node->left); - return result; + if (!shards_.has_table(table->table_name)) + return fail_dml("table not in shard map"); + if (!shards_.is_sharded(table->table_name)) { + sql_parser::StringRef sql = qb_.build_select( + table, ctx.where_expr, proj_exprs, proj_count, + nullptr, 0, nullptr, nullptr, 0, -1, true); + return make_remote_scan(shards_.get_backend(table->table_name), sql, table); } // Sharded DISTINCT: each shard computes DISTINCT, local DISTINCT deduplicates diff --git a/include/sql_engine/operators/set_op_op.h b/include/sql_engine/operators/set_op_op.h index 441e4e6..ea7e7c1 100644 --- a/include/sql_engine/operators/set_op_op.h +++ b/include/sql_engine/operators/set_op_op.h @@ -17,86 +17,82 @@ class SetOpOperator : public Operator { public: SetOpOperator(Operator* left, Operator* right, uint8_t op, bool all, bool parallel_open = false, ThreadPool* pool = nullptr) - : left_(left), right_(right), op_(op), all_(all), + : op_(op), all_(all), parallel_open_(parallel_open), pool_(pool) { + children_.push_back(left); + children_.push_back(right); + } + + explicit SetOpOperator(std::vector children, + bool parallel_open = false, ThreadPool* pool = nullptr) + : children_(std::move(children)), op_(SET_OP_UNION), all_(true), parallel_open_(parallel_open), pool_(pool) {} void open() override { - if (parallel_open_ && pool_) { - // Thread-pool parallel open: ~1-2us dispatch vs ~200us for std::async - auto fl = pool_->submit([this]{ left_->open(); }); - auto fr = pool_->submit([this]{ right_->open(); }); - fl.get(); - fr.get(); - } else if (parallel_open_) { - // Fallback: std::async when no pool available - auto fl = std::async(std::launch::async, [this]{ left_->open(); }); - auto fr = std::async(std::launch::async, [this]{ right_->open(); }); - fl.get(); - fr.get(); + if (children_.empty()) return; + if (parallel_open_ && children_.size() > 1) { + std::vector> futures; + futures.reserve(children_.size()); + for (size_t i = 0; i < children_.size(); ++i) { + auto launcher = [this, i]{ children_[i]->open(); }; + if (pool_) { + futures.push_back(pool_->submit(std::move(launcher))); + } else { + futures.push_back(std::async(std::launch::async, std::move(launcher))); + } + } + for (auto& f : futures) f.get(); } else { - left_->open(); - right_->open(); + for (auto* c : children_) c->open(); } - reading_left_ = true; + child_idx_ = 0; seen_.clear(); expected_col_count_ = -1; - if (op_ == SET_OP_INTERSECT || op_ == SET_OP_EXCEPT) { - // Materialize right side into a set + if ((op_ == SET_OP_INTERSECT || op_ == SET_OP_EXCEPT) && children_.size() >= 2) { right_set_.clear(); Row r{}; - while (right_->next(r)) { + while (children_[1]->next(r)) { check_col_count(r); check_operator_row_limit(right_set_.size(), kDefaultMaxOperatorRows, "SetOpOperator"); right_set_.insert(row_key(r)); } - right_->close(); + children_[1]->close(); + right_closed_ = true; } } bool next(Row& out) override { + if (children_.empty()) return false; + if (op_ == SET_OP_UNION && !all_) { - // UNION (deduplicated) - while (true) { - bool got = false; - if (reading_left_) { - got = left_->next(out); - if (!got) { reading_left_ = false; } - } - if (!reading_left_) { - got = right_->next(out); - if (!got) return false; + while (child_idx_ < children_.size()) { + if (!children_[child_idx_]->next(out)) { + ++child_idx_; + continue; } - if (got) { - check_col_count(out); - std::string key = row_key(out); - if (seen_.find(key) == seen_.end()) { - check_operator_row_limit(seen_.size(), kDefaultMaxOperatorRows, "SetOpOperator"); - } - if (seen_.insert(key).second) return true; + check_col_count(out); + std::string key = row_key(out); + if (seen_.find(key) == seen_.end()) { + check_operator_row_limit(seen_.size(), kDefaultMaxOperatorRows, "SetOpOperator"); } + if (seen_.insert(key).second) return true; } + return false; } if (op_ == SET_OP_UNION && all_) { - // UNION ALL: yield left then right - if (reading_left_) { - if (left_->next(out)) { + while (child_idx_ < children_.size()) { + if (children_[child_idx_]->next(out)) { check_col_count(out); return true; } - reading_left_ = false; - } - if (right_->next(out)) { - check_col_count(out); - return true; + ++child_idx_; } return false; } if (op_ == SET_OP_INTERSECT) { - // Yield left rows that also appear in right - while (left_->next(out)) { + while (children_[0]->next(out)) { check_col_count(out); std::string key = row_key(out); if (right_set_.count(key)) { @@ -108,8 +104,7 @@ class SetOpOperator : public Operator { } if (op_ == SET_OP_EXCEPT) { - // Yield left rows that don't appear in right - while (left_->next(out)) { + while (children_[0]->next(out)) { check_col_count(out); std::string key = row_key(out); if (!right_set_.count(key)) { @@ -124,20 +119,22 @@ class SetOpOperator : public Operator { } void close() override { - left_->close(); - right_->close(); + for (size_t i = 0; i < children_.size(); ++i) { + if (right_closed_ && i == 1) continue; + children_[i]->close(); + } seen_.clear(); right_set_.clear(); } private: - Operator* left_; - Operator* right_; + std::vector children_; uint8_t op_; bool all_; bool parallel_open_; ThreadPool* pool_ = nullptr; - bool reading_left_ = true; + size_t child_idx_ = 0; + bool right_closed_ = false; std::unordered_set seen_; std::unordered_set right_set_; // Column count established by the first row we see. Used to detect diff --git a/include/sql_engine/plan_executor.h b/include/sql_engine/plan_executor.h index b6395b1..8bff423 100644 --- a/include/sql_engine/plan_executor.h +++ b/include/sql_engine/plan_executor.h @@ -1171,11 +1171,29 @@ class PlanExecutor { } Operator* build_set_op(PlanNode* node) { + std::vector remote_leaves; + if (node->set_op.op == SET_OP_UNION && node->set_op.all && + collect_union_all_remotes(node, remote_leaves) && + remote_leaves.size() > 2) { + std::vector children; + children.reserve(remote_leaves.size()); + for (PlanNode* leaf : remote_leaves) { + Operator* child = build_remote_scan(leaf); + if (!child) return nullptr; + children.push_back(child); + } + bool parallel = parallel_open_enabled_ && children.size() > 1; + auto op = std::make_unique( + std::move(children), parallel, parallel ? pool_ : nullptr); + Operator* ptr = op.get(); + operators_.push_back(std::move(op)); + return ptr; + } + Operator* left = build_operator(node->left); Operator* right = build_operator(node->right); if (!left || !right) return nullptr; - // Enable parallel open when both children are remote scans and executor is thread-safe bool parallel = parallel_open_enabled_ && (node->left && node->left->type == PlanNodeType::REMOTE_SCAN && node->right && node->right->type == PlanNodeType::REMOTE_SCAN); @@ -1187,6 +1205,21 @@ class PlanExecutor { return ptr; } + static bool collect_union_all_remotes(const PlanNode* node, + std::vector& out) { + if (!node) return false; + if (node->type == PlanNodeType::REMOTE_SCAN) { + out.push_back(const_cast(node)); + return true; + } + if (node->type == PlanNodeType::SET_OP && + node->set_op.op == SET_OP_UNION && node->set_op.all) { + return collect_union_all_remotes(node->left, out) && + collect_union_all_remotes(node->right, out); + } + return false; + } + Operator* build_remote_scan(PlanNode* node) { if (!remote_executor_) return nullptr; sql_parser::StringRef sql{node->remote_scan.remote_sql, diff --git a/include/sql_engine/session.h b/include/sql_engine/session.h index a2dd90d..26d2494 100644 --- a/include/sql_engine/session.h +++ b/include/sql_engine/session.h @@ -126,15 +126,16 @@ class Session { std::string sql_key(sql, len); auto cache_it = plan_cache_.find(sql_key); if (cache_it != plan_cache_.end()) { - // Cache hit: move this entry to the front of the LRU list. plan_cache_order_.splice(plan_cache_order_.begin(), plan_cache_order_, cache_it->second); exec_arena_.reset(); auto& entry = *cache_it->second; + PlanNode* plan = maybe_distribute(entry.plan, exec_arena_); + if (!plan) return {}; PlanExecutor executor(functions_, catalog_, exec_arena_); wire_executor(executor); - return executor.execute(entry.plan); + return executor.execute(plan); } // Cache miss: full parse -> plan -> optimize -> distribute pipeline @@ -164,25 +165,21 @@ class Session { plan = optimizer_.optimize(plan, cached_parser->arena()); - // Distribute across shards if shard map is configured + ResultSet rs; if (shard_map_ && remote_executor_) { - DistributedPlanner dplanner(*shard_map_, catalog_, cached_parser->arena(), remote_executor_, &functions_); - plan = dplanner.distribute(plan); + exec_arena_.reset(); + PlanNode* dist = maybe_distribute(plan, exec_arena_); + if (!dist) return {}; + PlanExecutor executor(functions_, catalog_, exec_arena_); + wire_executor(executor); + rs = executor.execute(dist); + } else { + PlanExecutor executor(functions_, catalog_, cached_parser->arena()); + wire_executor(executor); + rs = executor.execute(plan); } - // Execute first using the parser's arena. preprocess_aggregates (called - // inside execute()) may allocate into the arena to modify the plan in-place. - // Those allocations must persist in the parser arena (not exec_arena_) so - // the cached plan remains valid across calls. - PlanExecutor executor(functions_, catalog_, cached_parser->arena()); - wire_executor(executor); - ResultSet rs = executor.execute(plan); - - // Cache the plan, enforcing the LRU bound. The parser arena is kept - // alive via unique_ptr stored in the CachedPlan so all plan/AST - // string pointers remain valid for subsequent cache hits. insert_into_plan_cache(std::move(sql_key), std::move(cached_parser), plan); - return rs; } @@ -378,6 +375,15 @@ class Session { std::unordered_map plan_cache_; size_t plan_cache_max_size_ = 1024; + PlanNode* maybe_distribute(PlanNode* plan, sql_parser::Arena& arena) { + if (!plan || !shard_map_ || !remote_executor_) return plan; + DistributedPlanner dplanner(*shard_map_, catalog_, arena, + remote_executor_, &functions_); + PlanNode* dist = dplanner.distribute(plan); + if (dplanner.last_error()) return nullptr; + return dist; + } + void insert_into_plan_cache(std::string key, std::unique_ptr> parser, PlanNode* plan) { diff --git a/include/sql_engine/shard_map.h b/include/sql_engine/shard_map.h index cbb6061..907a651 100644 --- a/include/sql_engine/shard_map.h +++ b/include/sql_engine/shard_map.h @@ -51,19 +51,34 @@ struct ShardListEntry { size_t shard_index = 0; }; +// One component of a (possibly composite) shard key. n==1 reuses the +// single-column HASH/RANGE/LIST path so existing routes stay identical. +struct ShardKeyPart { + bool is_int = true; + int64_t int_val = 0; + const char* str = nullptr; + uint32_t str_len = 0; +}; + struct TableShardConfig { std::string table_name; - std::string shard_key; // empty if unsharded + std::string shard_key; // empty if unsharded; "a+b" for composite std::vector shards; // 1 if unsharded, N if sharded RoutingStrategy strategy = RoutingStrategy::HASH; std::vector ranges; // used iff strategy == RANGE std::vector list; // used iff strategy == LIST + std::vector shard_keys; // filled by ShardMap::add_table }; class ShardMap { public: void add_table(const TableShardConfig& config) { TableShardConfig copy = config; + if (copy.shard_keys.empty() && !copy.shard_key.empty()) { + split_keys(copy.shard_key, copy.shard_keys); + } else if (copy.shard_key.empty() && !copy.shard_keys.empty()) { + copy.shard_key = join_keys(copy.shard_keys); + } if (copy.strategy == RoutingStrategy::RANGE) { std::sort(copy.ranges.begin(), copy.ranges.end(), [](const ShardRange& a, const ShardRange& b) { @@ -77,7 +92,7 @@ class ShardMap { bool is_sharded(sql_parser::StringRef table_name) const { const TableShardConfig* cfg = lookup(table_name); if (!cfg) return false; - return !cfg->shard_key.empty() && cfg->shards.size() > 1; + return !cfg->shard_keys.empty() && cfg->shards.size() > 1; } const std::vector& get_shards(sql_parser::StringRef table_name) const { @@ -89,20 +104,109 @@ class ShardMap { sql_parser::StringRef get_shard_key(sql_parser::StringRef table_name) const { const TableShardConfig* cfg = lookup(table_name); - if (cfg && !cfg->shard_key.empty()) { - const std::string& sk = cfg->shard_key; + if (cfg && !cfg->shard_keys.empty()) { + const std::string& sk = cfg->shard_keys[0]; return sql_parser::StringRef{sk.c_str(), static_cast(sk.size())}; } return sql_parser::StringRef{nullptr, 0}; } + const std::vector& get_shard_keys(sql_parser::StringRef table_name) const { + const TableShardConfig* cfg = lookup(table_name); + if (cfg) return cfg->shard_keys; + static const std::vector empty; + return empty; + } + bool has_table(sql_parser::StringRef table_name) const { return lookup(table_name) != nullptr; } + // False if the table is unknown, has no shards, or a LIST value is unmapped. + bool try_shard_index_for_int(sql_parser::StringRef table_name, int64_t value, + size_t& out) const { + const TableShardConfig* cfg = lookup(table_name); + if (!cfg || cfg->shards.empty()) return false; + size_t n = cfg->shards.size(); + switch (cfg->strategy) { + case RoutingStrategy::HASH: + out = fnv1a_int64(value) % n; + return true; + case RoutingStrategy::RANGE: + out = shard_index_for_int(table_name, value); + return true; + case RoutingStrategy::LIST: + for (const auto& e : cfg->list) { + if (e.is_int && e.int_val == value) { + out = clamp_index(e.shard_index, n); + return true; + } + } + return false; + } + return false; + } + + bool try_shard_index_for_string(sql_parser::StringRef table_name, + const char* val, uint32_t val_len, + size_t& out) const { + const TableShardConfig* cfg = lookup(table_name); + if (!cfg || cfg->shards.empty()) return false; + size_t n = cfg->shards.size(); + switch (cfg->strategy) { + case RoutingStrategy::HASH: + out = fnv1a_bytes(reinterpret_cast(val), val_len) % n; + return true; + case RoutingStrategy::RANGE: + return false; + case RoutingStrategy::LIST: + for (const auto& e : cfg->list) { + if (!e.is_int && e.str_val.size() == val_len && + std::memcmp(e.str_val.data(), val, val_len) == 0) { + out = clamp_index(e.shard_index, n); + return true; + } + } + return false; + } + return false; + } + + bool try_shard_index_for_parts(sql_parser::StringRef table_name, + const ShardKeyPart* parts, size_t n, + size_t& out) const { + const TableShardConfig* cfg = lookup(table_name); + if (!cfg || cfg->shards.empty() || !parts || n == 0) return false; + if (n == 1) { + return parts[0].is_int + ? try_shard_index_for_int(table_name, parts[0].int_val, out) + : try_shard_index_for_string(table_name, parts[0].str, + parts[0].str_len, out); + } + if (cfg->strategy != RoutingStrategy::HASH) return false; + uint64_t h = 0xcbf29ce484222325ULL; + for (size_t i = 0; i < n; ++i) { + if (parts[i].is_int) { + uint64_t u = static_cast(parts[i].int_val); + uint8_t bytes[8]; + for (int b = 0; b < 8; ++b) + bytes[b] = static_cast((u >> (b * 8)) & 0xff); + h = fnv1a_mix(h, bytes, 8); + } else { + h = fnv1a_mix(h, + reinterpret_cast(parts[i].str), + parts[i].str_len); + } + uint8_t sep = 0xff; + h = fnv1a_mix(h, &sep, 1); + } + out = static_cast(h % cfg->shards.size()); + return true; + } + // Determine which shard index a value maps to. Dispatches on the // configured RoutingStrategy. Returns 0 if the table is unknown or - // has no shards. + // has no shards — prefer try_shard_index_for_* which fails closed. size_t shard_index_for_int(sql_parser::StringRef table_name, int64_t value) const { const TableShardConfig* cfg = lookup(table_name); if (!cfg || cfg->shards.empty()) return 0; @@ -161,6 +265,21 @@ class ShardMap { return cfg ? cfg->strategy : RoutingStrategy::HASH; } + // LIST only. Inclusive [lo, hi]. Pushes every mapped int in the window. + void collect_int_list_shards(sql_parser::StringRef table_name, + int64_t lo, int64_t hi, + std::vector& out) const { + const TableShardConfig* cfg = lookup(table_name); + if (!cfg || cfg->strategy != RoutingStrategy::LIST || cfg->list.empty()) + return; + if (lo > hi) return; + size_t n = cfg->shards.size(); + for (const auto& e : cfg->list) { + if (e.is_int && e.int_val >= lo && e.int_val <= hi) + out.push_back(clamp_index(e.shard_index, n)); + } + } + // RANGE only. Inclusive [lo, hi]. HASH/LIST yield no indices (caller scatters). void collect_int_range_shards(sql_parser::StringRef table_name, int64_t lo, int64_t hi, @@ -238,10 +357,34 @@ class ShardMap { return idx < n ? idx : (n == 0 ? 0 : n - 1); } + static void split_keys(const std::string& spec, std::vector& out) { + size_t start = 0; + while (start <= spec.size()) { + size_t plus = spec.find('+', start); + if (plus == std::string::npos) { + out.push_back(spec.substr(start)); + break; + } + out.push_back(spec.substr(start, plus - start)); + start = plus + 1; + } + out.erase(std::remove_if(out.begin(), out.end(), + [](const std::string& s) { return s.empty(); }), + out.end()); + } + + static std::string join_keys(const std::vector& keys) { + std::string out; + for (size_t i = 0; i < keys.size(); ++i) { + if (i) out += '+'; + out += keys[i]; + } + return out; + } + // FNV-1a 64-bit. Deterministic across compilers, fast, good enough for // shard-key routing where adversarial inputs are not a concern. - static uint64_t fnv1a_bytes(const uint8_t* data, size_t len) { - uint64_t h = 0xcbf29ce484222325ULL; + static uint64_t fnv1a_mix(uint64_t h, const uint8_t* data, size_t len) { for (size_t i = 0; i < len; ++i) { h ^= static_cast(data[i]); h *= 0x100000001b3ULL; @@ -249,6 +392,10 @@ class ShardMap { return h; } + static uint64_t fnv1a_bytes(const uint8_t* data, size_t len) { + return fnv1a_mix(0xcbf29ce484222325ULL, data, len); + } + static uint64_t fnv1a_int64(int64_t v) { uint64_t u = static_cast(v); uint8_t bytes[8]; diff --git a/src/sql_engine/tool_config_parser.cpp b/src/sql_engine/tool_config_parser.cpp index 04fd35b..e9c4738 100644 --- a/src/sql_engine/tool_config_parser.cpp +++ b/src/sql_engine/tool_config_parser.cpp @@ -172,6 +172,10 @@ ParsedShard parse_shard_spec(const std::string& spec) { ps.config.table_name = spec.substr(0, c1); ps.config.shard_key = spec.substr(c1 + 1, c2 - c1 - 1); + if (ps.config.shard_key.empty()) { + ps.error = "Empty shard key in: " + spec; + return ps; + } // Look for an optional strategy qualifier in the next colon-separated // token: hash | range | list. If absent, default to HASH. @@ -204,6 +208,10 @@ ParsedShard parse_shard_spec(const std::string& spec) { return ps; } } else if (strategy_token == "range") { + if (ps.config.shard_key.find('+') != std::string::npos) { + ps.error = "composite shard keys require HASH strategy: " + spec; + return ps; + } ps.config.strategy = RoutingStrategy::RANGE; for (auto& entry : split_csv(body)) { std::string upper_str, backend; @@ -226,6 +234,10 @@ ParsedShard parse_shard_spec(const std::string& spec) { return ps; } } else if (strategy_token == "list") { + if (ps.config.shard_key.find('+') != std::string::npos) { + ps.error = "composite shard keys require HASH strategy: " + spec; + return ps; + } ps.config.strategy = RoutingStrategy::LIST; for (auto& entry : split_csv(body)) { std::string val_str, backend; diff --git a/tests/test_distributed_dml.cpp b/tests/test_distributed_dml.cpp index d9d32d6..513afee 100644 --- a/tests/test_distributed_dml.cpp +++ b/tests/test_distributed_dml.cpp @@ -12,6 +12,8 @@ #include "sql_engine/in_memory_catalog.h" #include "sql_engine/data_source.h" #include "sql_engine/function_registry.h" +#include "sql_engine/session.h" +#include "sql_engine/local_txn.h" #include "sql_parser/parser.h" #include #include @@ -305,7 +307,8 @@ class DistributedDmlTest : public ::testing::Test { return r; } - DistributedPlanner dist(shard_map, catalog, parser.arena()); + DistributedPlanner dist(shard_map, catalog, parser.arena(), + &mock_executor, &functions); PlanNode* dist_plan = dist.distribute_dml(plan); if (dist.last_error()) { DmlResult r; @@ -371,8 +374,10 @@ class DistributedDmlTest : public ::testing::Test { PlanNode* plan = builder.build(pr.ast); if (!plan) return {}; - DistributedPlanner dist(shard_map, catalog, p.arena()); + DistributedPlanner dist(shard_map, catalog, p.arena(), + &mock_executor, &functions); PlanNode* dist_plan = dist.distribute(plan); + if (dist.last_error()) return {}; PlanExecutor executor(functions, catalog, p.arena()); executor.set_remote_executor(&mock_executor); @@ -708,11 +713,88 @@ TEST_F(DistributedDmlTest, InsertNonLiteralShardKeyErrors) { EXPECT_EQ(mock_executor.total_row_count("users"), 0u); } -TEST_F(DistributedDmlTest, UpdateShardKeyErrors) { +TEST_F(DistributedDmlTest, UpdateShardKeyMovesRow) { execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (3, 'Carol', 17)"); + const char* src = backend_for_id(3); + const char* dst = backend_for_id(9); + ASSERT_STRNE(src, dst); + auto result = execute_distributed_dml("UPDATE users SET id = 9 WHERE id = 3"); + EXPECT_TRUE(result.success) << result.error_message; + EXPECT_EQ(row_count_on(src, "users"), 0u); + EXPECT_EQ(row_count_on(dst, "users"), 1u); + + auto got = execute_distributed_select("SELECT name FROM users WHERE id = 9"); + ASSERT_EQ(got.row_count(), 1u); + EXPECT_EQ(std::string(got.rows[0].get(0).str_val.ptr, got.rows[0].get(0).str_val.len), + "Carol"); +} + +TEST_F(DistributedDmlTest, UpdateShardKeySameShard) { + int64_t a = 3; + int64_t b = a; + for (int64_t i = 4; i < 200; ++i) { + if (shard_for_id(i) == shard_for_id(a)) { b = i; break; } + } + ASSERT_NE(a, b); + + execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (3, 'Carol', 17)"); + auto result = execute_distributed_dml( + ("UPDATE users SET id = " + std::to_string(b) + " WHERE id = 3").c_str()); + EXPECT_TRUE(result.success) << result.error_message; + EXPECT_EQ(row_count_on(backend_for_id(a), "users"), 1u); + auto got = execute_distributed_select( + ("SELECT id FROM users WHERE id = " + std::to_string(b)).c_str()); + ASSERT_EQ(got.row_count(), 1u); +} + +TEST_F(DistributedDmlTest, UnknownTableDmlErrors) { + catalog.add_table("", "ghost", { + {"id", SqlType::make_int(), false}, + }); + auto result = execute_distributed_dml("INSERT INTO ghost (id) VALUES (1)"); EXPECT_FALSE(result.success); - EXPECT_NE(result.error_message.find("shard key"), std::string::npos); + EXPECT_NE(result.error_message.find("shard map"), std::string::npos); +} + +TEST_F(DistributedDmlTest, CompositeShardKeyInsertSelect) { + catalog.add_table("", "kv", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(255), true}, + }); + TableShardConfig cfg; + cfg.table_name = "kv"; + cfg.shard_key = "tenant_id+id"; + cfg.shards = {{"shard0"}, {"shard1"}, {"shard2"}}; + shard_map.add_table(cfg); + mock_executor.add_table_to_all("kv", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(255), true}, + }); + + auto missing = execute_distributed_dml( + "INSERT INTO kv (id, name) VALUES (1, 'x')"); + EXPECT_FALSE(missing.success); + EXPECT_NE(missing.error_message.find("shard key"), std::string::npos); + + EXPECT_TRUE(execute_distributed_dml( + "INSERT INTO kv (tenant_id, id, name) VALUES (1, 3, 'Alice')").success); + EXPECT_TRUE(execute_distributed_dml( + "INSERT INTO kv (tenant_id, id, name) VALUES (2, 3, 'Bob')").success); + + auto alice = execute_distributed_select( + "SELECT name FROM kv WHERE tenant_id = 1 AND id = 3"); + ASSERT_EQ(alice.row_count(), 1u); + EXPECT_EQ(std::string(alice.rows[0].get(0).str_val.ptr, + alice.rows[0].get(0).str_val.len), "Alice"); + + auto bob = execute_distributed_select( + "SELECT name FROM kv WHERE tenant_id = 2 AND id = 3"); + ASSERT_EQ(bob.row_count(), 1u); + EXPECT_EQ(std::string(bob.rows[0].get(0).str_val.ptr, + bob.rows[0].get(0).str_val.len), "Bob"); } TEST_F(DistributedDmlTest, InsertThenPointSelectRange) { @@ -784,3 +866,93 @@ TEST_F(DistributedDmlTest, MultiTableDeleteOnShardedFails) { EXPECT_FALSE(result.success); EXPECT_NE(result.error_message.find("sharded"), std::string::npos); } + +TEST_F(DistributedDmlTest, InsertUnmappedListValueErrors) { + TableShardConfig cfg; + cfg.table_name = "users"; + cfg.shard_key = "id"; + cfg.shards = {{"shard0"}, {"shard1"}, {"shard2"}}; + cfg.strategy = RoutingStrategy::LIST; + cfg.list = { + {true, 3, "", 0}, + {true, 7, "", 1}, + }; + shard_map.add_table(cfg); + + auto result = execute_distributed_dml( + "INSERT INTO users (id, name, age) VALUES (99, 'Nope', 1)"); + EXPECT_FALSE(result.success); + EXPECT_NE(result.error_message.find("mapped"), std::string::npos); + EXPECT_EQ(mock_executor.total_row_count("users"), 0u); +} + +TEST_F(DistributedDmlTest, InsertThenListBetweenSelect) { + TableShardConfig cfg; + cfg.table_name = "users"; + cfg.shard_key = "id"; + cfg.shards = {{"shard0"}, {"shard1"}, {"shard2"}}; + cfg.strategy = RoutingStrategy::LIST; + cfg.list = { + {true, 3, "", 0}, + {true, 7, "", 1}, + {true, 20, "", 2}, + }; + shard_map.add_table(cfg); + + EXPECT_TRUE(execute_distributed_dml( + "INSERT INTO users (id, name, age) VALUES (3, 'A', 1)").success); + EXPECT_TRUE(execute_distributed_dml( + "INSERT INTO users (id, name, age) VALUES (7, 'B', 2)").success); + EXPECT_TRUE(execute_distributed_dml( + "INSERT INTO users (id, name, age) VALUES (20, 'C', 3)").success); + + auto mid = execute_distributed_select("SELECT name FROM users WHERE id BETWEEN 6 AND 8"); + ASSERT_EQ(mid.row_count(), 1u); + EXPECT_EQ(std::string(mid.rows[0].get(0).str_val.ptr, mid.rows[0].get(0).str_val.len), "B"); +} + +TEST_F(DistributedDmlTest, PlanCacheRedistributesSelect) { + execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (3, 'Carol', 17)"); + LocalTransactionManager txn(data_arena); + Session session(catalog, txn); + session.set_remote_executor(&mock_executor); + session.set_shard_map(&shard_map); + + auto first = session.execute_query("SELECT name FROM users WHERE id = 3"); + ASSERT_EQ(first.row_count(), 1u); + EXPECT_EQ(session.plan_cache_size(), 1u); + + auto second = session.execute_query("SELECT name FROM users WHERE id = 3"); + ASSERT_EQ(second.row_count(), 1u); + EXPECT_EQ(session.plan_cache_size(), 1u); + EXPECT_EQ(std::string(second.rows[0].get(0).str_val.ptr, + second.rows[0].get(0).str_val.len), "Carol"); +} + +TEST_F(DistributedDmlTest, PlanCacheSeesUpdatedShardMap) { + execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (3, 'Carol', 17)"); + std::string home = backend_for_id(3); + LocalTransactionManager txn(data_arena); + Session session(catalog, txn); + session.set_remote_executor(&mock_executor); + session.set_shard_map(&shard_map); + + const char* sql = "SELECT name FROM users WHERE id = 3"; + ASSERT_EQ(session.execute_query(sql).row_count(), 1u); + EXPECT_EQ(session.plan_cache_size(), 1u); + + TableShardConfig unsharded; + unsharded.table_name = "users"; + unsharded.shard_key = ""; + unsharded.shards = {{home}}; + shard_map.add_table(unsharded); + mock_executor.clear_sql_logs(); + + auto rs = session.execute_query(sql); + ASSERT_EQ(rs.row_count(), 1u); + EXPECT_EQ(mock_executor.get_executed_sqls(home).size(), 1u); + for (const char* s : {"shard0", "shard1", "shard2"}) { + if (s != home) + EXPECT_EQ(mock_executor.get_executed_sqls(s).size(), 0u) << s; + } +} diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index 62000b0..ca52c72 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -874,6 +874,40 @@ TEST_F(DistributedPlannerTest, ShardRouting_RangeInequalityPrunes) { EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id >= 1 AND id <= 15"), 3u); } +TEST_F(DistributedPlannerTest, ShardRouting_ListBetweenPrunes) { + TableShardConfig cfg; + cfg.table_name = "users"; + cfg.shard_key = "id"; + cfg.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + cfg.strategy = RoutingStrategy::LIST; + cfg.list = { + {true, 1, "", 0}, + {true, 6, "", 1}, + {true, 7, "", 1}, + {true, 15, "", 2}, + }; + shard_map.add_table(cfg); + + auto count_remotes = [&](const char* sql) { + Parser parser; + auto pr = parser.parse(sql, std::strlen(sql)); + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + return remotes.size(); + }; + + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id BETWEEN 6 AND 7"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id BETWEEN 1 AND 15"), 3u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id = 99"), 3u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id IN (6, 99)"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id IN (6, 15)"), 2u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id BETWEEN 2 AND 5"), 3u); +} + TEST_F(DistributedPlannerTest, ShardRouting_Correctness) { // The shard routing uses FNV-1a 64-bit hash of the key % num_shards // to pick a shard, but the test fixture uses sequential partitioning @@ -1165,3 +1199,111 @@ TEST_F(DistributedPlannerTest, ColocatedJoinPushedToShards) { EXPECT_NE(remote.find("JOIN"), std::string::npos) << remote; } } + +TEST_F(DistributedPlannerTest, CompositeColocatedJoinPushedToShards) { + catalog.add_table("", "kv", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(255), true}, + }); + catalog.add_table("", "kv_orders", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"amt", SqlType::make_int(), true}, + }); + TableShardConfig kv; + kv.table_name = "kv"; + kv.shard_key = "tenant_id+id"; + kv.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + shard_map.add_table(kv); + TableShardConfig kvo; + kvo.table_name = "kv_orders"; + kvo.shard_key = "tenant_id+id"; + kvo.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + shard_map.add_table(kvo); + + Parser parser; + const char* sql = + "SELECT * FROM kv JOIN kv_orders ON kv.tenant_id = kv_orders.tenant_id " + "AND kv.id = kv_orders.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()) << "composite co-located join should not stay local"; + ASSERT_EQ(remotes.size(), 3u); + 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; + } +} + +TEST_F(DistributedPlannerTest, IncompleteCompositeJoinStaysLocal) { + catalog.add_table("", "kv", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(255), true}, + }); + catalog.add_table("", "kv_orders", { + {"tenant_id", SqlType::make_int(), false}, + {"id", SqlType::make_int(), false}, + {"amt", SqlType::make_int(), true}, + }); + TableShardConfig kv; + kv.table_name = "kv"; + kv.shard_key = "tenant_id+id"; + kv.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + shard_map.add_table(kv); + TableShardConfig kvo; + kvo.table_name = "kv_orders"; + kvo.shard_key = "tenant_id+id"; + kvo.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + shard_map.add_table(kvo); + + Parser parser; + const char* sql = "SELECT * FROM kv JOIN kv_orders ON kv.tenant_id = kv_orders.tenant_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; + find_nodes(dist, PlanNodeType::JOIN, joins); + EXPECT_FALSE(joins.empty()) << "join on a partial composite key must gather"; +} + +TEST_F(DistributedPlannerTest, UnknownTableErrors) { + catalog.add_table("", "ghost", { + {"id", SqlType::make_int(), false}, + }); + Parser parser; + const char* sql = "SELECT * FROM ghost"; + 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); + EXPECT_EQ(dist, nullptr); + ASSERT_NE(dp.last_error(), nullptr); + EXPECT_NE(std::string(dp.last_error()).find("shard map"), std::string::npos); +} diff --git a/tests/test_operators.cpp b/tests/test_operators.cpp index 8f77f43..0a5a476 100644 --- a/tests/test_operators.cpp +++ b/tests/test_operators.cpp @@ -1001,6 +1001,26 @@ TEST_F(SetOpOpTest, UnionAll) { setop.close(); } +TEST_F(SetOpOpTest, UnionAllNary) { + std::vector a = {build_row(arena, {value_int(1)})}; + std::vector b = {build_row(arena, {value_int(2)})}; + std::vector c = {build_row(arena, {value_int(3)})}; + InMemoryDataSource ds_a(table, a); + InMemoryDataSource ds_b(table, b); + InMemoryDataSource ds_c(table, c); + ScanOperator scan_a(&ds_a); + ScanOperator scan_b(&ds_b); + ScanOperator scan_c(&ds_c); + + SetOpOperator setop({&scan_a, &scan_b, &scan_c}); + setop.open(); + Row out{}; + int count = 0; + while (setop.next(out)) count++; + EXPECT_EQ(count, 3); + setop.close(); +} + // Regression test for silent column-count mismatch in set operations. // Before the validation fix, SetOpOperator::row_key would iterate each row's // own column_count, producing truncated/misaligned keys and silently-wrong diff --git a/tests/test_shard_map.cpp b/tests/test_shard_map.cpp index ae7cf64..31dd0a9 100644 --- a/tests/test_shard_map.cpp +++ b/tests/test_shard_map.cpp @@ -179,13 +179,36 @@ TEST(ShardMapListTest, IntKeysRouteByExplicitMap) { EXPECT_EQ(map.shard_index_for_int(sref("users"), 7), 1u); } -TEST(ShardMapListTest, MissingKeyRoutesToShardZero) { +TEST(ShardMapListTest, MissingKeyIsUnroutable) { TableShardConfig cfg = make_two_shards(RoutingStrategy::LIST); cfg.list = {ShardListEntry{true, 1, "", 0}}; ShardMap map; map.add_table(cfg); - EXPECT_EQ(map.shard_index_for_int(sref("users"), 999), 0u); + size_t idx = 99; + EXPECT_FALSE(map.try_shard_index_for_int(sref("users"), 999, idx)); +} + +TEST(ShardMapListTest, BetweenCollectsMappedInts) { + TableShardConfig cfg = make_two_shards(RoutingStrategy::LIST); + cfg.list = { + ShardListEntry{true, 1, "", 0}, + ShardListEntry{true, 6, "", 1}, + ShardListEntry{true, 7, "", 1}, + ShardListEntry{true, 20, "", 0}, + }; + ShardMap map; + map.add_table(cfg); + + std::vector mid; + map.collect_int_list_shards(sref("users"), 5, 10, mid); + ASSERT_EQ(mid.size(), 2u); + EXPECT_EQ(mid[0], 1u); + EXPECT_EQ(mid[1], 1u); + + std::vector none; + map.collect_int_list_shards(sref("users"), 2, 5, none); + EXPECT_TRUE(none.empty()); } TEST(ShardMapListTest, StringKeysRouteByExplicitMap) { @@ -199,7 +222,8 @@ TEST(ShardMapListTest, StringKeysRouteByExplicitMap) { EXPECT_EQ(map.shard_index_for_string(sref("users"), "us-east", 7), 0u); EXPECT_EQ(map.shard_index_for_string(sref("users"), "us-west", 7), 1u); - EXPECT_EQ(map.shard_index_for_string(sref("users"), "eu-north", 8), 0u); + size_t miss = 99; + EXPECT_FALSE(map.try_shard_index_for_string(sref("users"), "eu-north", 8, miss)); } // ---------------------------------------------------------------------- @@ -225,8 +249,32 @@ TEST(ShardMapTest, ClampOnOutOfRangeShardIndex) { EXPECT_EQ(map.shard_index_for_int(sref("users"), 50), 1u); } -TEST(ShardMapTest, UnknownTableReturnsZero) { - ShardMap map; // empty - EXPECT_EQ(map.shard_index_for_int(sref("nope"), 5), 0u); - EXPECT_EQ(map.shard_index_for_string(sref("nope"), "x", 1), 0u); +TEST(ShardMapTest, UnknownTableIsUnroutable) { + ShardMap map; + size_t idx = 99; + EXPECT_FALSE(map.try_shard_index_for_int(sref("nope"), 5, idx)); + EXPECT_FALSE(map.try_shard_index_for_string(sref("nope"), "x", 1, idx)); +} + +TEST(ShardMapTest, CompositeHashIsDeterministicAndDiffersFromSingle) { + TableShardConfig cfg; + cfg.table_name = "kv"; + cfg.shard_key = "tenant_id+id"; + cfg.shards = {ShardInfo{"s0"}, ShardInfo{"s1"}, ShardInfo{"s2"}}; + ShardMap map; + map.add_table(cfg); + + EXPECT_EQ(map.get_shard_keys(sref("kv")).size(), 2u); + EXPECT_EQ(std::string(map.get_shard_key(sref("kv")).ptr, + map.get_shard_key(sref("kv")).len), "tenant_id"); + + ShardKeyPart a[] = {{true, 1, nullptr, 0}, {true, 3, nullptr, 0}}; + ShardKeyPart b[] = {{true, 2, nullptr, 0}, {true, 3, nullptr, 0}}; + size_t ia = 0, ib = 0, ia2 = 0; + ASSERT_TRUE(map.try_shard_index_for_parts(sref("kv"), a, 2, ia)); + ASSERT_TRUE(map.try_shard_index_for_parts(sref("kv"), b, 2, ib)); + ASSERT_TRUE(map.try_shard_index_for_parts(sref("kv"), a, 2, ia2)); + EXPECT_EQ(ia, ia2); + EXPECT_LT(ia, 3u); + EXPECT_LT(ib, 3u); } diff --git a/tests/test_ssl_config.cpp b/tests/test_ssl_config.cpp index 0395f7c..b60a56d 100644 --- a/tests/test_ssl_config.cpp +++ b/tests/test_ssl_config.cpp @@ -226,3 +226,19 @@ TEST(SSLConfigTest, ParseShardSpecRejectsEmptyList) { EXPECT_FALSE(ps.ok); EXPECT_NE(ps.error.find("at least one"), std::string::npos); } + +TEST(SSLConfigTest, ParseShardSpecCompositeHash) { + auto ps = parse_shard_spec("kv:tenant_id+id:hash:s0,s1,s2"); + + ASSERT_TRUE(ps.ok); + EXPECT_EQ(ps.config.table_name, "kv"); + EXPECT_EQ(ps.config.shard_key, "tenant_id+id"); + EXPECT_EQ(ps.config.strategy, RoutingStrategy::HASH); + ASSERT_EQ(ps.config.shards.size(), 3u); +} + +TEST(SSLConfigTest, ParseShardSpecRejectsCompositeRange) { + auto ps = parse_shard_spec("kv:tenant_id+id:range:5=s0,10=s1"); + EXPECT_FALSE(ps.ok); + EXPECT_NE(ps.error.find("HASH"), std::string::npos); +} From 7b1726be34bdc2458f74a0b5676ac34bf715cd52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9=20Canna=C3=B2?= Date: Thu, 20 Aug 2026 14:34:32 +0700 Subject: [PATCH 3/3] feat: composite RANGE, semi-join prune, PG LIST/RANGE + 2PC demo - Composite RANGE now routes and prunes on the first key component (LIST still rejects composites; HASH uses all parts). - Semi-join prune: when one side of a join is sharded and the other is not, the engine materializes the small side and pushes an IN-list onto the sharded probe side. - New demo scripts for PostgreSQL shards (16432/16433) exercising RANGE + LIST + cross-shard 2PC with DistributedTransactionManager in POSTGRESQL mode. - sqlengine now auto-selects PostgreSQL 2PC dialect when all backends are pgsql://. - Added planner and live-backend tests; full suite now 1344 passed. Stacked on PR #61. --- include/sql_engine/distributed_planner.h | 162 ++++++++++++++++++++++- include/sql_engine/shard_map.h | 2 +- scripts/run_pg_sharding_demo.sh | 81 ++++++++++++ scripts/start_pg_sharding_demo.sh | 98 ++++++++++++++ src/sql_engine/tool_config_parser.cpp | 4 - tests/test_distributed_planner.cpp | 51 +++++++ tests/test_distributed_real.cpp | 89 +++++++++++++ tests/test_shard_map.cpp | 19 +++ tests/test_ssl_config.cpp | 8 +- tools/sqlengine.cpp | 8 +- 10 files changed, 512 insertions(+), 10 deletions(-) create mode 100755 scripts/run_pg_sharding_demo.sh create mode 100755 scripts/start_pg_sharding_demo.sh diff --git a/include/sql_engine/distributed_planner.h b/include/sql_engine/distributed_planner.h index 1a8491f..d224065 100644 --- a/include/sql_engine/distributed_planner.h +++ b/include/sql_engine/distributed_planner.h @@ -357,7 +357,13 @@ class DistributedPlanner { if (keys.empty()) return all_shards; std::vector target_indices; - if (keys.size() > 1) { + if (keys.size() > 1 && + shards_.routing_strategy(table->table_name) == RoutingStrategy::RANGE) { + sql_parser::StringRef first{keys[0].c_str(), + static_cast(keys[0].size())}; + extract_shard_targets(where_expr, first, table->table_name, + all_shards.size(), target_indices); + } else if (keys.size() > 1) { extract_composite_targets(where_expr, keys, table->table_name, target_indices); } else { sql_parser::StringRef shard_key{keys[0].c_str(), @@ -1203,6 +1209,157 @@ class DistributedPlanner { return current ? current : join_node; } + bool column_on_table(const sql_parser::AstNode* node, const TableInfo* table, + const TableInfo* other) const { + if (!node || !table) return false; + if (node->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) { + const sql_parser::AstNode* t = node->first_child; + if (!t) return false; + sql_parser::StringRef tn = t->value(); + if (table->table_name.equals_ci(tn.ptr, tn.len)) return true; + if (table->alias.ptr && table->alias.equals_ci(tn.ptr, tn.len)) return true; + return false; + } + if (node->type == sql_parser::NodeType::NODE_COLUMN_REF || + node->type == sql_parser::NodeType::NODE_IDENTIFIER) { + if (!catalog_.get_column(table, node->value())) return false; + if (other && catalog_.get_column(other, node->value())) return false; + return true; + } + return false; + } + + const sql_parser::AstNode* probe_key_in_join(const sql_parser::AstNode* cond, + const TableInfo* probe, + const TableInfo* build) const { + if (!cond || !probe || !build) return nullptr; + const auto& keys = shards_.get_shard_keys(probe->table_name); + if (keys.size() != 1) return nullptr; + sql_parser::StringRef sk{keys[0].c_str(), static_cast(keys[0].size())}; + std::vector> eqs; + collect_eq_pairs(cond, eqs); + for (const auto& eq : eqs) { + if (is_shard_key_ref(eq.first, sk) && column_on_table(eq.second, build, probe)) + return eq.first; + if (is_shard_key_ref(eq.second, sk) && column_on_table(eq.first, build, probe)) + return eq.second; + } + return nullptr; + } + + sql_parser::AstNode* make_in_list_on_column(const sql_parser::AstNode* col, + const std::vector& values) { + if (!col || values.empty()) return nullptr; + sql_parser::AstNode* stub = sql_parser::make_node( + arena_, sql_parser::NodeType::NODE_IN_LIST, + sql_parser::StringRef{nullptr, 0}); + sql_parser::AstNode* col_copy = sql_parser::make_node( + arena_, col->type, col->value(), col->flags); + col_copy->first_child = col->first_child; + stub->add_child(col_copy); + return build_in_list_from_values(stub, values); + } + + sql_parser::AstNode* and_preds(const sql_parser::AstNode* a, + const sql_parser::AstNode* b) { + if (!a) return const_cast(b); + if (!b) return const_cast(a); + sql_parser::AstNode* n = sql_parser::make_node( + arena_, sql_parser::NodeType::NODE_BINARY_OP, + sql_parser::StringRef{"AND", 3}); + n->add_child(const_cast(a)); + n->add_child(const_cast(b)); + return n; + } + + std::vector collect_build_join_keys(const TableInfo* build, + const sql_parser::AstNode* where_expr, + const sql_parser::AstNode* join_eq_other) { + std::vector out; + if (!build || !join_eq_other || !remote_executor_) return out; + const sql_parser::AstNode* proj[1] = {join_eq_other}; + const auto& shards = shards_.get_shards(build->table_name); + if (shards.empty()) return out; + std::vector targets = shards; + if (shards_.is_sharded(build->table_name) && shards.size() > 1) + return out; + sql_parser::StringRef sql = qb_.build_select( + build, where_expr, proj, 1, nullptr, 0, + nullptr, nullptr, 0, -1, true); + ResultSet rs = remote_executor_->execute(shards[0].backend_name.c_str(), sql); + for (const auto& row : rs.rows) { + if (row.column_count > 0 && value_is_routable(row.get(0))) + out.push_back(copy_value_arena(row.get(0))); + } + return out; + } + + const sql_parser::AstNode* other_eq_side(const sql_parser::AstNode* cond, + const sql_parser::AstNode* probe_key) const { + std::vector> eqs; + collect_eq_pairs(cond, eqs); + for (const auto& eq : eqs) { + if (eq.first == probe_key) return eq.second; + if (eq.second == probe_key) return eq.first; + } + return nullptr; + } + + PlanNode* try_semijoin_prune(PlanNode* join_node, + const TableInfo* left_table, + const TableInfo* right_table) { + if (!join_node || !remote_executor_ || !join_node->join.condition) + return nullptr; + if (!left_table || !right_table) return nullptr; + + bool ls = shards_.is_sharded(left_table->table_name); + bool rs = shards_.is_sharded(right_table->table_name); + if (ls == rs) return nullptr; + + const TableInfo* probe = ls ? left_table : right_table; + const TableInfo* build = ls ? right_table : left_table; + bool probe_is_left = ls; + const sql_parser::AstNode* probe_key = + probe_key_in_join(join_node->join.condition, probe, build); + if (!probe_key) return nullptr; + const sql_parser::AstNode* build_col = + other_eq_side(join_node->join.condition, probe_key); + if (!build_col) return nullptr; + + ScanContext bctx = extract_scan_context( + probe_is_left ? join_node->right : join_node->left); + std::vector keys = collect_build_join_keys(build, bctx.where_expr, build_col); + if (keys.empty()) return nullptr; + + sql_parser::AstNode* in_list = make_in_list_on_column(probe_key, keys); + if (!in_list) return nullptr; + + ScanContext pctx = extract_scan_context( + probe_is_left ? join_node->left : join_node->right); + if (!pctx.scan) return nullptr; + const sql_parser::AstNode* probe_where = and_preds(pctx.where_expr, in_list); + PlanNode* probe_dist = distribute_scan(pctx.scan, probe_where, + nullptr, nullptr, nullptr, false); + + PlanNode* build_dist = nullptr; + if (bctx.scan && !shards_.is_sharded(build->table_name)) { + sql_parser::StringRef sql = qb_.build_select( + build, bctx.where_expr, nullptr, 0, nullptr, 0, + nullptr, nullptr, 0, -1, false); + build_dist = make_remote_scan( + shards_.get_backend(build->table_name), sql, build); + } else { + build_dist = distribute_node(probe_is_left ? join_node->right : join_node->left); + } + if (!probe_dist || !build_dist) return nullptr; + + PlanNode* result = make_plan_node(arena_, PlanNodeType::JOIN); + result->join = join_node->join; + result->left = probe_is_left ? probe_dist : build_dist; + result->right = probe_is_left ? build_dist : probe_dist; + return result; + } + PlanNode* distribute_join(PlanNode* join_node) { const TableInfo* left_table = find_table(join_node->left); const TableInfo* right_table = find_table(join_node->right); @@ -1217,6 +1374,9 @@ class DistributedPlanner { return distribute_colocated_join(join_node, left_table, right_table); } + if (PlanNode* sj = try_semijoin_prune(join_node, left_table, right_table)) + return sj; + PlanNode* left_dist = nullptr; PlanNode* right_dist = nullptr; diff --git a/include/sql_engine/shard_map.h b/include/sql_engine/shard_map.h index 907a651..894a5cd 100644 --- a/include/sql_engine/shard_map.h +++ b/include/sql_engine/shard_map.h @@ -177,7 +177,7 @@ class ShardMap { size_t& out) const { const TableShardConfig* cfg = lookup(table_name); if (!cfg || cfg->shards.empty() || !parts || n == 0) return false; - if (n == 1) { + if (n == 1 || cfg->strategy == RoutingStrategy::RANGE) { return parts[0].is_int ? try_shard_index_for_int(table_name, parts[0].int_val, out) : try_shard_index_for_string(table_name, parts[0].str, diff --git a/scripts/run_pg_sharding_demo.sh b/scripts/run_pg_sharding_demo.sh new file mode 100755 index 0000000..bbe2550 --- /dev/null +++ b/scripts/run_pg_sharding_demo.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# RANGE + LIST + 2PC against the two PostgreSQL shards. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_DIR="$(dirname "$SCRIPT_DIR")" +cd "$PROJECT_DIR" + +if ! docker exec parsersql-pg-shard1 pg_isready -Upostgres &>/dev/null 2>&1; then + echo "ERROR: PG shards not running. Start them with: ./scripts/start_pg_sharding_demo.sh" + exit 1 +fi + +if [ ! -f ./sqlengine ]; then + echo "Building sqlengine..." + make build-sqlengine +fi + +PG1='pgsql://postgres:test@127.0.0.1:16432/testdb?name=pg1' +PG2='pgsql://postgres:test@127.0.0.1:16433/testdb?name=pg2' +TXN_LOG="${TMPDIR:-/tmp}/parsersql-pg-demo.txn" + +run_sql() { + local desc="$1" + local sql="$2" + echo "----------------------------------------------" + echo "QUERY: $desc" + echo "SQL: $sql" + echo "" + echo "$sql" | ./sqlengine \ + --backend "$PG1" \ + --backend "$PG2" \ + --shard "users:id:range:5=pg1,10=pg2" \ + --shard "regions:name:list:us-east=pg1,us-west=pg2" \ + --shard "orders:id:range:105=pg1,110=pg2" \ + --txn-log "$TXN_LOG" \ + 2>&1 + echo "" +} + +echo "==============================================" +echo " PostgreSQL LIST + RANGE + 2PC demo" +echo "==============================================" +echo " pg1 :16432 users 1-5 / us-east" +echo " pg2 :16433 users 6-10 / us-west" +echo "" + +run_sql "RANGE point lookup" \ + "SELECT name FROM users WHERE id = 3" + +run_sql "RANGE BETWEEN prune" \ + "SELECT name FROM users WHERE id BETWEEN 6 AND 10" + +run_sql "LIST point lookup" \ + "SELECT tz FROM regions WHERE name = 'us-west'" + +run_sql "Scatter scan" \ + "SELECT COUNT(*) FROM users" + +echo "==============================================" +echo " 2PC write across both shards (one engine)" +echo "==============================================" +{ + echo "BEGIN" + echo "INSERT INTO users (id, name, age) VALUES (0, 'Zero', 1)" + echo "INSERT INTO users (id, name, age) VALUES (11, 'Eleven', 2)" + echo "COMMIT" +} | ./sqlengine \ + --backend "$PG1" \ + --backend "$PG2" \ + --shard "users:id:range:5=pg1,10=pg2" \ + --shard "regions:name:list:us-east=pg1,us-west=pg2" \ + --shard "orders:id:range:105=pg1,110=pg2" \ + --txn-log "$TXN_LOG" \ + 2>&1 +echo "" + +run_sql "Read back 2PC inserts" \ + "SELECT id, name FROM users WHERE id IN (0, 11)" + +echo "Demo complete. Stop: docker rm -f parsersql-pg-shard1 parsersql-pg-shard2" diff --git a/scripts/start_pg_sharding_demo.sh b/scripts/start_pg_sharding_demo.sh new file mode 100755 index 0000000..f1a1578 --- /dev/null +++ b/scripts/start_pg_sharding_demo.sh @@ -0,0 +1,98 @@ +#!/bin/bash +# Two PostgreSQL shards for LIST + RANGE + 2PC. Ports 16432/16433 +# (15432 is the unit-test backend; 13306 is the MySQL sharding demo). +set -e + +echo "=== Starting 2-shard PostgreSQL demo ===" + +docker rm -f parsersql-pg-shard1 parsersql-pg-shard2 2>/dev/null || true + +docker run -d --name parsersql-pg-shard1 \ + -p 16432:5432 \ + -e POSTGRES_PASSWORD=test \ + -e POSTGRES_DB=testdb \ + postgres:16 \ + -c max_prepared_transactions=16 + +docker run -d --name parsersql-pg-shard2 \ + -p 16433:5432 \ + -e POSTGRES_PASSWORD=test \ + -e POSTGRES_DB=testdb \ + postgres:16 \ + -c max_prepared_transactions=16 + +echo "Waiting for PG shard 1..." +until docker exec parsersql-pg-shard1 pg_isready -Upostgres &>/dev/null 2>&1; do sleep 1; done +echo "PG shard 1 ready" + +echo "Waiting for PG shard 2..." +until docker exec parsersql-pg-shard2 pg_isready -Upostgres &>/dev/null 2>&1; do sleep 1; done +echo "PG shard 2 ready" + +echo "Loading RANGE users 1-5 + LIST region us-east on shard 1..." +docker exec -i parsersql-pg-shard1 psql -Upostgres testdb <<'SQL' +DROP TABLE IF EXISTS orders; +DROP TABLE IF EXISTS users; +DROP TABLE IF EXISTS regions; + +CREATE TABLE users ( + id INT PRIMARY KEY, + name VARCHAR(255) NOT NULL, + age INT +); +CREATE TABLE regions ( + name VARCHAR(64) PRIMARY KEY, + tz VARCHAR(32) +); +CREATE TABLE orders ( + id INT PRIMARY KEY, + user_id INT, + total NUMERIC(10,2) +); + +INSERT INTO users VALUES + (1, 'Alice', 30), + (2, 'Bob', 25), + (3, 'Carol', 35), + (4, 'Dave', 28), + (5, 'Eve', 32); +INSERT INTO regions VALUES ('us-east', 'EST'); +INSERT INTO orders VALUES (101, 1, 150.00), (102, 3, 50.00); +SQL + +echo "Loading RANGE users 6-10 + LIST region us-west on shard 2..." +docker exec -i parsersql-pg-shard2 psql -Upostgres testdb <<'SQL' +DROP TABLE IF EXISTS orders; +DROP TABLE IF EXISTS users; +DROP TABLE IF EXISTS regions; + +CREATE TABLE users ( + id INT PRIMARY KEY, + name VARCHAR(255) NOT NULL, + age INT +); +CREATE TABLE regions ( + name VARCHAR(64) PRIMARY KEY, + tz VARCHAR(32) +); +CREATE TABLE orders ( + id INT PRIMARY KEY, + user_id INT, + total NUMERIC(10,2) +); + +INSERT INTO users VALUES + (6, 'Frank', 40), + (7, 'Grace', 22), + (8, 'Hank', 31), + (9, 'Ivy', 27), + (10, 'Jack', 36); +INSERT INTO regions VALUES ('us-west', 'PST'); +INSERT INTO orders VALUES (106, 6, 80.00), (107, 8, 120.00); +SQL + +echo "PostgreSQL shards ready:" +echo " pg1 127.0.0.1:16432 users 1-5, region us-east" +echo " pg2 127.0.0.1:16433 users 6-10, region us-west" +echo "Run: ./scripts/run_pg_sharding_demo.sh" +echo "Stop: docker rm -f parsersql-pg-shard1 parsersql-pg-shard2" diff --git a/src/sql_engine/tool_config_parser.cpp b/src/sql_engine/tool_config_parser.cpp index e9c4738..8cf06b8 100644 --- a/src/sql_engine/tool_config_parser.cpp +++ b/src/sql_engine/tool_config_parser.cpp @@ -208,10 +208,6 @@ ParsedShard parse_shard_spec(const std::string& spec) { return ps; } } else if (strategy_token == "range") { - if (ps.config.shard_key.find('+') != std::string::npos) { - ps.error = "composite shard keys require HASH strategy: " + spec; - return ps; - } ps.config.strategy = RoutingStrategy::RANGE; for (auto& entry : split_csv(body)) { std::string upper_str, backend; diff --git a/tests/test_distributed_planner.cpp b/tests/test_distributed_planner.cpp index ca52c72..f5f1c88 100644 --- a/tests/test_distributed_planner.cpp +++ b/tests/test_distributed_planner.cpp @@ -1288,6 +1288,57 @@ TEST_F(DistributedPlannerTest, IncompleteCompositeJoinStaysLocal) { EXPECT_FALSE(joins.empty()) << "join on a partial composite key must gather"; } +TEST_F(DistributedPlannerTest, CompositeRangePrunesOnFirstKey) { + TableShardConfig cfg; + cfg.table_name = "users"; + cfg.shard_key = "id+name"; + cfg.shards = {{"shard_1"}, {"shard_2"}, {"shard_3"}}; + cfg.strategy = RoutingStrategy::RANGE; + cfg.ranges = {{5, 0}, {10, 1}, {100000, 2}}; + shard_map.add_table(cfg); + + auto count_remotes = [&](const char* sql) { + Parser parser; + auto pr = parser.parse(sql, std::strlen(sql)); + PlanBuilder builder(catalog, parser.arena()); + PlanNode* plan = builder.build(pr.ast); + DistributedPlanner dp(shard_map, catalog, parser.arena()); + PlanNode* dist = dp.distribute(plan); + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + return remotes.size(); + }; + + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id <= 5"), 1u); + EXPECT_EQ(count_remotes("SELECT * FROM users WHERE id BETWEEN 6 AND 10"), 1u); +} + +TEST_F(DistributedPlannerTest, SemiJoinPrunesProbeShards) { + 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(), + &mock_executor, &functions); + PlanNode* dist = dp.distribute(plan); + ASSERT_NE(dist, nullptr); + + std::vector remotes; + find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes); + bool saw_users_in = false; + for (auto* rs : remotes) { + std::string remote(rs->remote_scan.remote_sql, rs->remote_scan.remote_sql_len); + if (remote.find("users") != std::string::npos && + remote.find("IN") != std::string::npos) + saw_users_in = true; + } + EXPECT_TRUE(saw_users_in) << "semi-join should push IN on the sharded probe"; +} + TEST_F(DistributedPlannerTest, UnknownTableErrors) { catalog.add_table("", "ghost", { {"id", SqlType::make_int(), false}, diff --git a/tests/test_distributed_real.cpp b/tests/test_distributed_real.cpp index ef904f0..9bf397b 100644 --- a/tests/test_distributed_real.cpp +++ b/tests/test_distributed_real.cpp @@ -17,6 +17,7 @@ #include "sql_engine/function_registry.h" #include "sql_engine/session.h" #include "sql_engine/local_txn.h" +#include "sql_engine/distributed_txn.h" #include "sql_parser/parser.h" #include @@ -365,4 +366,92 @@ TEST_F(LiveShardedWriteTest, InsertThenPointSelect) { session.execute_statement("DELETE FROM shard_write_t WHERE id = 9"); } +bool pg_port_available(uint16_t port) { + std::string conninfo = std::string("host=127.0.0.1 port=") + std::to_string(port) + + " user=postgres password=test dbname=testdb connect_timeout=2"; + PGconn* conn = PQconnectdb(conninfo.c_str()); + bool ok = (PQstatus(conn) == CONNECTION_OK); + PQfinish(conn); + return ok; +} + +TEST(LivePgShardedWriteTest, RangeListAndTwoPhaseCommit) { + if (!pg_port_available(16432) || !pg_port_available(16433)) { + GTEST_SKIP() << "Need PG on 16432 and 16433 (start_pg_sharding_demo.sh)"; + } + + MultiRemoteExecutor exec; + BackendConfig p1; + p1.name = "pg1"; + p1.host = "127.0.0.1"; + p1.port = 16432; + p1.user = "postgres"; + p1.password = "test"; + p1.database = "testdb"; + p1.dialect = Dialect::PostgreSQL; + BackendConfig p2 = p1; + p2.name = "pg2"; + p2.port = 16433; + exec.add_backend(p1); + exec.add_backend(p2); + + InMemoryCatalog catalog; + catalog.add_table("", "users", { + {"id", SqlType::make_int(), false}, + {"name", SqlType::make_varchar(255), true}, + {"age", SqlType::make_int(), true}, + }); + catalog.add_table("", "regions", { + {"name", SqlType::make_varchar(64), false}, + {"tz", SqlType::make_varchar(32), true}, + }); + + ShardMap shards; + TableShardConfig users; + users.table_name = "users"; + users.shard_key = "id"; + users.shards = {{"pg1"}, {"pg2"}}; + users.strategy = RoutingStrategy::RANGE; + users.ranges = {{5, 0}, {100000, 1}}; + shards.add_table(users); + TableShardConfig regions; + regions.table_name = "regions"; + regions.shard_key = "name"; + regions.shards = {{"pg1"}, {"pg2"}}; + regions.strategy = RoutingStrategy::LIST; + regions.list = { + {false, 0, "us-east", 0}, + {false, 0, "us-west", 1}, + }; + shards.add_table(regions); + + Arena txn_arena{65536, 1048576}; + DistributedTransactionManager txn( + exec, DistributedTransactionManager::BackendDialect::POSTGRESQL); + Session session(catalog, txn); + session.set_remote_executor(&exec); + session.set_shard_map(&shards); + + auto east = session.execute_query("SELECT tz FROM regions WHERE name = 'us-east'"); + auto west = session.execute_query("SELECT tz FROM regions WHERE name = 'us-west'"); + ASSERT_EQ(east.row_count(), 1u); + ASSERT_EQ(west.row_count(), 1u); + + EXPECT_TRUE(session.begin()); + auto i0 = session.execute_statement( + "INSERT INTO users (id, name, age) VALUES (0, 'Zero', 1)"); + auto i11 = session.execute_statement( + "INSERT INTO users (id, name, age) VALUES (11, 'Eleven', 2)"); + EXPECT_TRUE(i0.success) << i0.error_message; + EXPECT_TRUE(i11.success) << i11.error_message; + EXPECT_TRUE(session.commit()); + + auto back = session.execute_query("SELECT name FROM users WHERE id IN (0, 11)"); + EXPECT_EQ(back.row_count(), 2u); + + session.execute_statement("DELETE FROM users WHERE id = 0"); + session.execute_statement("DELETE FROM users WHERE id = 11"); + exec.disconnect_all(); +} + } // namespace diff --git a/tests/test_shard_map.cpp b/tests/test_shard_map.cpp index 31dd0a9..a656e82 100644 --- a/tests/test_shard_map.cpp +++ b/tests/test_shard_map.cpp @@ -278,3 +278,22 @@ TEST(ShardMapTest, CompositeHashIsDeterministicAndDiffersFromSingle) { EXPECT_LT(ia, 3u); EXPECT_LT(ib, 3u); } + +TEST(ShardMapTest, CompositeRangeRoutesOnFirstKey) { + TableShardConfig cfg; + cfg.table_name = "kv"; + cfg.shard_key = "tenant_id+id"; + cfg.shards = {ShardInfo{"s0"}, ShardInfo{"s1"}}; + cfg.strategy = RoutingStrategy::RANGE; + cfg.ranges = {ShardRange{5, 0}, ShardRange{100, 1}}; + ShardMap map; + map.add_table(cfg); + + ShardKeyPart low[] = {{true, 3, nullptr, 0}, {true, 99, nullptr, 0}}; + ShardKeyPart high[] = {{true, 9, nullptr, 0}, {true, 1, nullptr, 0}}; + size_t il = 99, ih = 99; + ASSERT_TRUE(map.try_shard_index_for_parts(sref("kv"), low, 2, il)); + ASSERT_TRUE(map.try_shard_index_for_parts(sref("kv"), high, 2, ih)); + EXPECT_EQ(il, 0u); + EXPECT_EQ(ih, 1u); +} diff --git a/tests/test_ssl_config.cpp b/tests/test_ssl_config.cpp index b60a56d..6516e7c 100644 --- a/tests/test_ssl_config.cpp +++ b/tests/test_ssl_config.cpp @@ -237,8 +237,10 @@ TEST(SSLConfigTest, ParseShardSpecCompositeHash) { ASSERT_EQ(ps.config.shards.size(), 3u); } -TEST(SSLConfigTest, ParseShardSpecRejectsCompositeRange) { +TEST(SSLConfigTest, ParseShardSpecCompositeRangeUsesFirstKey) { auto ps = parse_shard_spec("kv:tenant_id+id:range:5=s0,10=s1"); - EXPECT_FALSE(ps.ok); - EXPECT_NE(ps.error.find("HASH"), std::string::npos); + ASSERT_TRUE(ps.ok); + EXPECT_EQ(ps.config.strategy, RoutingStrategy::RANGE); + EXPECT_EQ(ps.config.shard_key, "tenant_id+id"); + ASSERT_EQ(ps.config.ranges.size(), 2u); } diff --git a/tools/sqlengine.cpp b/tools/sqlengine.cpp index ea6adc1..d14ebb4 100644 --- a/tools/sqlengine.cpp +++ b/tools/sqlengine.cpp @@ -356,8 +356,14 @@ int main(int argc, char* argv[]) { } if (remote_exec) { + bool all_pg = !backends.empty(); + for (const auto& b : backends) { + if (b.dialect != Dialect::PostgreSQL) { all_pg = false; break; } + } dtxn.reset(new DistributedTransactionManager( - *remote_exec, DistributedTransactionManager::BackendDialect::MYSQL)); + *remote_exec, + all_pg ? DistributedTransactionManager::BackendDialect::POSTGRESQL + : 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;