Extract UnnestNode's source predicate

In EffectivePredicateExtractor, extract and propagate UnnestNode's source predicate for INNER and LEFT join type.
G-Research · Nov 14, 2024 · 3d1aac6 · 3d1aac6
1 parent cd79163
commit 3d1aac6
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/EffectivePredicateExtractor.java b/core/trino-main/src/main/java/io/trino/sql/planner/EffectivePredicateExtractor.java
@@ -301,7 +301,10 @@ public Expression visitUnion(UnionNode node, Void context)
         @Override
         public Expression visitUnnest(UnnestNode node, Void context)
         {
-            return TRUE;
+            return switch (node.getJoinType()) {
+                case INNER, LEFT -> pullExpressionThroughSymbols(node.getSource().accept(this, context), node.getOutputSymbols());
+                case RIGHT, FULL -> TRUE;
+            };
         }
 
         @Override

diff --git a/core/trino-main/src/test/java/io/trino/sql/planner/TestLogicalPlanner.java b/core/trino-main/src/test/java/io/trino/sql/planner/TestLogicalPlanner.java
@@ -166,6 +166,7 @@
 import static io.trino.sql.planner.assertions.PlanMatchPattern.tableScan;
 import static io.trino.sql.planner.assertions.PlanMatchPattern.topN;
 import static io.trino.sql.planner.assertions.PlanMatchPattern.topNRanking;
+import static io.trino.sql.planner.assertions.PlanMatchPattern.unnest;
 import static io.trino.sql.planner.assertions.PlanMatchPattern.values;
 import static io.trino.sql.planner.assertions.PlanMatchPattern.windowFunction;
 import static io.trino.sql.planner.optimizations.PlanNodeSearcher.searchFrom;
@@ -2631,6 +2632,30 @@ public void testPruneWindow()
                         values("c")));
     }
 
+    @Test
+    public void testCrossJoinUnnest()
+    {
+        assertPlan(
+                """
+                WITH RECURSIVE recursive_call (level) AS (
+                    SELECT 1 AS level
+                    FROM (SELECT ARRAY[] AS array) AS t1 CROSS JOIN UNNEST(array) AS t2 (x)
+                    UNION ALL
+                    SELECT recursive_call.level + 1 AS level
+                    FROM recursive_call INNER JOIN nation AS c ON recursive_call.level <= 1
+                )
+                SELECT * FROM recursive_call
+                """,
+                output(exchange(
+                        LOCAL,
+                        any(unnest(values("array"))),
+                        any(join(
+                                INNER,
+                                builder -> builder
+                                        .left(any(unnest(values("array"))))
+                                        .right(exchange(tableScan("nation"))))))));
+    }
+
     private Session noJoinReordering()
     {
         return Session.builder(getPlanTester().getDefaultSession())