From 5589e6d741d36a3859be3bdb9d72da8f7ce0a987 Mon Sep 17 00:00:00 2001
From: Amit Adhikari <amitadhikari@fb.com>
Date: Tue, 5 Apr 2022 14:03:09 -0700
Subject: [PATCH] Add missing MergeJoinNode visit methods

---
 .../scheduler/AllAtOnceExecutionSchedule.java        |  9 +++++++++
 .../presto/sql/planner/SplitSourceFactory.java       | 12 ++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/scheduler/AllAtOnceExecutionSchedule.java b/presto-main/src/main/java/com/facebook/presto/execution/scheduler/AllAtOnceExecutionSchedule.java
index 60948e6a10c49..f5b0b649c3f49 100644
--- a/presto-main/src/main/java/com/facebook/presto/execution/scheduler/AllAtOnceExecutionSchedule.java
+++ b/presto-main/src/main/java/com/facebook/presto/execution/scheduler/AllAtOnceExecutionSchedule.java
@@ -22,6 +22,7 @@
 import com.facebook.presto.sql.planner.plan.IndexJoinNode;
 import com.facebook.presto.sql.planner.plan.InternalPlanVisitor;
 import com.facebook.presto.sql.planner.plan.JoinNode;
+import com.facebook.presto.sql.planner.plan.MergeJoinNode;
 import com.facebook.presto.sql.planner.plan.PlanFragmentId;
 import com.facebook.presto.sql.planner.plan.RemoteSourceNode;
 import com.facebook.presto.sql.planner.plan.SemiJoinNode;
@@ -139,6 +140,14 @@ public Void visitJoin(JoinNode node, Void context)
             return null;
         }
 
+        @Override
+        public Void visitMergeJoin(MergeJoinNode node, Void context)
+        {
+            node.getRight().accept(this, context);
+            node.getLeft().accept(this, context);
+            return null;
+        }
+
         @Override
         public Void visitSemiJoin(SemiJoinNode node, Void context)
         {
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java
index 24560f59799b4..3a9f9f0aed294 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/SplitSourceFactory.java
@@ -45,6 +45,7 @@
 import com.facebook.presto.sql.planner.plan.IndexJoinNode;
 import com.facebook.presto.sql.planner.plan.InternalPlanVisitor;
 import com.facebook.presto.sql.planner.plan.JoinNode;
+import com.facebook.presto.sql.planner.plan.MergeJoinNode;
 import com.facebook.presto.sql.planner.plan.MetadataDeleteNode;
 import com.facebook.presto.sql.planner.plan.OutputNode;
 import com.facebook.presto.sql.planner.plan.RemoteSourceNode;
@@ -176,6 +177,17 @@ public Map<PlanNodeId, SplitSource> visitJoin(JoinNode node, Context context)
                     .build();
         }
 
+        @Override
+        public Map<PlanNodeId, SplitSource> visitMergeJoin(MergeJoinNode node, Context context)
+        {
+            Map<PlanNodeId, SplitSource> leftSplits = node.getLeft().accept(this, context);
+            Map<PlanNodeId, SplitSource> rightSplits = node.getRight().accept(this, context);
+            return ImmutableMap.<PlanNodeId, SplitSource>builder()
+                    .putAll(leftSplits)
+                    .putAll(rightSplits)
+                    .build();
+        }
+
         @Override
         public Map<PlanNodeId, SplitSource> visitSemiJoin(SemiJoinNode node, Context context)
         {