From 101b2beeffadfdf60931640193e61adac2e04503 Mon Sep 17 00:00:00 2001
From: Daniel Hunte <danielhunte@meta.com>
Date: Mon, 23 Dec 2024 12:07:28 -0800
Subject: [PATCH] feat(fuzzer): Update "tryFlipJoinSides" functions to handle
 multi-joins (#11938)

Summary:

The function should traverse the plan tree and recursively the sides of all join nodes that are eligible to be flipped.

Differential Revision: D67606686
---
 velox/exec/fuzzer/JoinFuzzer.cpp | 65 +++++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 14 deletions(-)

diff --git a/velox/exec/fuzzer/JoinFuzzer.cpp b/velox/exec/fuzzer/JoinFuzzer.cpp
index d23b7a1cd72d..5d3f424f9e5c 100644
--- a/velox/exec/fuzzer/JoinFuzzer.cpp
+++ b/velox/exec/fuzzer/JoinFuzzer.cpp
@@ -606,7 +606,9 @@ std::optional<core::JoinType> tryFlipJoinType(core::JoinType joinType) {
 }
 
 // Returns a plan with flipped join sides of the input hash join node. If the
-// join type doesn't allow flipping, returns a nullptr.
+// inputs of the join node are other hash join nodes, recursively flip the join
+// sides of those join nodes as well. If the join type doesn't allow flipping,
+// returns a nullptr.
 core::PlanNodePtr tryFlipJoinSides(const core::HashJoinNode& joinNode) {
   //  Null-aware right semi project join doesn't support filter.
   if (joinNode.filter() &&
@@ -614,24 +616,36 @@ core::PlanNodePtr tryFlipJoinSides(const core::HashJoinNode& joinNode) {
       joinNode.isNullAware()) {
     return nullptr;
   }
+  core::PlanNodePtr left = joinNode.sources()[0];
+  core::PlanNodePtr right = joinNode.sources()[1];
+  if (auto leftJoinInput = std::dynamic_pointer_cast<const core::HashJoinNode>(
+          joinNode.sources()[0])) {
+    left = tryFlipJoinSides(*leftJoinInput);
+  }
+  if (auto rightJoinInput = std::dynamic_pointer_cast<const core::HashJoinNode>(
+          joinNode.sources()[1])) {
+    right = tryFlipJoinSides(*rightJoinInput);
+  }
   auto flippedJoinType = tryFlipJoinType(joinNode.joinType());
-  if (!flippedJoinType.has_value()) {
+  if (!flippedJoinType) {
     return nullptr;
   }
 
   return std::make_shared<core::HashJoinNode>(
       joinNode.id(),
-      flippedJoinType.value(),
+      *flippedJoinType,
       joinNode.isNullAware(),
       joinNode.rightKeys(),
       joinNode.leftKeys(),
       joinNode.filter(),
-      joinNode.sources()[1],
-      joinNode.sources()[0],
+      right,
+      left,
       joinNode.outputType());
 }
 
 // Returns a plan with flipped join sides of the input merge join node. If the
+// inputs of the join node are other merge join nodes, recursively flip the join
+// sides of those join nodes as well. If the
 // join type doesn't allow flipping, returns a nullptr.
 core::PlanNodePtr tryFlipJoinSides(const core::MergeJoinNode& joinNode) {
   // Merge join only supports inner and left join, so only inner join can be
@@ -639,33 +653,56 @@ core::PlanNodePtr tryFlipJoinSides(const core::MergeJoinNode& joinNode) {
   if (joinNode.joinType() != core::JoinType::kInner) {
     return nullptr;
   }
-  auto flippedJoinType = core::JoinType::kInner;
+  core::PlanNodePtr left = joinNode.sources()[0];
+  core::PlanNodePtr right = joinNode.sources()[1];
+  if (auto leftJoinInput = std::dynamic_pointer_cast<const core::MergeJoinNode>(
+          joinNode.sources()[0])) {
+    left = tryFlipJoinSides(*leftJoinInput);
+  }
+  if (auto rightJoinInput =
+          std::dynamic_pointer_cast<const core::MergeJoinNode>(
+              joinNode.sources()[1])) {
+    right = tryFlipJoinSides(*rightJoinInput);
+  }
 
   return std::make_shared<core::MergeJoinNode>(
       joinNode.id(),
-      flippedJoinType,
+      core::JoinType::kInner,
       joinNode.rightKeys(),
       joinNode.leftKeys(),
       joinNode.filter(),
-      joinNode.sources()[1],
-      joinNode.sources()[0],
+      right,
+      left,
       joinNode.outputType());
 }
 
 // Returns a plan with flipped join sides of the input nested loop join node. If
-// the join type doesn't allow flipping, returns a nullptr.
+// the inputs of the join node are other nested loop join nodes, recursively
+// flip the join sides of those join nodes as well. If the join type doesn't
+// allow flipping, returns a nullptr.
 core::PlanNodePtr tryFlipJoinSides(const core::NestedLoopJoinNode& joinNode) {
   auto flippedJoinType = tryFlipJoinType(joinNode.joinType());
-  if (!flippedJoinType.has_value()) {
+  if (!flippedJoinType) {
     return nullptr;
   }
-
+  core::PlanNodePtr left = joinNode.sources()[0];
+  core::PlanNodePtr right = joinNode.sources()[1];
+  if (auto leftJoinInput =
+          std::dynamic_pointer_cast<const core::NestedLoopJoinNode>(
+              joinNode.sources()[0])) {
+    left = tryFlipJoinSides(*leftJoinInput);
+  }
+  if (auto rightJoinInput =
+          std::dynamic_pointer_cast<const core::NestedLoopJoinNode>(
+              joinNode.sources()[1])) {
+    right = tryFlipJoinSides(*rightJoinInput);
+  }
   return std::make_shared<core::NestedLoopJoinNode>(
       joinNode.id(),
       flippedJoinType.value(),
       joinNode.joinCondition(),
-      joinNode.sources()[1],
-      joinNode.sources()[0],
+      right,
+      left,
       joinNode.outputType());
 }