fix: distribute derived tables and stop silent sort/SQL bugs

Rewrite DERIVED_SCAN inner plans so FROM (SELECT ...) hits remote
shards. Store remote_sql_len as uint32_t so statements longer than
64KB are not truncated. Resolve ORDER BY position/alias in the plan
builder, and only MERGE_SORT when every key is a table column —
expressions gather and sort locally instead of comparing column 0.

Author: renecannao

include/sql_engine/distributed_planner.h

@@ -206,6 +206,13 @@ class DistributedPlanner {
                 return result;
             }
 
+            case PlanNodeType::DERIVED_SCAN: {
+                PlanNode* result = make_plan_node(arena_, PlanNodeType::DERIVED_SCAN);
+                result->derived_scan = node->derived_scan;
+                result->derived_scan.inner_plan = distribute(node->derived_scan.inner_plan);
+                return result;
+            }
+
             case PlanNodeType::SET_OP: {
                 PlanNode* result = make_plan_node(arena_, PlanNodeType::SET_OP);
                 result->set_op = node->set_op;
@@ -293,6 +300,9 @@ class DistributedPlanner {
                 if (contains_type(node->merge_sort.children[i], type)) return true;
             }
         }
+        if (node->type == PlanNodeType::DERIVED_SCAN) {
+            return contains_type(node->derived_scan.inner_plan, type);
+        }
         return false;
     }
 
@@ -545,7 +555,7 @@ class DistributedPlanner {
         std::memcpy(bn, backend, blen + 1);
         node->remote_scan.backend_name = bn;
         node->remote_scan.remote_sql = sql.ptr;
-        node->remote_scan.remote_sql_len = static_cast<uint16_t>(sql.len);
+        node->remote_scan.remote_sql_len = sql.len;
         node->remote_scan.table = table;
         // Caller is responsible for setting output_exprs when the remote SQL
         // is not a passthrough SELECT *. make_plan_node() already zero-fills
@@ -782,15 +792,53 @@ class DistributedPlanner {
     }
 
     // Case 4: Distributed sort + limit
+    PlanNode* local_sort(PlanNode* sort_node) {
+        PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT);
+        result->sort = sort_node->sort;
+        result->left = distribute_node(sort_node->left);
+        return result;
+    }
+
+    int sort_key_table_ordinal(const sql_parser::AstNode* key, const TableInfo* table) const {
+        if (!key || !table) return -1;
+        if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = key->value();
+            if (!sv.ptr || sv.len == 0) return -1;
+            int64_t n = std::strtoll(sv.ptr, nullptr, 10);
+            if (n < 1 || n > static_cast<int64_t>(table->column_count)) return -1;
+            return static_cast<int>(n - 1);
+        }
+        sql_parser::StringRef col_name;
+        if (key->type == sql_parser::NodeType::NODE_COLUMN_REF ||
+            key->type == sql_parser::NodeType::NODE_IDENTIFIER) {
+            col_name = key->value();
+        } else if (key->type == sql_parser::NodeType::NODE_QUALIFIED_NAME) {
+            const sql_parser::AstNode* c = key->first_child;
+            if (c && c->next_sibling) col_name = c->next_sibling->value();
+            else if (c) col_name = c->value();
+        } else {
+            return -1;
+        }
+        if (!col_name.ptr) return -1;
+        const ColumnInfo* col = catalog_.get_column(table, col_name);
+        if (!col) return -1;
+        return static_cast<int>(col->ordinal);
+    }
+
+    bool all_sort_keys_are_table_columns(const PlanNode* sort_node, const TableInfo* table) const {
+        if (!sort_node || !table) return false;
+        for (uint16_t i = 0; i < sort_node->sort.count; ++i) {
+            if (sort_key_table_ordinal(sort_node->sort.keys[i], table) < 0) return false;
+        }
+        return true;
+    }
+
     PlanNode* distribute_sort(PlanNode* sort_node) {
         if (contains_type(sort_node->left, PlanNodeType::WINDOW) ||
             contains_type(sort_node->left, PlanNodeType::DERIVED_SCAN) ||
             contains_type(sort_node->left, PlanNodeType::AGGREGATE) ||
             contains_type(sort_node->left, PlanNodeType::MERGE_AGGREGATE)) {
-            PlanNode* result = make_plan_node(arena_, PlanNodeType::SORT);
-            result->sort = sort_node->sort;
-            result->left = distribute_node(sort_node->left);
-            return result;
+            return local_sort(sort_node);
         }
 
         ScanContext ctx = extract_scan_context(sort_node->left);
@@ -809,6 +857,10 @@ class DistributedPlanner {
             return result;
         }
 
+        if (!all_sort_keys_are_table_columns(sort_node, table)) {
+            return local_sort(sort_node);
+        }
+
         if (!shards_.is_sharded(table->table_name)) {
             // Unsharded -- push sort to remote
             sql_parser::StringRef sql = qb_.build_select(
@@ -875,8 +927,12 @@ class DistributedPlanner {
                 const TableInfo* table = ctx.scan->scan.table;
                 if (shards_.has_table(table->table_name) &&
                     shards_.is_sharded(table->table_name)) {
-                    // Case 4: Sharded sort + limit
-                    // Each shard: ORDER BY + LIMIT, MergeSort, then outer Limit
+                    if (!all_sort_keys_are_table_columns(sort_node, table)) {
+                        PlanNode* result = make_plan_node(arena_, PlanNodeType::LIMIT);
+                        result->limit = limit_node->limit;
+                        result->left = distribute_node(limit_node->left);
+                        return result;
+                    }
                     int64_t remote_limit = limit_node->limit.count + limit_node->limit.offset;
 
                     PlanNode* merge = make_sharded_merge_sort(

include/sql_engine/plan_builder.h

@@ -25,6 +25,7 @@
 #include "sql_parser/common.h"
 #include "sql_parser/arena.h"
 #include <cstring>
+#include <cstdlib>
 #include <vector>
 
 namespace sql_engine {
@@ -116,6 +117,55 @@ class PlanBuilder {
         return false;
     }
 
+    static const sql_parser::AstNode* select_item_expr(const sql_parser::AstNode* item) {
+        return item ? item->first_child : nullptr;
+    }
+
+    static sql_parser::StringRef select_item_alias(const sql_parser::AstNode* item) {
+        if (!item) return {};
+        for (const sql_parser::AstNode* c = item->first_child; c; c = c->next_sibling) {
+            if (c->type == sql_parser::NodeType::NODE_ALIAS) return c->value();
+        }
+        return {};
+    }
+
+    static const sql_parser::AstNode* resolve_order_key(
+            const sql_parser::AstNode* key, const sql_parser::AstNode* select_items) {
+        if (!key || !select_items) return key;
+        uint16_t n = count_children(select_items);
+        if (n == 0) return key;
+
+        if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = key->value();
+            if (!sv.ptr || sv.len == 0) return key;
+            int64_t pos = std::strtoll(sv.ptr, nullptr, 10);
+            if (pos < 1 || pos > static_cast<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.
@@ -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);

include/sql_engine/plan_executor.h

@@ -56,6 +56,7 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <cstdlib>
 
 namespace sql_engine {
 
@@ -1170,12 +1171,21 @@ class PlanExecutor {
 
     uint16_t resolve_column_index(const sql_parser::AstNode* key, const TableInfo* table) {
         if (!key || !table) return 0;
+        if (key->type == sql_parser::NodeType::NODE_LITERAL_INT) {
+            sql_parser::StringRef sv = key->value();
+            if (!sv.ptr || sv.len == 0) return 0;
+            int64_t n = std::strtoll(sv.ptr, nullptr, 10);
+            if (n < 1) return 0;
+            if (n > static_cast<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();

include/sql_engine/plan_node.h

@@ -101,7 +101,7 @@ struct PlanNode {
         struct {
             const char* backend_name;
             const char* remote_sql;
-            uint16_t remote_sql_len;
+            uint32_t remote_sql_len;
             const TableInfo* table;       // expected result schema (for SELECT *)
             // Optional projection expressions used to derive result column
             // names when the remote SQL is not a passthrough SELECT *. When

tests/test_distributed_planner.cpp

@@ -390,6 +390,10 @@ class DistributedPlannerTest : public ::testing::Test {
                 find_nodes(node->merge_sort.children[i], type, out);
             return;
         }
+        if (node->type == PlanNodeType::DERIVED_SCAN) {
+            find_nodes(node->derived_scan.inner_plan, type, out);
+            return;
+        }
         find_nodes(node->left, type, out);
         find_nodes(node->right, type, out);
     }
@@ -964,3 +968,96 @@ TEST_F(DistributedPlannerTest, WindowGatherCorrectness) {
     EXPECT_EQ(dist_rs.row_count(), 15u);
     EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs));
 }
+
+TEST_F(DistributedPlannerTest, DerivedScanIsDistributed) {
+    Parser<Dialect::MySQL> parser;
+    const char* sql = "SELECT name FROM (SELECT name, age FROM users WHERE age > 20) AS t";
+    auto pr = parser.parse(sql, std::strlen(sql));
+    ASSERT_EQ(pr.status, ParseResult::OK);
+
+    PlanBuilder<Dialect::MySQL> builder(catalog, parser.arena());
+    PlanNode* plan = builder.build(pr.ast);
+    ASSERT_NE(plan, nullptr);
+
+    DistributedPlanner<Dialect::MySQL> dp(shard_map, catalog, parser.arena());
+    PlanNode* dist = dp.distribute(plan);
+    ASSERT_NE(dist, nullptr);
+
+    std::vector<PlanNode*> scans, remotes;
+    find_nodes(dist, PlanNodeType::SCAN, scans);
+    find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes);
+    EXPECT_TRUE(scans.empty()) << "inner SCAN must be rewritten";
+    EXPECT_FALSE(remotes.empty());
+}
+
+TEST_F(DistributedPlannerTest, DerivedScanCorrectness) {
+    const char* sql = "SELECT name FROM (SELECT name, age FROM users WHERE age > 20) AS t";
+    auto local_rs = execute_local(sql);
+    auto dist_rs = execute_distributed(sql);
+    EXPECT_GT(local_rs.row_count(), 0u);
+    EXPECT_EQ(local_rs.row_count(), dist_rs.row_count());
+    EXPECT_TRUE(compare_results_unordered(local_rs, dist_rs));
+}
+
+TEST_F(DistributedPlannerTest, OrderByExpressionDoesNotMergeOnColumnZero) {
+    Parser<Dialect::MySQL> parser;
+    const char* sql = "SELECT name, age FROM users ORDER BY age + 1";
+    auto pr = parser.parse(sql, std::strlen(sql));
+    ASSERT_EQ(pr.status, ParseResult::OK);
+
+    PlanBuilder<Dialect::MySQL> builder(catalog, parser.arena());
+    PlanNode* plan = builder.build(pr.ast);
+    ASSERT_NE(plan, nullptr);
+
+    DistributedPlanner<Dialect::MySQL> dp(shard_map, catalog, parser.arena());
+    PlanNode* dist = dp.distribute(plan);
+    ASSERT_NE(dist, nullptr);
+
+    std::vector<PlanNode*> merges;
+    find_nodes(dist, PlanNodeType::MERGE_SORT, merges);
+    EXPECT_TRUE(merges.empty()) << "expression ORDER BY must not MERGE_SORT on col 0";
+
+    auto local_rs = execute_local(sql);
+    auto dist_rs = execute_distributed(sql);
+    EXPECT_EQ(local_rs.row_count(), dist_rs.row_count());
+    EXPECT_TRUE(compare_results_ordered(local_rs, dist_rs));
+}
+
+TEST_F(DistributedPlannerTest, OrderByAliasAndPositionCorrectness) {
+    const char* sql = "SELECT name AS n, age AS a FROM users ORDER BY a DESC";
+    auto local_rs = execute_local(sql);
+    auto dist_rs = execute_distributed(sql);
+    EXPECT_EQ(local_rs.row_count(), 15u);
+    EXPECT_TRUE(compare_results_ordered(local_rs, dist_rs));
+
+    const char* sql2 = "SELECT name, age FROM users ORDER BY 2 DESC";
+    auto local2 = execute_local(sql2);
+    auto dist2 = execute_distributed(sql2);
+    EXPECT_TRUE(compare_results_ordered(local2, dist2));
+}
+
+TEST_F(DistributedPlannerTest, RemoteSqlLenIsNotTruncated) {
+    std::string name(70000, 'x');
+    std::string sql = "SELECT * FROM users WHERE name = '" + name + "'";
+    Parser<Dialect::MySQL> parser;
+    auto pr = parser.parse(sql.c_str(), sql.size());
+    ASSERT_EQ(pr.status, ParseResult::OK);
+
+    PlanBuilder<Dialect::MySQL> builder(catalog, parser.arena());
+    PlanNode* plan = builder.build(pr.ast);
+    ASSERT_NE(plan, nullptr);
+
+    DistributedPlanner<Dialect::MySQL> dp(shard_map, catalog, parser.arena());
+    PlanNode* dist = dp.distribute(plan);
+    ASSERT_NE(dist, nullptr);
+
+    std::vector<PlanNode*> remotes;
+    find_nodes(dist, PlanNodeType::REMOTE_SCAN, remotes);
+    ASSERT_FALSE(remotes.empty());
+    for (auto* rs : remotes) {
+        EXPECT_GT(rs->remote_scan.remote_sql_len, 65535u);
+        std::string remote(rs->remote_scan.remote_sql, rs->remote_scan.remote_sql_len);
+        EXPECT_NE(remote.find(name), std::string::npos);
+        EXPECT_EQ(rs->remote_scan.remote_sql_len, remote.size());
+    }
+}

tests/test_plan_executor.cpp

@@ -147,6 +147,16 @@ TEST_F(PlanExecutorTest, SelectDistinctDept) {
     EXPECT_EQ(depts.size(), 2u);
 }
 
+TEST_F(PlanExecutorTest, OrderByAliasAndPosition) {
+    auto by_alias = run_query("SELECT name AS n, age AS a FROM users ORDER BY a DESC");
+    ASSERT_EQ(by_alias.row_count(), 5u);
+    EXPECT_EQ(by_alias.rows[0].get(1).int_val, 35);
+
+    auto by_pos = run_query("SELECT name, age FROM users ORDER BY 2 DESC");
+    ASSERT_EQ(by_pos.row_count(), 5u);
+    EXPECT_EQ(by_pos.rows[0].get(1).int_val, 35);
+}
+
 TEST_F(PlanExecutorTest, CountDistinctDept) {
     parser.reset();
     const char* sql = "SELECT COUNT(DISTINCT dept) FROM users";