Skip to content

Commit a61dc13

Browse files
renecannaorenecannao
authored
Merge pull request #62 from ProxySQL/feat/shard-demo-bloom-composite-range
feat: composite RANGE, semi-join prune, PG LIST/RANGE + 2PC demo
2 parents 7ade921 + 7b1726b commit a61dc13

10 files changed

Lines changed: 512 additions & 10 deletions

‎include/sql_engine/distributed_planner.h‎

Lines changed: 161 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,13 @@ class DistributedPlanner {
357357
if (keys.empty()) return all_shards;
358358

359359
std::vector<size_t> target_indices;
360-
if (keys.size() > 1) {
360+
if (keys.size() > 1 &&
361+
shards_.routing_strategy(table->table_name) == RoutingStrategy::RANGE) {
362+
sql_parser::StringRef first{keys[0].c_str(),
363+
static_cast<uint32_t>(keys[0].size())};
364+
extract_shard_targets(where_expr, first, table->table_name,
365+
all_shards.size(), target_indices);
366+
} else if (keys.size() > 1) {
361367
extract_composite_targets(where_expr, keys, table->table_name, target_indices);
362368
} else {
363369
sql_parser::StringRef shard_key{keys[0].c_str(),
@@ -1203,6 +1209,157 @@ class DistributedPlanner {
12031209
return current ? current : join_node;
12041210
}
12051211

1212+
bool column_on_table(const sql_parser::AstNode* node, const TableInfo* table,
1213+
const TableInfo* other) const {
1214+
if (!node || !table) return false;
1215+
if (node->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
1216+
const sql_parser::AstNode* t = node->first_child;
1217+
if (!t) return false;
1218+
sql_parser::StringRef tn = t->value();
1219+
if (table->table_name.equals_ci(tn.ptr, tn.len)) return true;
1220+
if (table->alias.ptr && table->alias.equals_ci(tn.ptr, tn.len)) return true;
1221+
return false;
1222+
}
1223+
if (node->type == sql_parser::NodeType::NODE_COLUMN_REF ||
1224+
node->type == sql_parser::NodeType::NODE_IDENTIFIER) {
1225+
if (!catalog_.get_column(table, node->value())) return false;
1226+
if (other && catalog_.get_column(other, node->value())) return false;
1227+
return true;
1228+
}
1229+
return false;
1230+
}
1231+
1232+
const sql_parser::AstNode* probe_key_in_join(const sql_parser::AstNode* cond,
1233+
const TableInfo* probe,
1234+
const TableInfo* build) const {
1235+
if (!cond || !probe || !build) return nullptr;
1236+
const auto& keys = shards_.get_shard_keys(probe->table_name);
1237+
if (keys.size() != 1) return nullptr;
1238+
sql_parser::StringRef sk{keys[0].c_str(), static_cast<uint32_t>(keys[0].size())};
1239+
std::vector<std::pair<const sql_parser::AstNode*, const sql_parser::AstNode*>> eqs;
1240+
collect_eq_pairs(cond, eqs);
1241+
for (const auto& eq : eqs) {
1242+
if (is_shard_key_ref(eq.first, sk) && column_on_table(eq.second, build, probe))
1243+
return eq.first;
1244+
if (is_shard_key_ref(eq.second, sk) && column_on_table(eq.first, build, probe))
1245+
return eq.second;
1246+
}
1247+
return nullptr;
1248+
}
1249+
1250+
sql_parser::AstNode* make_in_list_on_column(const sql_parser::AstNode* col,
1251+
const std::vector<Value>& values) {
1252+
if (!col || values.empty()) return nullptr;
1253+
sql_parser::AstNode* stub = sql_parser::make_node(
1254+
arena_, sql_parser::NodeType::NODE_IN_LIST,
1255+
sql_parser::StringRef{nullptr, 0});
1256+
sql_parser::AstNode* col_copy = sql_parser::make_node(
1257+
arena_, col->type, col->value(), col->flags);
1258+
col_copy->first_child = col->first_child;
1259+
stub->add_child(col_copy);
1260+
return build_in_list_from_values(stub, values);
1261+
}
1262+
1263+
sql_parser::AstNode* and_preds(const sql_parser::AstNode* a,
1264+
const sql_parser::AstNode* b) {
1265+
if (!a) return const_cast<sql_parser::AstNode*>(b);
1266+
if (!b) return const_cast<sql_parser::AstNode*>(a);
1267+
sql_parser::AstNode* n = sql_parser::make_node(
1268+
arena_, sql_parser::NodeType::NODE_BINARY_OP,
1269+
sql_parser::StringRef{"AND", 3});
1270+
n->add_child(const_cast<sql_parser::AstNode*>(a));
1271+
n->add_child(const_cast<sql_parser::AstNode*>(b));
1272+
return n;
1273+
}
1274+
1275+
std::vector<Value> collect_build_join_keys(const TableInfo* build,
1276+
const sql_parser::AstNode* where_expr,
1277+
const sql_parser::AstNode* join_eq_other) {
1278+
std::vector<Value> out;
1279+
if (!build || !join_eq_other || !remote_executor_) return out;
1280+
const sql_parser::AstNode* proj[1] = {join_eq_other};
1281+
const auto& shards = shards_.get_shards(build->table_name);
1282+
if (shards.empty()) return out;
1283+
std::vector<ShardInfo> targets = shards;
1284+
if (shards_.is_sharded(build->table_name) && shards.size() > 1)
1285+
return out;
1286+
sql_parser::StringRef sql = qb_.build_select(
1287+
build, where_expr, proj, 1, nullptr, 0,
1288+
nullptr, nullptr, 0, -1, true);
1289+
ResultSet rs = remote_executor_->execute(shards[0].backend_name.c_str(), sql);
1290+
for (const auto& row : rs.rows) {
1291+
if (row.column_count > 0 && value_is_routable(row.get(0)))
1292+
out.push_back(copy_value_arena(row.get(0)));
1293+
}
1294+
return out;
1295+
}
1296+
1297+
const sql_parser::AstNode* other_eq_side(const sql_parser::AstNode* cond,
1298+
const sql_parser::AstNode* probe_key) const {
1299+
std::vector<std::pair<const sql_parser::AstNode*, const sql_parser::AstNode*>> eqs;
1300+
collect_eq_pairs(cond, eqs);
1301+
for (const auto& eq : eqs) {
1302+
if (eq.first == probe_key) return eq.second;
1303+
if (eq.second == probe_key) return eq.first;
1304+
}
1305+
return nullptr;
1306+
}
1307+
1308+
PlanNode* try_semijoin_prune(PlanNode* join_node,
1309+
const TableInfo* left_table,
1310+
const TableInfo* right_table) {
1311+
if (!join_node || !remote_executor_ || !join_node->join.condition)
1312+
return nullptr;
1313+
if (!left_table || !right_table) return nullptr;
1314+
1315+
bool ls = shards_.is_sharded(left_table->table_name);
1316+
bool rs = shards_.is_sharded(right_table->table_name);
1317+
if (ls == rs) return nullptr;
1318+
1319+
const TableInfo* probe = ls ? left_table : right_table;
1320+
const TableInfo* build = ls ? right_table : left_table;
1321+
bool probe_is_left = ls;
1322+
const sql_parser::AstNode* probe_key =
1323+
probe_key_in_join(join_node->join.condition, probe, build);
1324+
if (!probe_key) return nullptr;
1325+
const sql_parser::AstNode* build_col =
1326+
other_eq_side(join_node->join.condition, probe_key);
1327+
if (!build_col) return nullptr;
1328+
1329+
ScanContext bctx = extract_scan_context(
1330+
probe_is_left ? join_node->right : join_node->left);
1331+
std::vector<Value> keys = collect_build_join_keys(build, bctx.where_expr, build_col);
1332+
if (keys.empty()) return nullptr;
1333+
1334+
sql_parser::AstNode* in_list = make_in_list_on_column(probe_key, keys);
1335+
if (!in_list) return nullptr;
1336+
1337+
ScanContext pctx = extract_scan_context(
1338+
probe_is_left ? join_node->left : join_node->right);
1339+
if (!pctx.scan) return nullptr;
1340+
const sql_parser::AstNode* probe_where = and_preds(pctx.where_expr, in_list);
1341+
PlanNode* probe_dist = distribute_scan(pctx.scan, probe_where,
1342+
nullptr, nullptr, nullptr, false);
1343+
1344+
PlanNode* build_dist = nullptr;
1345+
if (bctx.scan && !shards_.is_sharded(build->table_name)) {
1346+
sql_parser::StringRef sql = qb_.build_select(
1347+
build, bctx.where_expr, nullptr, 0, nullptr, 0,
1348+
nullptr, nullptr, 0, -1, false);
1349+
build_dist = make_remote_scan(
1350+
shards_.get_backend(build->table_name), sql, build);
1351+
} else {
1352+
build_dist = distribute_node(probe_is_left ? join_node->right : join_node->left);
1353+
}
1354+
if (!probe_dist || !build_dist) return nullptr;
1355+
1356+
PlanNode* result = make_plan_node(arena_, PlanNodeType::JOIN);
1357+
result->join = join_node->join;
1358+
result->left = probe_is_left ? probe_dist : build_dist;
1359+
result->right = probe_is_left ? build_dist : probe_dist;
1360+
return result;
1361+
}
1362+
12061363
PlanNode* distribute_join(PlanNode* join_node) {
12071364
const TableInfo* left_table = find_table(join_node->left);
12081365
const TableInfo* right_table = find_table(join_node->right);
@@ -1217,6 +1374,9 @@ class DistributedPlanner {
12171374
return distribute_colocated_join(join_node, left_table, right_table);
12181375
}
12191376

1377+
if (PlanNode* sj = try_semijoin_prune(join_node, left_table, right_table))
1378+
return sj;
1379+
12201380
PlanNode* left_dist = nullptr;
12211381
PlanNode* right_dist = nullptr;
12221382

‎include/sql_engine/shard_map.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ class ShardMap {
177177
size_t& out) const {
178178
const TableShardConfig* cfg = lookup(table_name);
179179
if (!cfg || cfg->shards.empty() || !parts || n == 0) return false;
180-
if (n == 1) {
180+
if (n == 1 || cfg->strategy == RoutingStrategy::RANGE) {
181181
return parts[0].is_int
182182
? try_shard_index_for_int(table_name, parts[0].int_val, out)
183183
: try_shard_index_for_string(table_name, parts[0].str,

‎scripts/run_pg_sharding_demo.sh‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
#!/bin/bash
2+
# RANGE + LIST + 2PC against the two PostgreSQL shards.
3+
set -e
4+
5+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
6+
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
7+
cd "$PROJECT_DIR"
8+
9+
if ! docker exec parsersql-pg-shard1 pg_isready -Upostgres &>/dev/null 2>&1; then
10+
echo "ERROR: PG shards not running. Start them with: ./scripts/start_pg_sharding_demo.sh"
11+
exit 1
12+
fi
13+
14+
if [ ! -f ./sqlengine ]; then
15+
echo "Building sqlengine..."
16+
make build-sqlengine
17+
fi
18+
19+
PG1='pgsql://postgres:test@127.0.0.1:16432/testdb?name=pg1'
20+
PG2='pgsql://postgres:test@127.0.0.1:16433/testdb?name=pg2'
21+
TXN_LOG="${TMPDIR:-/tmp}/parsersql-pg-demo.txn"
22+
23+
run_sql() {
24+
local desc="$1"
25+
local sql="$2"
26+
echo "----------------------------------------------"
27+
echo "QUERY: $desc"
28+
echo "SQL: $sql"
29+
echo ""
30+
echo "$sql" | ./sqlengine \
31+
--backend "$PG1" \
32+
--backend "$PG2" \
33+
--shard "users:id:range:5=pg1,10=pg2" \
34+
--shard "regions:name:list:us-east=pg1,us-west=pg2" \
35+
--shard "orders:id:range:105=pg1,110=pg2" \
36+
--txn-log "$TXN_LOG" \
37+
2>&1
38+
echo ""
39+
}
40+
41+
echo "=============================================="
42+
echo " PostgreSQL LIST + RANGE + 2PC demo"
43+
echo "=============================================="
44+
echo " pg1 :16432 users 1-5 / us-east"
45+
echo " pg2 :16433 users 6-10 / us-west"
46+
echo ""
47+
48+
run_sql "RANGE point lookup" \
49+
"SELECT name FROM users WHERE id = 3"
50+
51+
run_sql "RANGE BETWEEN prune" \
52+
"SELECT name FROM users WHERE id BETWEEN 6 AND 10"
53+
54+
run_sql "LIST point lookup" \
55+
"SELECT tz FROM regions WHERE name = 'us-west'"
56+
57+
run_sql "Scatter scan" \
58+
"SELECT COUNT(*) FROM users"
59+
60+
echo "=============================================="
61+
echo " 2PC write across both shards (one engine)"
62+
echo "=============================================="
63+
{
64+
echo "BEGIN"
65+
echo "INSERT INTO users (id, name, age) VALUES (0, 'Zero', 1)"
66+
echo "INSERT INTO users (id, name, age) VALUES (11, 'Eleven', 2)"
67+
echo "COMMIT"
68+
} | ./sqlengine \
69+
--backend "$PG1" \
70+
--backend "$PG2" \
71+
--shard "users:id:range:5=pg1,10=pg2" \
72+
--shard "regions:name:list:us-east=pg1,us-west=pg2" \
73+
--shard "orders:id:range:105=pg1,110=pg2" \
74+
--txn-log "$TXN_LOG" \
75+
2>&1
76+
echo ""
77+
78+
run_sql "Read back 2PC inserts" \
79+
"SELECT id, name FROM users WHERE id IN (0, 11)"
80+
81+
echo "Demo complete. Stop: docker rm -f parsersql-pg-shard1 parsersql-pg-shard2"

‎scripts/start_pg_sharding_demo.sh‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
#!/bin/bash
2+
# Two PostgreSQL shards for LIST + RANGE + 2PC. Ports 16432/16433
3+
# (15432 is the unit-test backend; 13306 is the MySQL sharding demo).
4+
set -e
5+
6+
echo "=== Starting 2-shard PostgreSQL demo ==="
7+
8+
docker rm -f parsersql-pg-shard1 parsersql-pg-shard2 2>/dev/null || true
9+
10+
docker run -d --name parsersql-pg-shard1 \
11+
-p 16432:5432 \
12+
-e POSTGRES_PASSWORD=test \
13+
-e POSTGRES_DB=testdb \
14+
postgres:16 \
15+
-c max_prepared_transactions=16
16+
17+
docker run -d --name parsersql-pg-shard2 \
18+
-p 16433:5432 \
19+
-e POSTGRES_PASSWORD=test \
20+
-e POSTGRES_DB=testdb \
21+
postgres:16 \
22+
-c max_prepared_transactions=16
23+
24+
echo "Waiting for PG shard 1..."
25+
until docker exec parsersql-pg-shard1 pg_isready -Upostgres &>/dev/null 2>&1; do sleep 1; done
26+
echo "PG shard 1 ready"
27+
28+
echo "Waiting for PG shard 2..."
29+
until docker exec parsersql-pg-shard2 pg_isready -Upostgres &>/dev/null 2>&1; do sleep 1; done
30+
echo "PG shard 2 ready"
31+
32+
echo "Loading RANGE users 1-5 + LIST region us-east on shard 1..."
33+
docker exec -i parsersql-pg-shard1 psql -Upostgres testdb <<'SQL'
34+
DROP TABLE IF EXISTS orders;
35+
DROP TABLE IF EXISTS users;
36+
DROP TABLE IF EXISTS regions;
37+
38+
CREATE TABLE users (
39+
id INT PRIMARY KEY,
40+
name VARCHAR(255) NOT NULL,
41+
age INT
42+
);
43+
CREATE TABLE regions (
44+
name VARCHAR(64) PRIMARY KEY,
45+
tz VARCHAR(32)
46+
);
47+
CREATE TABLE orders (
48+
id INT PRIMARY KEY,
49+
user_id INT,
50+
total NUMERIC(10,2)
51+
);
52+
53+
INSERT INTO users VALUES
54+
(1, 'Alice', 30),
55+
(2, 'Bob', 25),
56+
(3, 'Carol', 35),
57+
(4, 'Dave', 28),
58+
(5, 'Eve', 32);
59+
INSERT INTO regions VALUES ('us-east', 'EST');
60+
INSERT INTO orders VALUES (101, 1, 150.00), (102, 3, 50.00);
61+
SQL
62+
63+
echo "Loading RANGE users 6-10 + LIST region us-west on shard 2..."
64+
docker exec -i parsersql-pg-shard2 psql -Upostgres testdb <<'SQL'
65+
DROP TABLE IF EXISTS orders;
66+
DROP TABLE IF EXISTS users;
67+
DROP TABLE IF EXISTS regions;
68+
69+
CREATE TABLE users (
70+
id INT PRIMARY KEY,
71+
name VARCHAR(255) NOT NULL,
72+
age INT
73+
);
74+
CREATE TABLE regions (
75+
name VARCHAR(64) PRIMARY KEY,
76+
tz VARCHAR(32)
77+
);
78+
CREATE TABLE orders (
79+
id INT PRIMARY KEY,
80+
user_id INT,
81+
total NUMERIC(10,2)
82+
);
83+
84+
INSERT INTO users VALUES
85+
(6, 'Frank', 40),
86+
(7, 'Grace', 22),
87+
(8, 'Hank', 31),
88+
(9, 'Ivy', 27),
89+
(10, 'Jack', 36);
90+
INSERT INTO regions VALUES ('us-west', 'PST');
91+
INSERT INTO orders VALUES (106, 6, 80.00), (107, 8, 120.00);
92+
SQL
93+
94+
echo "PostgreSQL shards ready:"
95+
echo " pg1 127.0.0.1:16432 users 1-5, region us-east"
96+
echo " pg2 127.0.0.1:16433 users 6-10, region us-west"
97+
echo "Run: ./scripts/run_pg_sharding_demo.sh"
98+
echo "Stop: docker rm -f parsersql-pg-shard1 parsersql-pg-shard2"

‎src/sql_engine/tool_config_parser.cpp‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -208,10 +208,6 @@ ParsedShard parse_shard_spec(const std::string& spec) {
208208
return ps;
209209
}
210210
} else if (strategy_token == "range") {
211-
if (ps.config.shard_key.find('+') != std::string::npos) {
212-
ps.error = "composite shard keys require HASH strategy: " + spec;
213-
return ps;
214-
}
215211
ps.config.strategy = RoutingStrategy::RANGE;
216212
for (auto& entry : split_csv(body)) {
217213
std::string upper_str, backend;

0 commit comments

Comments
 (0)