Flatten coalesce expressions

Flatten coalesce(colA, coalesce(colB, colC)) to coalesce(colA, colB, colC).
trinodb · Jan 25, 2019 · cf17d08 · cf17d08
1 parent ac02f3d
commit cf17d08
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/sql/planner/ExpressionInterpreter.java b/presto-main/src/main/java/io/prestosql/sql/planner/ExpressionInterpreter.java
@@ -549,6 +549,12 @@ protected Object visitCoalesceExpression(CoalesceExpression node, Object context
             List<Object> values = node.getOperands().stream()
                     .map(value -> processWithExceptionHandling(value, context))
                     .filter(Objects::nonNull)
+                    .flatMap(expression -> {
+                        if (expression instanceof CoalesceExpression) {
+                            return ((CoalesceExpression) expression).getOperands().stream();
+                        }
+                        return Stream.of(expression);
+                    })
                     .collect(Collectors.toList());
 
             if ((!values.isEmpty() && !(values.get(0) instanceof Expression)) || values.size() == 1) {

diff --git a/presto-main/src/test/java/io/prestosql/sql/TestExpressionInterpreter.java b/presto-main/src/test/java/io/prestosql/sql/TestExpressionInterpreter.java
@@ -1145,6 +1145,11 @@ public void testCoalesce()
         assertOptimizedEquals("coalesce(6, unbound_long2, unbound_long, unbound_long3)", "6");
         assertOptimizedEquals("coalesce(2 * 3, unbound_long2, unbound_long, unbound_long3)", "6");
         assertOptimizedMatches("coalesce(random(), random(), 5)", "coalesce(random(), random(), 5E0)");
+        assertOptimizedMatches("coalesce(unbound_long, coalesce(unbound_long, 1))", "coalesce(unbound_long, BIGINT '1')");
+        assertOptimizedMatches("coalesce(coalesce(unbound_long, coalesce(unbound_long, 1)), unbound_long2)", "coalesce(unbound_long, BIGINT '1')");
+        assertOptimizedMatches("coalesce(unbound_long, 2, coalesce(unbound_long, 1))", "coalesce(unbound_long, BIGINT '2')");
+        assertOptimizedMatches("coalesce(coalesce(unbound_long, coalesce(unbound_long2, unbound_long3)), 1)", "coalesce(unbound_long, unbound_long2, unbound_long3, BIGINT '1')");
+        assertOptimizedMatches("coalesce(unbound_double, coalesce(random(), unbound_double))", "coalesce(unbound_double, random())");
     }
 
     @Test