SQL: Fix bug with optimization of null related conditionals (#41355)

The SimplifyConditional rule is removing NULL literals from those
functions to simplify their evaluation. This happens in the Optimizer
and a new instance of the conditional function is generated. Previously,
the dataType was not set properly (defaulted to DataType.NULL) for
those new instances and since the resolveType() wasn't called again
it resulted in returning always null.

E.g.:

SELECT COALESCE(null, 'foo', null, 'bar')

COALESCE(null, 'foo', null, 'bar')
-----------------
null

This issue was not visible before because the tests always used an alias
for the conditional function which caused the resolveType() to be
called which sets the dataType properly.

E.g.:

SELECT COALESCE(null, 'foo', null, 'bar') as c

c
-----------------
foo

(cherry picked from commit c39980a)

Author: matriv

x-pack/plugin/sql/qa/src/main/resources/null.csv-spec

@@ -61,6 +61,13 @@ c:i
 ;
 
 coalesceMixed
+SELECT COALESCE(null, 123, null, 321);
+
+COALESCE(null, 123, null, 321):i
+123
+;
+
+coalesceMixedWithAlias
 SELECT COALESCE(null, 123, null, 321) AS c;
 
 c:i

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/predicate/conditional/ArbitraryConditionalFunction.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.ConditionalProcessor.ConditionalOperation;
 import org.elasticsearch.xpack.sql.tree.Source;
+import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.type.DataTypeConversion;
 
 import java.util.ArrayList;
@@ -35,6 +36,7 @@ public abstract class ArbitraryConditionalFunction extends ConditionalFunction {
 
     @Override
     protected TypeResolution resolveType() {
+        dataType = DataType.NULL;
         for (Expression e : children()) {
             dataType = DataTypeConversion.commonType(dataType, e.dataType());
         }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/predicate/conditional/ConditionalFunction.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.type.DataType;
+import org.elasticsearch.xpack.sql.type.DataTypeConversion;
 
 import java.util.List;
 
@@ -20,14 +21,20 @@
  */
 public abstract class ConditionalFunction extends ScalarFunction {
 
-    protected DataType dataType = DataType.NULL;
+    protected DataType dataType = null;
 
     ConditionalFunction(Source source, List<Expression> fields) {
         super(source, fields);
     }
 
     @Override
     public DataType dataType() {
+        if (dataType == null) {
+            dataType = DataType.NULL;
+            for (Expression exp : children()) {
+                dataType = DataTypeConversion.commonType(dataType, exp.dataType());
+            }
+        }
         return dataType;
     }
 

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java

@@ -481,6 +481,7 @@ public void testSimplifyCoalesceRandomNullsWithValue() {
                         randomListOfNulls())));
         assertEquals(1, e.children().size());
         assertEquals(Literal.TRUE, e.children().get(0));
+        assertEquals(DataType.BOOLEAN, e.dataType());
     }
 
     private List<Expression> randomListOfNulls() {
@@ -494,6 +495,7 @@ public void testSimplifyCoalesceFirstLiteral() {
         assertEquals(Coalesce.class, e.getClass());
         assertEquals(1, e.children().size());
         assertEquals(Literal.TRUE, e.children().get(0));
+        assertEquals(DataType.BOOLEAN, e.dataType());
     }
 
     public void testSimplifyIfNullNulls() {
@@ -507,11 +509,13 @@ public void testSimplifyIfNullWithNullAndValue() {
         assertEquals(IfNull.class, e.getClass());
         assertEquals(1, e.children().size());
         assertEquals(ONE, e.children().get(0));
+        assertEquals(DataType.INTEGER, e.dataType());
 
         e = new SimplifyConditional().rule(new IfNull(EMPTY, ONE, Literal.NULL));
         assertEquals(IfNull.class, e.getClass());
         assertEquals(1, e.children().size());
         assertEquals(ONE, e.children().get(0));
+        assertEquals(DataType.INTEGER, e.dataType());
     }
 
     public void testFoldNullNotAppliedOnNullIf() {
@@ -539,6 +543,7 @@ public void testSimplifyGreatestRandomNullsWithValue() {
         assertEquals(2, e.children().size());
         assertEquals(ONE, e.children().get(0));
         assertEquals(TWO, e.children().get(1));
+        assertEquals(DataType.INTEGER, e.dataType());
     }
 
     public void testSimplifyLeastNulls() {
@@ -560,6 +565,7 @@ public void testSimplifyLeastRandomNullsWithValue() {
         assertEquals(2, e.children().size());
         assertEquals(ONE, e.children().get(0));
         assertEquals(TWO, e.children().get(1));
+        assertEquals(DataType.INTEGER, e.dataType());
     }
 
     public void testConcatFoldingIsNotNull() {