fix: resolve distributed subquery materialization, ORDER BY with PROJECT, derived table aggregate (#22, #23, #24)

Bug #22: DELETE/UPDATE parsers didn't set the subquery parse callback,
so subqueries in DML WHERE clauses (e.g., DELETE ... WHERE id IN
(SELECT ...)) were skipped instead of parsed. Enable callback on
DeleteParser and UpdateParser. Re-enable 2 disabled tests.

Bug #23: PlanBuilder placed Sort above Project, causing ORDER BY keys
to reference wrong column ordinals after projection. Move Sort before
Project in the plan tree so sorting operates on full rows. Update
preprocess_aggregates and distributed planner to look through SORT
when finding AGGREGATE nodes below PROJECT.

Bug #24: DerivedScanOperator now deep-copies inner query result rows
into the outer arena, preventing dangling pointers when inner arena
memory could be invalidated before the outer executor reads results.

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -129,9 +129,11 @@ class DistributedPlanner {
             }
 
             case PlanNodeType::PROJECT: {
-                // Check for PROJECT -> AGGREGATE pattern (or PROJECT -> FILTER -> AGGREGATE)
+                // Check for PROJECT -> [SORT ->] [FILTER ->] AGGREGATE pattern
                 // For aggregate queries, we want to handle the whole thing in distribute_aggregate
                 PlanNode* agg_child = node->left;
+                if (agg_child && agg_child->type == PlanNodeType::SORT)
+                    agg_child = agg_child->left;
                 if (agg_child && agg_child->type == PlanNodeType::FILTER)
                     agg_child = agg_child->left;
                 if (agg_child && agg_child->type == PlanNodeType::AGGREGATE) {

include/sql_engine/operators/derived_scan_op.h

@@ -2,22 +2,26 @@
 //
 // Materializes the inner plan's result set on open(), then yields
 // rows one at a time from the materialized buffer on next().
+// Deep-copies all row data into the outer arena to avoid dangling
+// pointers into arenas that may be reset after inner execution.
 
 #ifndef SQL_ENGINE_OPERATORS_DERIVED_SCAN_OP_H
 #define SQL_ENGINE_OPERATORS_DERIVED_SCAN_OP_H
 
 #include "sql_engine/operator.h"
 #include "sql_engine/row.h"
+#include "sql_parser/arena.h"
 #include <vector>
+#include <cstring>
 
 namespace sql_engine {
 
 class DerivedScanOperator : public Operator {
 public:
-    // Takes ownership of the inner operator. On open(), pulls all rows
-    // from the inner operator and stores them. On next(), yields them.
-    explicit DerivedScanOperator(Operator* inner)
-        : inner_(inner) {}
+    // Takes the inner operator and an arena for deep-copying row data.
+    // The arena must outlive any result set produced by the outer query.
+    DerivedScanOperator(Operator* inner, sql_parser::Arena& arena)
+        : inner_(inner), arena_(arena) {}
 
     void open() override {
         rows_.clear();
@@ -26,7 +30,7 @@ class DerivedScanOperator : public Operator {
             inner_->open();
             Row row{};
             while (inner_->next(row)) {
-                rows_.push_back(row);
+                rows_.push_back(deep_copy_row(row));
             }
             inner_->close();
         }
@@ -45,8 +49,35 @@ class DerivedScanOperator : public Operator {
 
 private:
     Operator* inner_;
+    sql_parser::Arena& arena_;
     std::vector<Row> rows_;
     size_t cursor_ = 0;
+
+    Row deep_copy_row(const Row& src) {
+        if (!src.values || src.column_count == 0) {
+            Row result;
+            result.values = nullptr;
+            result.column_count = 0;
+            return result;
+        }
+        uint16_t cc = src.column_count;
+        Row result = make_row(arena_, cc);
+        for (uint16_t i = 0; i < cc; ++i) {
+            Value v = src.values[i];
+            // Deep-copy string data into the arena
+            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* buf = static_cast<char*>(arena_.allocate(v.str_val.len));
+                std::memcpy(buf, v.str_val.ptr, v.str_val.len);
+                v.str_val.ptr = buf;
+            }
+            result.set(i, v);
+        }
+        return result;
+    }
 };
 
 } // namespace sql_engine

include/sql_engine/plan_builder.h

@@ -206,7 +206,37 @@ class PlanBuilder {
             current = filter;
         }
 
-        // 5. SELECT list -> Project
+        // 5. ORDER BY -> Sort (before Project so sort keys resolve against full row)
+        const sql_parser::AstNode* order_by = find_child(select_ast, sql_parser::NodeType::NODE_ORDER_BY_CLAUSE);
+        if (order_by) {
+            PlanNode* sort = make_plan_node(arena_, PlanNodeType::SORT);
+            uint16_t cnt = count_children(order_by);
+            sort->sort.count = cnt;
+
+            auto** keys = static_cast<const sql_parser::AstNode**>(
+                arena_.allocate(sizeof(sql_parser::AstNode*) * cnt));
+            auto* dirs = static_cast<uint8_t*>(arena_.allocate(cnt));
+
+            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;
+                // 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);
+                if (dir_node) {
+                    sql_parser::StringRef dir_val = dir_node->value();
+                    if (dir_val.equals_ci("DESC", 4)) dirs[idx] = 1;
+                }
+                ++idx;
+            }
+            sort->sort.keys = keys;
+            sort->sort.directions = dirs;
+            sort->left = current;
+            current = sort;
+        }
+
+        // 6. SELECT list -> Project
         const sql_parser::AstNode* item_list = find_child(select_ast, sql_parser::NodeType::NODE_SELECT_ITEM_LIST);
         if (item_list) {
             // Check if this is "SELECT *" with a single asterisk and no aliases -- skip Project for bare scan
@@ -246,43 +276,13 @@ class PlanBuilder {
             }
         }
 
-        // 6. DISTINCT -> Distinct
+        // 7. DISTINCT -> Distinct
         if (has_distinct(select_ast)) {
             PlanNode* dist = make_plan_node(arena_, PlanNodeType::DISTINCT);
             dist->left = current;
             current = dist;
         }
 
-        // 7. ORDER BY -> Sort
-        const sql_parser::AstNode* order_by = find_child(select_ast, sql_parser::NodeType::NODE_ORDER_BY_CLAUSE);
-        if (order_by) {
-            PlanNode* sort = make_plan_node(arena_, PlanNodeType::SORT);
-            uint16_t cnt = count_children(order_by);
-            sort->sort.count = cnt;
-
-            auto** keys = static_cast<const sql_parser::AstNode**>(
-                arena_.allocate(sizeof(sql_parser::AstNode*) * cnt));
-            auto* dirs = static_cast<uint8_t*>(arena_.allocate(cnt));
-
-            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;
-                // 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);
-                if (dir_node) {
-                    sql_parser::StringRef dir_val = dir_node->value();
-                    if (dir_val.equals_ci("DESC", 4)) dirs[idx] = 1;
-                }
-                ++idx;
-            }
-            sort->sort.keys = keys;
-            sort->sort.directions = dirs;
-            sort->left = current;
-            current = sort;
-        }
-
         // 8. LIMIT -> Limit
         const sql_parser::AstNode* limit_clause = find_child(select_ast, sql_parser::NodeType::NODE_LIMIT_CLAUSE);
         if (limit_clause) {

include/sql_engine/plan_executor.h

@@ -390,8 +390,10 @@ class PlanExecutor {
 
         if (node->type != PlanNodeType::PROJECT) return;
 
-        // Find AGGREGATE child (possibly through a FILTER for HAVING)
+        // Find AGGREGATE child (possibly through SORT and/or FILTER for HAVING)
         PlanNode* agg_node = node->left;
+        if (agg_node && agg_node->type == PlanNodeType::SORT)
+            agg_node = agg_node->left;
         if (agg_node && agg_node->type == PlanNodeType::FILTER)
             agg_node = agg_node->left;
         if (!agg_node || agg_node->type != PlanNodeType::AGGREGATE) return;
@@ -435,9 +437,17 @@ class PlanExecutor {
         // output matches what we need.
 
         // Overwrite this PROJECT node to become a pass-through:
-        // Change it into the node below it.
+        // Change it into the node below it, preserving any intermediate
+        // SORT and/or FILTER (HAVING) nodes.
         PlanNode* child = node->left;
-        if (child && child->type == PlanNodeType::FILTER) {
+        if (child && child->type == PlanNodeType::SORT) {
+            // PROJECT -> SORT -> [FILTER ->] AGGREGATE
+            // Replace PROJECT with SORT (preserving the sort)
+            node->type = PlanNodeType::SORT;
+            node->sort = child->sort;
+            node->left = child->left;
+            node->right = child->right;
+        } else if (child && child->type == PlanNodeType::FILTER) {
             // PROJECT -> FILTER -> AGGREGATE: keep FILTER, remove PROJECT
             node->type = PlanNodeType::FILTER;
             node->filter.expr = child->filter.expr;
@@ -576,10 +586,11 @@ class PlanExecutor {
     }
 
     Operator* build_derived_scan_op(PlanNode* node) {
-        // Build the inner plan's operator tree and wrap in DerivedScanOperator
+        // Build the inner plan's operator tree and wrap in DerivedScanOperator.
+        // Pass the arena so deep-copied row data persists in the outer arena.
         Operator* inner = build_operator(node->derived_scan.inner_plan);
         if (!inner) return nullptr;
-        auto op = std::make_unique<DerivedScanOperator>(inner);
+        auto op = std::make_unique<DerivedScanOperator>(inner, arena_);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));
         return ptr;

include/sql_parser/delete_parser.h

@@ -22,6 +22,10 @@ class DeleteParser {
         : tok_(tokenizer), arena_(arena), expr_parser_(tokenizer, arena),
           table_ref_parser_(tokenizer, arena, expr_parser_) {}
 
+    void set_subquery_callback(SubqueryParseCallback<D> cb) {
+        expr_parser_.set_subquery_callback(cb);
+    }
+
     // Parse DELETE statement (DELETE keyword already consumed).
     AstNode* parse() {
         AstNode* root = make_node(arena_, NodeType::NODE_DELETE_STMT);

include/sql_parser/update_parser.h

@@ -18,6 +18,10 @@ class UpdateParser {
         : tok_(tokenizer), arena_(arena), expr_parser_(tokenizer, arena),
           table_ref_parser_(tokenizer, arena, expr_parser_) {}
 
+    void set_subquery_callback(SubqueryParseCallback<D> cb) {
+        expr_parser_.set_subquery_callback(cb);
+    }
+
     // Parse UPDATE statement (UPDATE keyword already consumed).
     AstNode* parse() {
         AstNode* root = make_node(arena_, NodeType::NODE_UPDATE_STMT);

src/sql_parser/parser.cpp

@@ -347,6 +347,7 @@ ParseResult Parser<D>::parse_update() {
     r.stmt_type = StmtType::UPDATE;
 
     UpdateParser<D> update_parser(tokenizer_, arena_);
+    update_parser.set_subquery_callback(&parse_subquery_select<D>);
     AstNode* ast = update_parser.parse();
 
     if (ast) {
@@ -382,6 +383,7 @@ ParseResult Parser<D>::parse_delete() {
     r.stmt_type = StmtType::DELETE_STMT;
 
     DeleteParser<D> delete_parser(tokenizer_, arena_);
+    delete_parser.set_subquery_callback(&parse_subquery_select<D>);
     AstNode* ast = delete_parser.parse();
 
     if (ast) {

tests/test_distributed_dml.cpp

@@ -524,7 +524,7 @@ TEST_F(DistributedDmlTest, InsertThenSelectVerify) {
 
 // DELETE FROM users WHERE id IN (SELECT user_id FROM orders)
 // users on shards, orders on orders_backend
-TEST_F(DistributedDmlTest, DISABLED_DmlWithCrossShardSubquery) {
+TEST_F(DistributedDmlTest, DmlWithCrossShardSubquery) {
     // Insert users: id=0 -> shard0, id=1 -> shard1, id=2 -> shard2
     execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (0, 'Alice', 25)");
     execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (1, 'Bob', 30)");
@@ -568,7 +568,7 @@ TEST_F(DistributedDmlTest, DISABLED_DmlWithCrossShardSubquery) {
 }
 
 // UPDATE users SET age = 99 WHERE id IN (SELECT user_id FROM orders)
-TEST_F(DistributedDmlTest, DISABLED_UpdateWithCrossShardSubquery) {
+TEST_F(DistributedDmlTest, UpdateWithCrossShardSubquery) {
     // Insert users
     execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (0, 'Alice', 25)");
     execute_distributed_dml("INSERT INTO users (id, name, age) VALUES (1, 'Bob', 30)");

tests/test_plan_executor.cpp

@@ -155,3 +155,34 @@ TEST_F(PlanExecutorTest, SelectWithLike) {
     std::string name(rs.rows[0].get(0).str_val.ptr, rs.rows[0].get(0).str_val.len);
     EXPECT_EQ(name, "Alice");
 }
+
+// Bug #23: SELECT name FROM users ORDER BY age DESC -- sort by age, return only name
+TEST_F(PlanExecutorTest, OrderByColumnNotInSelectList) {
+    auto rs = run_query("SELECT name FROM users ORDER BY age DESC");
+    EXPECT_EQ(rs.row_count(), 5u);
+    EXPECT_EQ(rs.column_count, 1);
+    // Sorted by age DESC: Eve(35), Bob(30), Alice(25), Dave(22), Carol(17)
+    std::string first(rs.rows[0].get(0).str_val.ptr, rs.rows[0].get(0).str_val.len);
+    std::string last(rs.rows[4].get(0).str_val.ptr, rs.rows[4].get(0).str_val.len);
+    EXPECT_EQ(first, "Eve");
+    EXPECT_EQ(last, "Carol");
+}
+
+// Bug #24: SELECT * FROM (SELECT COUNT(*) AS cnt FROM users) AS t
+TEST_F(PlanExecutorTest, DerivedTableWithAggregate) {
+    auto rs = run_query("SELECT * FROM (SELECT COUNT(*) AS cnt FROM users) AS t");
+    EXPECT_EQ(rs.row_count(), 1u);
+    EXPECT_FALSE(rs.rows[0].get(0).is_null());
+    EXPECT_EQ(rs.rows[0].get(0).int_val, 5);
+}
+
+// Bug #24: SELECT * FROM (SELECT dept, COUNT(*) AS cnt FROM users GROUP BY dept) AS t
+TEST_F(PlanExecutorTest, DerivedTableWithGroupByAggregate) {
+    auto rs = run_query("SELECT * FROM (SELECT dept, COUNT(*) AS cnt FROM users GROUP BY dept) AS t");
+    EXPECT_EQ(rs.row_count(), 2u);
+    int64_t total = 0;
+    for (const auto& row : rs.rows) {
+        total += row.get(1).int_val;
+    }
+    EXPECT_EQ(total, 5); // 3 Engineering + 2 Sales
+}