Fix joins, rename, union

Co-authored-by: Marcos Magueta <maguetamarcos@gmail.com>

2 files changed, 195 insertions, 5 deletions

diff --git a/src/RNT_C_API.cpp b/src/RNT_C_API.cpp
index abd683c..a632ab9 100644
--- a/src/RNT_C_API.cpp
+++ b/src/RNT_C_API.cpp
@@ -503,6 +503,10 @@ namespace {
         node->left = l->root;
         node->right = r->root;
 
+        if (a.attrs)
+            for (const char** p = a.attrs; *p; ++p)
+                node->join_attrs.emplace(*p);
+
         auto* pw = new PlanWrapper();
         pw->root = node.get();
         pw->nodes.push_back(std::move(node));
@@ -577,6 +581,85 @@ namespace {
         delete s;
         return pw;
     }
+
+    PlanWrapper* build_materialize(const PlanArgsMaterialize& a) {
+        auto* s = static_cast<PlanWrapper*>(a.source);
+        if (!s)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_MATERIALIZE;
+        node->left = s->root;
+
+        auto* pw = new PlanWrapper();
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+        for (auto& n : s->nodes)
+            pw->nodes.push_back(std::move(n));
+        for (auto* c : s->cursors)
+            pw->cursors.push_back(c);
+        for (auto* h : s->handles)
+            pw->handles.push_back(h);
+
+        delete s;
+        return pw;
+    }
+
+    PlanWrapper* build_rename(const PlanArgsRename& a) {
+        auto* s = static_cast<PlanWrapper*>(a.source);
+        if (!s)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_RENAME;
+        node->upstream = s->root;
+
+        for (const char** k = a.pairs; *k; k++)
+            node->attrs[*k] = *++k;
+
+        auto* pw = new PlanWrapper();
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+
+        for (auto& n : s->nodes)
+            pw->nodes.push_back(std::move(n));
+        for (auto* c : s->cursors)
+            pw->cursors.push_back(c);
+        for (auto* h : s->handles)
+            pw->handles.push_back(h);
+
+        delete s;
+        return pw;
+    }
+
+    PlanWrapper* build_union(const PlanArgsUnion& a) {
+        if (!*a.sources)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_UNION;
+
+        auto* pw = new PlanWrapper();
+        for (auto* s = a.sources; *s; s++) {
+            auto* source = static_cast<PlanWrapper*>(*s);
+
+            node->nodes.push_back(source->root);
+
+            for (auto& n : source->nodes)
+                pw->nodes.push_back(std::move(n));
+            for (auto* c : source->cursors)
+                pw->cursors.push_back(std::move(c));
+            for (auto* h : source->handles)
+                pw->handles.push_back(std::move(h));
+
+            delete source;
+        }
+
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+
+        return pw;
+    }
 } // namespace
 
 // ---------------------------------------------------------------------------
@@ -1019,6 +1102,12 @@ rnt_plan_t rnt_plan_assemble(PlanAction action) {
         return build_take(action.take);
     case nt::FOL_OPERATION_PROJECT:
         return build_project(action.project);
+    case nt::FOL_OPERATION_MATERIALIZE:
+        return build_materialize(action.materialize);
+    case nt::FOL_OPERATION_RENAME:
+        return build_rename(action.rename);
+    case nt::FOL_OPERATION_UNION:
+        return build_union(action.union_);
     }
 
     return nullptr;
diff --git a/src/VM.cpp b/src/VM.cpp
index 5f66f83..1f509ff 100644
--- a/src/VM.cpp
+++ b/src/VM.cpp
@@ -32,6 +32,49 @@ namespace nt {
         c->exhausted = false;
     }
 
+    void VM::Rewind(PlanNode* node, Tuple* outer) {
+        if (node == nullptr)
+            return;
+
+        switch (node->op) {
+        case Operation::FOL_OPERATION_SCAN:
+            ResetInner(node->scan_cursor, ResolveArgs(node->scan_args, outer));
+            break;
+
+        case Operation::FOL_OPERATION_PROJECT:
+            Rewind(node->left, outer);
+            break;
+
+        case Operation::FOL_OPERATION_TAKE:
+            node->take_count = 0;
+            Rewind(node->left, outer);
+            break;
+
+        case Operation::FOL_OPERATION_JOIN:
+            node->join_left = nullptr;
+            Rewind(node->left, outer);
+            Rewind(node->right, outer);
+            break;
+
+        case Operation::FOL_OPERATION_MATERIALIZE:
+            // Rewind the replay cursor only, and don't pass `outer` to the child.
+            // The point of this node is to reuse the cached tuples for every
+            // outer tuple, so it must not re-run or rebind what it cached.
+            node->mat_pos = 0;
+            break;
+
+        case Operation::FOL_OPERATION_RENAME:
+            Rewind(node->upstream, outer);
+            break;
+
+        case Operation::FOL_OPERATION_UNION:
+            node->current_node = node->nodes.begin();
+            for (auto& node : node->nodes)
+                Rewind(node, outer);
+            break;
+        }
+    }
+
     Tuple* VM::MergeInto(PlanNode* node, Tuple* left, Tuple* right) {
         std::vector<Attribute> merged;
         for (const auto& a : left->attrs())
@@ -67,18 +110,22 @@ namespace nt {
 
         case Operation::FOL_OPERATION_JOIN: {
             while (true) {
+            again:
                 if (node->join_left == nullptr) {
                     node->join_left = Next(node->left);
-                    if (node->join_left == nullptr)
-                        return nullptr;
+                    if (node->join_left == nullptr) return nullptr;
 
-                    auto args = ResolveArgs(node->right->scan_args, node->join_left);
-                    ResetInner(node->right->scan_cursor, std::move(args));
+                    Rewind(node->right, node->join_left);
                 }
 
                 Tuple* right = Next(node->right);
-                if (right != nullptr)
+                if (right != nullptr) {
+                    for (auto &attr : node->join_attrs)
+                        if ((*node->join_left)[attr] != (*right)[attr])
+                            goto again;
+
                     return MergeInto(node, node->join_left, right);
+                }
 
                 node->join_left = nullptr;
             }
@@ -92,6 +139,60 @@ namespace nt {
                 ++node->take_count;
             return t;
         }
+
+        case Operation::FOL_OPERATION_MATERIALIZE: {
+            if (!node->mat_done) {
+                // First scan: pull a tuple and keep our own copy. We rebuild it
+                // from attrs() rather than copying the Tuple, because a Tuple
+                // holds an iterator into its own map and a plain copy would leave
+                // that iterator dangling. The child also overwrites its output
+                // buffer on the next pull, so we cannot just keep its pointer.
+                Tuple* t = Next(node->left);
+                if (t != nullptr) {
+                    std::vector<Attribute> attrs;
+                    for (const auto& a : t->attrs())
+                        attrs.push_back(a);
+                    node->mat_buffer.emplace_back(std::move(attrs));
+                    return &node->mat_buffer.back();
+                }
+                node->mat_done = true;
+                return nullptr; // child exhausted; do not fall into replay
+            }
+
+            // Later scans: hand back cached tuples until the cache runs out.
+            if (node->mat_pos < node->mat_buffer.size())
+                return &node->mat_buffer[node->mat_pos++];
+            return nullptr;
+        }
+
+        case Operation::FOL_OPERATION_RENAME: {
+            Tuple* t = Next(node->upstream);
+            if (!t)
+                return nullptr;
+
+            std::vector<Attribute> attrs;
+            for (const auto& a : t->attrs()) {
+                std::string name = a.name;
+                if (node->attrs.contains(name))
+                    name = node->attrs[name];
+                Attribute new_attr { name, a.value };
+                attrs.push_back(new_attr);
+            }
+            delete t;
+            return new Tuple(attrs);
+        }
+
+        case Operation::FOL_OPERATION_UNION: {
+            if (node->current_node == node->nodes.end())
+                return nullptr;
+
+            Tuple* t = Next(*node->current_node);
+            if (!t) {
+                node->current_node++;
+                return Next(node);
+            }
+            return t;
+        }
         }
 
         return nullptr;