DRILL-7041: CompileException happens if a nested coalesce function returns null

- Made `NullExpression`s in `IfExpression` with nested `IfExpression`s to be rewritten to typed ones recursively if necessary
closes #1668

2 files changed, 42 insertions, 22 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/expr/ExpressionTreeMaterializer.java b/exec/java-exec/src/main/java/org/apache/drill/exec/expr/ExpressionTreeMaterializer.java
index 4e35d6dea..4a4f39d7e 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/expr/ExpressionTreeMaterializer.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/expr/ExpressionTreeMaterializer.java
@@ -23,6 +23,7 @@ import java.util.Collections;
 import java.util.Deque;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Queue;
 
 import org.apache.drill.shaded.guava.com.google.common.base.Preconditions;
@@ -82,10 +83,7 @@ import org.apache.drill.exec.resolver.FunctionResolver;
 import org.apache.drill.exec.resolver.FunctionResolverFactory;
 import org.apache.drill.exec.resolver.TypeCastRules;
 
-import org.apache.drill.shaded.guava.com.google.common.base.Optional;
-import org.apache.drill.shaded.guava.com.google.common.base.Predicate;
 import org.apache.drill.shaded.guava.com.google.common.collect.ImmutableList;
-import org.apache.drill.shaded.guava.com.google.common.collect.Iterables;
 import org.apache.drill.shaded.guava.com.google.common.collect.Lists;
 import org.apache.drill.exec.store.parquet.stat.ColumnStatistics;
 import org.apache.drill.exec.util.DecimalUtility;
@@ -660,27 +658,16 @@ public class ExpressionTreeMaterializer {
       allExpressions.add(conditions.expression);
       allExpressions.add(newElseExpr);
 
-      boolean containsNullExpr = Iterables.any(allExpressions, new Predicate<LogicalExpression>() {
-        @Override
-        public boolean apply(LogicalExpression input) {
-          return input instanceof NullExpression;
-        }
-      });
-
-      if (containsNullExpr) {
-        Optional<LogicalExpression> nonNullExpr = Iterables.tryFind(allExpressions,
-          new Predicate<LogicalExpression>() {
-            @Override
-            public boolean apply(LogicalExpression input) {
-              return !input.getMajorType().getMinorType().equals(TypeProtos.MinorType.NULL);
-            }
-          }
-        );
+      boolean containsNullType = allExpressions.stream()
+          .anyMatch(expression -> expression.getMajorType().getMinorType() == MinorType.NULL);
+      if (containsNullType) {
+        Optional<LogicalExpression> nonNullExpr = allExpressions.stream()
+            .filter(expression -> expression.getMajorType().getMinorType() != MinorType.NULL)
+            .findAny();
 
-        if(nonNullExpr.isPresent()) {
+        if (nonNullExpr.isPresent()) {
           MajorType type = nonNullExpr.get().getMajorType();
           conditions = new IfExpression.IfCondition(conditions.condition, rewriteNullExpression(conditions.expression, type));
-
           newElseExpr = rewriteNullExpression(newElseExpr, type);
         }
       }
@@ -730,11 +717,24 @@ public class ExpressionTreeMaterializer {
     private LogicalExpression rewriteNullExpression(LogicalExpression expr, MajorType type) {
       if(expr instanceof NullExpression) {
         return new TypedNullConstant(type);
+      } else if (expr instanceof IfExpression) {
+        return rewriteIfWithNullExpression((IfExpression) expr, type);
       } else {
         return expr;
       }
     }
 
+    private LogicalExpression rewriteIfWithNullExpression(IfExpression expr, MajorType type) {
+      IfCondition condition = new IfExpression.IfCondition(expr.ifCondition.condition, rewriteNullExpression(expr.ifCondition.expression, type));
+      LogicalExpression elseExpression = rewriteNullExpression(expr.elseExpression, type);
+      return new IfExpression.Builder()
+          .setPosition(expr.getPosition())
+          .setOutputType(expr.outputType)
+          .setIfCondition(condition)
+          .setElse(elseExpression)
+          .build();
+    }
+
     @Override
     public LogicalExpression visitIntConstant(IntExpression intExpr, FunctionLookupContext functionLookupContext) {
       return intExpr;
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestCaseNullableTypes.java b/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestCaseNullableTypes.java
index 4523b45d0..cf598275e 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestCaseNullableTypes.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestCaseNullableTypes.java
@@ -26,7 +26,6 @@ import org.junit.experimental.categories.Category;
 import java.math.BigDecimal;
 
 /**
- * DRILL-4906
  * Tests for handling nullable types in CASE function
  */
 @Category(SqlTest.class)
@@ -158,4 +157,25 @@ public class TestCaseNullableTypes extends BaseTestQuery {
       resetSessionOption(PlannerSettings.ENABLE_DECIMAL_DATA_TYPE_KEY);
     }
   }
+
+  // Coalesce is being transformed to if-else cases
+  @Test
+  public void testCoalesceWithUntypedNullValues() throws Exception {
+    testBuilder()
+        .sqlQuery("select coalesce(coalesce(n_name1, n_name2, n_name), coalesce(n_name3, n_name4), n_name3) res from cp.`tpch/nation.parquet` limit 1")
+        .ordered()
+        .baselineColumns("res")
+        .baselineValues("ALGERIA")
+        .go();
+  }
+
+  @Test
+  public void testCoalesceWithUntypedNullValues_2() throws Exception {
+    testBuilder()
+        .sqlQuery("select coalesce(coalesce(n_name1, n_name2), n_name) res from cp.`tpch/nation.parquet` limit 1")
+        .ordered()
+        .baselineColumns("res")
+        .baselineValues("ALGERIA")
+        .go();
+  }
 }