From 6853991193b296a964b30564b62c8ad5071aeed4 Mon Sep 17 00:00:00 2001
From: "Xiaotian (Jackie) Jiang"
 <17555551+Jackie-Jiang@users.noreply.github.com>
Date: Fri, 27 Oct 2023 08:54:41 -0700
Subject: [PATCH] Fix multiple consecutive Exchange returning empty response
 (#11885)

---
 .../pinot/query/planner/logical/PlanFragmenter.java       | 2 ++
 .../src/test/resources/queries/WithStatements.json        | 8 ++++++++
 2 files changed, 10 insertions(+)

diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
index afc139fd99f4..fdb5858fdd89 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
@@ -128,6 +128,8 @@ public PlanNode visitExchange(ExchangeNode node, Context context) {
     }
     int currentPlanFragmentId = context._previousPlanFragmentId;
     int nextPlanFragmentId = ++context._currentPlanFragmentId;
+    // Set previous PlanFragment ID in the context to be the next PlanFragment ID to be used by the child node.
+    context._previousPlanFragmentId = nextPlanFragmentId;
     PlanNode nextPlanFragmentRoot = node.getInputs().get(0).visit(this, context);
 
     PinotRelExchangeType exchangeType = node.getExchangeType();
diff --git a/pinot-query-runtime/src/test/resources/queries/WithStatements.json b/pinot-query-runtime/src/test/resources/queries/WithStatements.json
index 401bf39e0bc8..0d8063e1c5fc 100644
--- a/pinot-query-runtime/src/test/resources/queries/WithStatements.json
+++ b/pinot-query-runtime/src/test/resources/queries/WithStatements.json
@@ -109,6 +109,14 @@
           ["b", true, true, 2, 1]
         ]
       },
+      {
+        "description": "multi 'with' table with multiple semi-joins",
+        "sql": "WITH t1 AS ( SELECT * FROM {tbl1} WHERE intCol > 1 ), t2 AS ( SELECT * FROM {tbl2} WHERE strCol1 IN (SELECT strCol FROM t1) AND intCol > 2 ) SELECT strCol1, strCol2, intCol FROM t2 WHERE strCol1 IN (SELECT strCol FROM t1)",
+        "outputs": [
+          ["b", "alice", 42],
+          ["b", "bob", 196883]
+        ]
+      },
       {
         "description": "nested 'with' on agg table: (with a as ( ... ), select ... ",
         "sql": "WITH agg1 AS (SELECT strCol1, strCol2, sum(intCol) AS sumVal FROM {tbl2} GROUP BY strCol1, strCol2) SELECT strCol1, avg(sumVal) AS avgVal FROM agg1 GROUP BY strCol1",