Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
331 changes: 291 additions & 40 deletions include/sql_engine/distributed_planner.h

Large diffs are not rendered by default.

12 changes: 12 additions & 0 deletions include/sql_engine/distributed_txn.h
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,18 @@ class DistributedTransactionManager : public TransactionManager {
return executor_.execute_dml(backend_name, sql);
}

ResultSet route_query(const char* backend_name,
sql_parser::StringRef sql) override {
if (!active_) return executor_.execute(backend_name, sql);
auto it = sessions_.find(backend_name);
if (it != sessions_.end() && it->second) {
return it->second->execute(sql);
}
return executor_.execute(backend_name, sql);
}

bool route_query_supported() const override { return true; }

bool commit() override {
if (!active_) return false;
if (participants_.empty()) {
Expand Down
23 changes: 17 additions & 6 deletions include/sql_engine/operators/aggregate_op.h
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
#include <string>
#include <cstring>
#include <cmath>
#include <unordered_set>
#include "sql_parser/common.h"

namespace sql_engine {

Expand Down Expand Up @@ -131,6 +133,8 @@ class AggregateOperator : public Operator {
Value max_val{};
bool has_value = false;
bool count_star = false; // COUNT(*)
bool distinct = false;
std::unordered_set<std::string> seen;
};

struct GroupState {
Expand Down Expand Up @@ -186,9 +190,9 @@ class AggregateOperator : public Operator {

if (expr->type == sql_parser::NodeType::NODE_FUNCTION_CALL) {
sql_parser::StringRef name = expr->value();
state.distinct = (expr->flags & sql_parser::FLAG_FUNC_DISTINCT) != 0;
if (name.equals_ci("COUNT", 5)) {
state.type = AggType::COUNT;
// Check for COUNT(*)
const sql_parser::AstNode* arg = expr->first_child;
if (arg && arg->type == sql_parser::NodeType::NODE_ASTERISK) {
state.count_star = true;
Expand All @@ -203,25 +207,32 @@ class AggregateOperator : public Operator {
state.type = AggType::EXPR;
}

static bool note_distinct(AggState& state, const sql_parser::AstNode* expr,
const Value& v) {
bool distinct = state.distinct ||
(expr && (expr->flags & sql_parser::FLAG_FUNC_DISTINCT));
if (!distinct) return true;
return state.seen.insert(value_to_string(v)).second;
}

void update_agg(AggState& state, const sql_parser::AstNode* expr,
const std::function<Value(sql_parser::StringRef)>& resolver) {
switch (state.type) {
case AggType::COUNT: {
if (state.count_star) {
state.count++;
} else {
// COUNT(expr) - count non-null values
const sql_parser::AstNode* arg = expr->first_child;
Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
if (!v.is_null()) state.count++;
if (!v.is_null() && note_distinct(state, expr, v)) state.count++;
}
break;
}
case AggType::SUM:
case AggType::AVG: {
const sql_parser::AstNode* arg = expr->first_child;
Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
if (!v.is_null()) {
if (!v.is_null() && note_distinct(state, expr, v)) {
state.sum += v.to_double();
state.count++;
state.has_value = true;
Expand All @@ -231,7 +242,7 @@ class AggregateOperator : public Operator {
case AggType::MIN: {
const sql_parser::AstNode* arg = expr->first_child;
Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
if (!v.is_null()) {
if (!v.is_null() && note_distinct(state, expr, v)) {
if (!state.has_value || compare_values(v, state.min_val) < 0) {
state.min_val = v;
state.has_value = true;
Expand All @@ -242,7 +253,7 @@ class AggregateOperator : public Operator {
case AggType::MAX: {
const sql_parser::AstNode* arg = expr->first_child;
Value v = evaluate_expression<D>(arg, resolver, functions_, arena_);
if (!v.is_null()) {
if (!v.is_null() && note_distinct(state, expr, v)) {
if (!state.has_value || compare_values(v, state.max_val) > 0) {
state.max_val = v;
state.has_value = true;
Expand Down
56 changes: 54 additions & 2 deletions include/sql_engine/plan_builder.h
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include "sql_parser/common.h"
#include "sql_parser/arena.h"
#include <cstring>
#include <cstdlib>
#include <vector>

namespace sql_engine {
Expand Down Expand Up @@ -116,6 +117,55 @@ class PlanBuilder {
return false;
}

static const sql_parser::AstNode* select_item_expr(const sql_parser::AstNode* item) {
return item ? item->first_child : nullptr;
}

static sql_parser::StringRef select_item_alias(const sql_parser::AstNode* item) {
if (!item) return {};
for (const sql_parser::AstNode* c = item->first_child; c; c = c->next_sibling) {
if (c->type == sql_parser::NodeType::NODE_ALIAS) return c->value();
}
return {};
}

static const sql_parser::AstNode* resolve_order_key(
const sql_parser::AstNode* key, const sql_parser::AstNode* select_items) {
if (!key || !select_items) return key;
uint16_t n = count_children(select_items);
if (n == 0) return key;

if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
sql_parser::StringRef sv = key->value();
if (!sv.ptr || sv.len == 0) return key;
int64_t pos = std::strtoll(sv.ptr, nullptr, 10);
if (pos < 1 || pos > static_cast<int64_t>(n)) return key;
uint16_t idx = 0;
for (const sql_parser::AstNode* item = select_items->first_child; item;
item = item->next_sibling, ++idx) {
if (idx + 1 == static_cast<uint16_t>(pos)) {
const sql_parser::AstNode* expr = select_item_expr(item);
return expr ? expr : key;
}
}
return key;
}

if (key->type == sql_parser::NodeType::NODE_COLUMN_REF ||
key->type == sql_parser::NodeType::NODE_IDENTIFIER) {
sql_parser::StringRef name = key->value();
for (const sql_parser::AstNode* item = select_items->first_child; item;
item = item->next_sibling) {
sql_parser::StringRef alias = select_item_alias(item);
if (alias.ptr && alias.equals_ci(name.ptr, name.len)) {
const sql_parser::AstNode* expr = select_item_expr(item);
return expr ? expr : key;
}
}
}
return key;
}

// Check if an expression (or any descendant) contains an aggregate function call.
// Does NOT recurse into subqueries -- aggregates inside subqueries belong
// to the subquery's own aggregation, not the outer query.
Expand Down Expand Up @@ -279,10 +329,12 @@ class PlanBuilder {
arena_.allocate(sizeof(sql_parser::AstNode*) * cnt));
auto* dirs = static_cast<uint8_t*>(arena_.allocate(cnt));

const sql_parser::AstNode* select_items =
find_child(select_ast, sql_parser::NodeType::NODE_SELECT_ITEM_LIST);

uint16_t idx = 0;
for (const sql_parser::AstNode* item = order_by->first_child; item; item = item->next_sibling) {
// First child is the key expression
keys[idx] = item->first_child;
keys[idx] = resolve_order_key(item->first_child, select_items);
// Check for DESC direction (second child with "DESC" value)
dirs[idx] = 0; // ASC by default
const sql_parser::AstNode* dir_node = find_child(item, sql_parser::NodeType::NODE_IDENTIFIER);
Expand Down
130 changes: 127 additions & 3 deletions include/sql_engine/plan_executor.h
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@
#include <string>
#include <vector>
#include <memory>
#include <cstdlib>
#include <cstring>

namespace sql_engine {

Expand Down Expand Up @@ -846,15 +848,128 @@ class PlanExecutor {
return ptr;
}

static bool same_agg_call(const sql_parser::AstNode* a, const sql_parser::AstNode* b) {
if (!a || !b) return false;
if (a->type != sql_parser::NodeType::NODE_FUNCTION_CALL ||
b->type != sql_parser::NodeType::NODE_FUNCTION_CALL) return false;
return a->value().equals_ci(b->value().ptr, b->value().len);
}

const sql_parser::AstNode* rewrite_having_expr(const sql_parser::AstNode* expr,
PlanNode* agg_node) {
if (!expr || !agg_node) return expr;
uint16_t group_count = 0;
uint16_t agg_count = 0;
const sql_parser::AstNode** agg_exprs = nullptr;
const sql_parser::AstNode** group_by = nullptr;
if (agg_node->type == PlanNodeType::AGGREGATE) {
group_count = agg_node->aggregate.group_count;
agg_count = agg_node->aggregate.agg_count;
agg_exprs = agg_node->aggregate.agg_exprs;
group_by = agg_node->aggregate.group_by;
} else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE) {
group_count = agg_node->merge_aggregate.group_key_count;
if (agg_node->merge_aggregate.output_exprs &&
agg_node->merge_aggregate.output_expr_count > group_count) {
agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
agg_count = static_cast<uint16_t>(
agg_node->merge_aggregate.output_expr_count - group_count);
group_by = agg_node->merge_aggregate.output_exprs;
}
} else {
return expr;
}

if (expr->type == sql_parser::NodeType::NODE_FUNCTION_CALL && agg_exprs) {
for (uint16_t i = 0; i < agg_count; ++i) {
if (same_agg_call(expr, agg_exprs[i])) {
sql_parser::StringRef name = expr->value();
return sql_parser::make_node(
arena_, sql_parser::NodeType::NODE_IDENTIFIER, name);
}
}
}

bool changed = false;
sql_parser::AstNode* clone = sql_parser::make_node(
arena_, expr->type, expr->value(), expr->flags);
for (const sql_parser::AstNode* c = expr->first_child; c; c = c->next_sibling) {
const sql_parser::AstNode* rw = rewrite_having_expr(c, agg_node);
if (rw != c) changed = true;
if (rw) clone->add_child(const_cast<sql_parser::AstNode*>(rw));
}
(void)group_count;
(void)group_by;
return changed ? clone : expr;
}

const TableInfo* make_agg_output_table(PlanNode* agg_node) {
if (!agg_node) return nullptr;
uint16_t group_count = 0;
uint16_t agg_count = 0;
const sql_parser::AstNode** group_by = nullptr;
const sql_parser::AstNode** agg_exprs = nullptr;
if (agg_node->type == PlanNodeType::AGGREGATE) {
group_count = agg_node->aggregate.group_count;
agg_count = agg_node->aggregate.agg_count;
group_by = agg_node->aggregate.group_by;
agg_exprs = agg_node->aggregate.agg_exprs;
} else if (agg_node->type == PlanNodeType::MERGE_AGGREGATE &&
agg_node->merge_aggregate.output_exprs) {
group_count = agg_node->merge_aggregate.group_key_count;
group_by = agg_node->merge_aggregate.output_exprs;
if (agg_node->merge_aggregate.output_expr_count > group_count) {
agg_exprs = agg_node->merge_aggregate.output_exprs + group_count;
agg_count = static_cast<uint16_t>(
agg_node->merge_aggregate.output_expr_count - group_count);
}
} else {
return nullptr;
}

uint16_t n = static_cast<uint16_t>(group_count + agg_count);
if (n == 0) return nullptr;
auto* cols = static_cast<ColumnInfo*>(arena_.allocate(sizeof(ColumnInfo) * n));
if (!cols) return nullptr;
for (uint16_t i = 0; i < group_count; ++i) {
cols[i].ordinal = i;
cols[i].nullable = true;
cols[i].type = SqlType::make_int();
cols[i].name = (group_by && group_by[i]) ? group_by[i]->value()
: sql_parser::StringRef{};
}
for (uint16_t i = 0; i < agg_count; ++i) {
cols[group_count + i].ordinal = static_cast<uint16_t>(group_count + i);
cols[group_count + i].nullable = true;
cols[group_count + i].type = SqlType::make_int();
cols[group_count + i].name = (agg_exprs && agg_exprs[i])
? agg_exprs[i]->value() : sql_parser::StringRef{};
}
auto* ti = static_cast<TableInfo*>(arena_.allocate(sizeof(TableInfo)));
if (!ti) return nullptr;
std::memset(ti, 0, sizeof(TableInfo));
ti->columns = cols;
ti->column_count = n;
return ti;
}

Operator* build_filter(PlanNode* node) {
Operator* child = build_operator(node->left);
if (!child && node->left) return nullptr;

std::vector<const TableInfo*> tables;
collect_tables(node->left, tables);
const sql_parser::AstNode* expr = node->filter.expr;
if (node->left && (node->left->type == PlanNodeType::AGGREGATE ||
node->left->type == PlanNodeType::MERGE_AGGREGATE)) {
expr = rewrite_having_expr(expr, node->left);
const TableInfo* synth = make_agg_output_table(node->left);
if (synth) tables.push_back(synth);
} else {
collect_tables(node->left, tables);
}

auto op = std::make_unique<FilterOperator<D>>(
child, node->filter.expr, catalog_, tables, functions_, arena_,
child, expr, catalog_, tables, functions_, arena_,
&subquery_exec_, outer_resolver_);
Operator* ptr = op.get();
operators_.push_back(std::move(op));
Expand Down Expand Up @@ -1170,12 +1285,21 @@ class PlanExecutor {

uint16_t resolve_column_index(const sql_parser::AstNode* key, const TableInfo* table) {
if (!key || !table) return 0;
if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
sql_parser::StringRef sv = key->value();
if (!sv.ptr || sv.len == 0) return 0;
int64_t n = std::strtoll(sv.ptr, nullptr, 10);
if (n < 1) return 0;
if (n > static_cast<int64_t>(table->column_count)) {
return static_cast<uint16_t>(table->column_count - 1);
}
return static_cast<uint16_t>(n - 1);
}
sql_parser::StringRef col_name;
if (key->type == sql_parser::NodeType::NODE_COLUMN_REF ||
key->type == sql_parser::NodeType::NODE_IDENTIFIER) {
col_name = key->value();
} else if (key->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
// table.column -- get the column part
const sql_parser::AstNode* c = key->first_child;
if (c && c->next_sibling) col_name = c->next_sibling->value();
else if (c) col_name = c->value();
Expand Down
2 changes: 1 addition & 1 deletion include/sql_engine/plan_node.h
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ struct PlanNode {
struct {
const char* backend_name;
const char* remote_sql;
uint16_t remote_sql_len;
uint32_t remote_sql_len;
const TableInfo* table; // expected result schema (for SELECT *)
// Optional projection expressions used to derive result column
// names when the remote SQL is not a passthrough SELECT *. When
Expand Down
22 changes: 22 additions & 0 deletions include/sql_engine/remote_query_builder.h
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,28 @@ class RemoteQueryBuilder {
return sb.finish();
}

sql_parser::StringRef build_select_join(
const TableInfo* left,
const TableInfo* right,
const sql_parser::AstNode* on_expr,
const sql_parser::AstNode* where_expr)
{
sql_parser::StringBuilder sb(arena_, 512);
sb.append("SELECT * FROM ");
if (left) sb.append(left->table_name.ptr, left->table_name.len);
sb.append(" JOIN ");
if (right) sb.append(right->table_name.ptr, right->table_name.len);
if (on_expr) {
sb.append(" ON ");
emit_expr(on_expr, sb);
}
if (where_expr) {
sb.append(" WHERE ");
emit_expr(where_expr, sb);
}
return sb.finish();
}

// Build an INSERT statement string.
sql_parser::StringRef build_insert(
const TableInfo* table,
Expand Down
Loading