From 83e732e19d55a1572c00ca27ecf15cfcbafe6f22 Mon Sep 17 00:00:00 2001
From: David Sisson <EpsilonPrime@users.noreply.github.com>
Date: Wed, 14 Feb 2024 01:06:59 -0800
Subject: [PATCH] Added to output mapping.

---
 .../converter/InitialPlanProtoVisitor.cpp     | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/substrait/textplan/converter/InitialPlanProtoVisitor.cpp b/src/substrait/textplan/converter/InitialPlanProtoVisitor.cpp
index 2b783996..d044acc6 100644
--- a/src/substrait/textplan/converter/InitialPlanProtoVisitor.cpp
+++ b/src/substrait/textplan/converter/InitialPlanProtoVisitor.cpp
@@ -68,6 +68,14 @@ void eraseInputs(::substrait::proto::Rel* relation) {
       relation->mutable_cross()->clear_left();
       relation->mutable_cross()->clear_right();
       break;
+    case ::substrait::proto::Rel::kReference:
+      break;
+    case ::substrait::proto::Rel::kWrite:
+      relation->mutable_write()->clear_input();
+      break;
+    case ::substrait::proto::Rel::kDdl:
+      relation->mutable_ddl()->clear_view_definition();
+      break;
     case ::substrait::proto::Rel::kHashJoin:
       relation->mutable_hash_join()->clear_left();
       relation->mutable_hash_join()->clear_right();
@@ -76,6 +84,10 @@ void eraseInputs(::substrait::proto::Rel* relation) {
       relation->mutable_merge_join()->clear_left();
       relation->mutable_merge_join()->clear_right();
       break;
+    case ::substrait::proto::Rel::kNestedLoopJoin:
+      relation->mutable_nested_loop_join()->clear_left();
+      relation->mutable_nested_loop_join()->clear_right();
+      break;
     case ::substrait::proto::Rel::kWindow:
       relation->mutable_window()->clear_input();
       break;
@@ -117,10 +129,19 @@ ::google::protobuf::RepeatedField<int32_t> getOutputMapping(
       return relation.extension_leaf().common().emit().output_mapping();
     case ::substrait::proto::Rel::kCross:
       return relation.cross().common().emit().output_mapping();
+    case ::substrait::proto::Rel::kReference:
+      // There is no common message in a ReferenceRel.
+      break;
+    case ::substrait::proto::Rel::kWrite:
+      return relation.write().common().emit().output_mapping();
+    case ::substrait::proto::Rel::kDdl:
+      return relation.ddl().common().emit().output_mapping();
     case ::substrait::proto::Rel::kHashJoin:
       return relation.hash_join().common().emit().output_mapping();
     case ::substrait::proto::Rel::kMergeJoin:
       return relation.merge_join().common().emit().output_mapping();
+    case ::substrait::proto::Rel::kNestedLoopJoin:
+      return relation.nested_loop_join().common().emit().output_mapping();
     case ::substrait::proto::Rel::kWindow:
       return relation.window().common().emit().output_mapping();
     case ::substrait::proto::Rel::kExchange:
@@ -566,6 +587,14 @@ void InitialPlanProtoVisitor::updateLocalSchema(
       addFieldsToRelation(
           relationData, relation.cross().left(), relation.cross().right());
       break;
+    case ::substrait::proto::Rel::kReference:
+      break;
+    case ::substrait::proto::Rel::kWrite:
+      addFieldsToRelation(relationData, relation.write().input());
+      break;
+    case ::substrait::proto::Rel::kDdl:
+      addFieldsToRelation(relationData, relation.ddl().view_definition());
+      break;
     case ::substrait::proto::Rel::RelTypeCase::kHashJoin:
       addFieldsToRelation(
           relationData,
@@ -578,6 +607,12 @@ void InitialPlanProtoVisitor::updateLocalSchema(
           relation.merge_join().left(),
           relation.merge_join().right());
       break;
+    case ::substrait::proto::Rel::kNestedLoopJoin:
+      addFieldsToRelation(
+          relationData,
+          relation.nested_loop_join().left(),
+          relation.nested_loop_join().right());
+      break;
     case ::substrait::proto::Rel::kWindow:
       addFieldsToRelation(relationData, relation.window().input());
       break;