Fix cte filter pushdown wrong results by splitting SpecialFormExpress…

…ions
tdcmeehan · May 9, 2024 · 8e34250 · 8e34250
1 parent 6c2f50d
commit 8e34250
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 5 deletions.
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java
@@ -113,6 +113,72 @@ public void testPersistentCteWithTimeStampWithTimeZoneType()
                         testQuery));
     }
 
+    @Test
+    public void testComplexCommonFilterPushdown()
+    {
+        QueryRunner queryRunner = getQueryRunner();
+        String testQuery = "WITH order_platform_data AS (\n" +
+                "  SELECT\n" +
+                "    o.orderkey AS order_key,\n" +
+                "    o.orderdate AS datestr,\n" +
+                "    o.orderpriority AS event_type\n" +
+                "  FROM\n" +
+                "    orders o\n" +
+                "  WHERE\n" +
+                "    o.orderdate BETWEEN DATE '1995-01-01' AND DATE '1995-01-31'\n" +
+                "    AND o.orderpriority IN ('1-URGENT', '3-MEDIUM')\n" +
+                "  UNION ALL\n" +
+                "  SELECT\n" +
+                "    l.orderkey AS order_key,\n" +
+                "    o.orderdate AS datestr,\n" +
+                "    o.orderpriority AS event_type\n" +
+                "  FROM\n" +
+                "    lineitem l\n" +
+                "    JOIN orders o ON l.orderkey = o.orderkey\n" +
+                "  WHERE\n" +
+                "    o.orderdate BETWEEN DATE '1995-01-01' AND DATE '1995-01-31'\n" +
+                "    AND o.orderpriority IN ('2-HIGH', '5-LOW')\n" +
+                "),\n" +
+                "urgent AS (\n" +
+                "    SELECT order_key, datestr\n" +
+                "    FROM order_platform_data\n" +
+                "    WHERE event_type = '1-URGENT'\n" +
+                "),\n" +
+                "medium AS (\n" +
+                "    SELECT order_key, datestr\n" +
+                "    FROM order_platform_data\n" +
+                "    WHERE event_type = '3-MEDIUM'\n" +
+                "),\n" +
+                "high AS (\n" +
+                "    SELECT order_key, datestr\n" +
+                "    FROM order_platform_data\n" +
+                "    WHERE event_type = '2-HIGH'\n" +
+                "),\n" +
+                "low AS (\n" +
+                "    SELECT order_key, datestr\n" +
+                "    FROM order_platform_data\n" +
+                "    WHERE event_type = '5-LOW'\n" +
+                ")\n" +
+                "SELECT\n" +
+                "    ofin.order_key AS order_key,\n" +
+                "    ofin.datestr AS order_date\n" +
+                " FROM " +
+                "    urgent ofin\n" +
+                "    LEFT JOIN medium oproc ON ofin.datestr = oproc.datestr\n" +
+                "   LEFT JOIN low on oproc.datestr = low.datestr" +
+                "  LEFT JOIN high on low.datestr = high.datestr" +
+                " ORDER BY\n" +
+                "    ofin.order_key\n";
+        compareResults(queryRunner.execute(Session.builder(super.getSession())
+                                .setSystemProperty(PUSHDOWN_SUBFIELDS_ENABLED, "true")
+                                .setSystemProperty(CTE_MATERIALIZATION_STRATEGY, "HEURISTIC_COMPLEX_QUERIES_ONLY")
+                                .setSystemProperty(CTE_FILTER_AND_PROJECTION_PUSHDOWN_ENABLED, "true")
+                                .build(),
+                        testQuery),
+                queryRunner.execute(getSession(),
+                        testQuery));
+    }
+
     @Test
     public void testPersistentCteWithStructTypes()
     {

diff --git a/...java/com/facebook/presto/sql/planner/optimizations/CteProjectionAndPredicatePushDown.java b/...java/com/facebook/presto/sql/planner/optimizations/CteProjectionAndPredicatePushDown.java
@@ -31,6 +31,7 @@
 import com.facebook.presto.sql.planner.SimplePlanVisitor;
 import com.facebook.presto.sql.planner.TypeProvider;
 import com.facebook.presto.sql.planner.VariablesExtractor;
+import com.facebook.presto.sql.planner.iterative.rule.SimplifyRowExpressions;
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.SimplePlanRewriter;
 import com.google.common.collect.ImmutableList;
@@ -48,7 +49,6 @@
 import static com.facebook.presto.SystemSessionProperties.getCteFilterAndProjectionPushdownEnabled;
 import static com.facebook.presto.SystemSessionProperties.isCteMaterializationApplicable;
 import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
-import static com.facebook.presto.spi.relation.SpecialFormExpression.Form.OR;
 import static com.facebook.presto.sql.planner.PlannerUtils.isConstant;
 import static com.facebook.presto.sql.planner.plan.ChildReplacer.replaceChildren;
 import static com.facebook.presto.sql.relational.Expressions.constant;
@@ -370,14 +370,21 @@ private PlanNode addFilter(PlanNode node, List<RowExpression> predicates)
                 return node;
             }
 
-            RowExpression predicate;
+            RowExpression resultPredicate;
             if (predicates.size() == 1) {
-                predicate = predicates.get(0);
+                resultPredicate = predicates.get(0);
             }
             else {
-                predicate = new SpecialFormExpression(OR, BOOLEAN, predicates);
+                resultPredicate = predicates.get(0);
+                for (int i = 1; i < predicates.size(); i++) {
+                    resultPredicate = new SpecialFormExpression(
+                            SpecialFormExpression.Form.OR,
+                            BOOLEAN,
+                            resultPredicate, predicates.get(i));
+                }
             }
-            return new FilterNode(node.getSourceLocation(), idAllocator.getNextId(), node, predicate);
+            resultPredicate = SimplifyRowExpressions.rewrite(resultPredicate, metadata, session.toConnectorSession());
+            return new FilterNode(node.getSourceLocation(), idAllocator.getNextId(), node, resultPredicate);
         }
 
         private boolean isConstTrue(List<RowExpression> predicates)