From 448fb70f687bf584ab9a51868a1b894c1dd1cd74 Mon Sep 17 00:00:00 2001
From: starocean999 <40539150+starocean999@users.noreply.github.com>
Date: Thu, 8 Feb 2024 13:48:13 +0800
Subject: [PATCH] [fix](nereids) collect all correlated slots from subquery in
 correct way (#30908)

---
 .../org/apache/doris/nereids/util/Utils.java  | 21 ++-----------------
 .../subquery/subquery_unnesting.out           | 12 +++++++++++
 .../subquery/subquery_unnesting.groovy        |  1 +
 3 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/Utils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/Utils.java
index 31b29eb325..bc53075151 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/Utils.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/Utils.java
@@ -17,10 +17,8 @@
 
 package org.apache.doris.nereids.util;
 
-import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.expressions.Cast;
 import org.apache.doris.nereids.trees.expressions.Expression;
-import org.apache.doris.nereids.trees.expressions.Not;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
 
@@ -167,23 +165,8 @@ public class Utils {
      */
     public static List<Expression> getCorrelatedSlots(List<Expression> correlatedPredicates,
             List<Expression> correlatedSlots) {
-        List<Expression> slots = new ArrayList<>();
-        correlatedPredicates.forEach(predicate -> {
-            if (!(predicate instanceof BinaryExpression)
-                    && (!(predicate instanceof Not) || !(predicate.child(0) instanceof BinaryExpression))) {
-                throw new AnalysisException("Unsupported correlated subquery with"
-                        + " non-equals correlated predicate " + predicate.toSql());
-            }
-
-            BinaryExpression binaryExpression;
-            if (predicate instanceof Not) {
-                binaryExpression = (BinaryExpression) ((Not) predicate).child();
-            } else {
-                binaryExpression = (BinaryExpression) predicate;
-            }
-            slots.addAll(collectCorrelatedSlotsFromChildren(binaryExpression, correlatedSlots));
-        });
-        return slots;
+        return ExpressionUtils.getInputSlotSet(correlatedPredicates).stream()
+                .filter(slot -> !correlatedSlots.contains(slot)).collect(Collectors.toList());
     }
 
     private static List<Expression> collectCorrelatedSlotsFromChildren(
diff --git a/regression-test/data/nereids_p0/subquery/subquery_unnesting.out b/regression-test/data/nereids_p0/subquery/subquery_unnesting.out
index f809cbf8bc..df30f66818 100644
--- a/regression-test/data/nereids_p0/subquery/subquery_unnesting.out
+++ b/regression-test/data/nereids_p0/subquery/subquery_unnesting.out
@@ -1492,3 +1492,15 @@
 22	3
 24	4
 
+-- !select60 --
+\N	1
+\N	2
+1	2
+1	3
+2	4
+3	3
+3	4
+20	2
+22	3
+24	4
+
diff --git a/regression-test/suites/nereids_p0/subquery/subquery_unnesting.groovy b/regression-test/suites/nereids_p0/subquery/subquery_unnesting.groovy
index 0a830f1a36..2bae4cc7bb 100644
--- a/regression-test/suites/nereids_p0/subquery/subquery_unnesting.groovy
+++ b/regression-test/suites/nereids_p0/subquery/subquery_unnesting.groovy
@@ -129,4 +129,5 @@ suite ("subquery_unnesting") {
     qt_select57 """select t1.* from t1 left join t2 on t1.k2 = t2.k3 and not exists ( select t3.k1 from t3 where t1.k2 < t3.k2 ) order by t1.k1, t1.k2;"""
     qt_select58 """select t1.* from t1 left join t2 on t1.k2 = t2.k3 and not exists ( select t3.k1 from t3 ) or t1.k1 < 10 order by t1.k1, t1.k2;"""
     qt_select59 """select t1.* from t1 left join t2 on t1.k2 = t2.k3 and not exists ( select t3.k1 from t3 ) order by t1.k1, t1.k2;"""
+    qt_select60 """select * from t1 where exists(select distinct k1 from t2 where t1.k1 > t2.k3 or t1.k2 < t2.v1) order by t1.k1, t1.k2;"""
 }