SQL: Fix issue with IIF function when condition folds

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

@@ -52,6 +52,21 @@ ORDER BY emp_no LIMIT 10;
 10014   | null
 ;
 
+caseWithConditionsFolded
+schema::emp_no:i|CASE_1:byte|CASE_2:i
+SELECT emp_no, CASE WHEN NULL = 1 THEN emp_no WHEN 10 < 5 THEN emp_no ELSE languages END AS "CASE_1",
+CASE WHEN NULL = 1 THEN languages WHEN 10 > 5 THEN emp_no ELSE languages END AS "CASE_2"
+FROM test_emp ORDER BY 1 LIMIT 5;
+
+ emp_no | CASE_1 | CASE_2
+--------+--------+-------
+10001   | 2      | 10001
+10002   | 5      | 10002
+10003   | 4      | 10003
+10004   | 5      | 10004
+10005   | 1      | 10005
+;
+
 caseWhere
 SELECT last_name FROM test_emp WHERE CASE WHEN LENGTH(last_name) < 7 THEN 'ShortName' ELSE 'LongName' END = 'LongName'
 ORDER BY emp_no LIMIT 10;
@@ -235,6 +250,19 @@ ORDER BY emp_no LIMIT 10;
 10014   | null
 ;
 
+iifWithConditionFolded
+schema::emp_no:i|IIF_1:i|IIF_2:byte|IIF_3:i
+SELECT emp_no, IIF(NULL, emp_no) AS IIF_1, IIF(NULL, emp_no, languages) AS IIF_2, IIF(10 > 5, emp_no, languages) AS IIF_3 FROM test_emp ORDER BY 1 LIMIT 5;
+
+ emp_no | IIF_1 | IIF_2 | IIF_3
+--------+-------+-------+------
+10001   | null  | 2     | 10001
+10002   | null  | 5     | 10002
+10003   | null  | 4     | 10003
+10004   | null  | 5     | 10004
+10005   | null  | 1     | 10005
+;
+
 iifWhere
 SELECT last_name FROM test_emp WHERE IIF(LENGTH(last_name) < 7, 'ShortName') IS NOT NULL ORDER BY emp_no LIMIT 10;
 

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

@@ -18,20 +18,21 @@
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.sql.expression.TypeResolutions.isBoolean;
+import static org.elasticsearch.xpack.sql.util.CollectionUtils.combine;
 
 public class Iif extends Case {
 
     public Iif(Source source, Expression condition, Expression thenResult, Expression elseResult) {
         super(source, Arrays.asList(new IfConditional(source, condition, thenResult), elseResult != null ? elseResult : Literal.NULL));
     }
 
-    private Iif(Source source, List<Expression> expressions) {
+    Iif(Source source, List<Expression> expressions) {
         super(source, expressions);
     }
 
     @Override
     protected NodeInfo<? extends Iif> info() {
-        return NodeInfo.create(this, Iif::new, conditions().get(0).condition(), conditions().get(0).result(), elseResult());
+        return NodeInfo.create(this, Iif::new, combine(conditions(), elseResult()));
     }
 
     @Override
@@ -41,6 +42,10 @@ public Expression replaceChildren(List<Expression> newChildren) {
 
     @Override
     protected TypeResolution resolveType() {
+        if (conditions().isEmpty()) {
+            return TypeResolution.TYPE_RESOLVED;
+        }
+
         TypeResolution conditionTypeResolution = isBoolean(conditions().get(0).condition(), sourceText(), Expressions.ParamOrdinal.FIRST);
         if (conditionTypeResolution.unresolved()) {
             return conditionTypeResolution;

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/predicate/conditional/CaseTests.java

@@ -17,9 +17,11 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 
+import static org.elasticsearch.xpack.sql.expression.Expression.TypeResolution;
 import static org.elasticsearch.xpack.sql.expression.function.scalar.FunctionTestUtils.randomIntLiteral;
 import static org.elasticsearch.xpack.sql.expression.function.scalar.FunctionTestUtils.randomStringLiteral;
 import static org.elasticsearch.xpack.sql.tree.Source.EMPTY;
@@ -117,6 +119,13 @@ public void testDataTypes() {
         assertEquals(DataType.KEYWORD, c.dataType());
     }
 
+    public void testAllConditionsFolded() {
+        Case c = new Case(EMPTY, Collections.singletonList(Literal.of(EMPTY, "foo")));
+        assertEquals(DataType.KEYWORD, c.dataType());
+        assertEquals(TypeResolution.TYPE_RESOLVED, c.typeResolved());
+        assertNotNull(c.info());
+    }
+
     private List<Expression> mutateChildren(Case c) {
         boolean removeConditional = randomBoolean();
         List<Expression> expressions = new ArrayList<>(c.children().size());

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/predicate/conditional/IifTests.java

@@ -6,20 +6,24 @@
 package org.elasticsearch.xpack.sql.expression.predicate.conditional;
 
 import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.expression.Literal;
 import org.elasticsearch.xpack.sql.expression.function.scalar.FunctionTestUtils;
 import org.elasticsearch.xpack.sql.expression.predicate.operator.comparison.Equals;
 import org.elasticsearch.xpack.sql.tree.AbstractNodeTestCase;
 import org.elasticsearch.xpack.sql.tree.NodeSubclassTests;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.SourceTests;
+import org.elasticsearch.xpack.sql.type.DataType;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 
 import static org.elasticsearch.xpack.sql.expression.function.scalar.FunctionTestUtils.randomIntLiteral;
 import static org.elasticsearch.xpack.sql.expression.function.scalar.FunctionTestUtils.randomStringLiteral;
+import static org.elasticsearch.xpack.sql.tree.Source.EMPTY;
 import static org.elasticsearch.xpack.sql.tree.SourceTests.randomSource;
 
 /**
@@ -74,6 +78,13 @@ public void testReplaceChildren() {
                 newChildren.get(2))));
     }
 
+    public void testConditionFolded() {
+        Iif iif = new Iif(EMPTY, Collections.singletonList(Literal.of(EMPTY, "foo")));
+        assertEquals(DataType.KEYWORD, iif.dataType());
+        assertEquals(Expression.TypeResolution.TYPE_RESOLVED, iif.typeResolved());
+        assertNotNull(iif.info());
+    }
+
     private List<Expression> mutateChildren(Iif iif) {
         List<Expression> expressions = new ArrayList<>(3);
         Equals eq = (Equals) iif.conditions().get(0).condition();

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

@@ -49,9 +49,9 @@
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.Coalesce;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.ConditionalFunction;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.Greatest;
-import org.elasticsearch.xpack.sql.expression.predicate.conditional.Iif;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.IfConditional;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.IfNull;
+import org.elasticsearch.xpack.sql.expression.predicate.conditional.Iif;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.Least;
 import org.elasticsearch.xpack.sql.expression.predicate.conditional.NullIf;
 import org.elasticsearch.xpack.sql.expression.predicate.logical.And;
@@ -112,6 +112,7 @@
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonList;
+import static org.elasticsearch.xpack.sql.expression.Expression.TypeResolution;
 import static org.elasticsearch.xpack.sql.expression.Literal.FALSE;
 import static org.elasticsearch.xpack.sql.expression.Literal.NULL;
 import static org.elasticsearch.xpack.sql.expression.Literal.TRUE;
@@ -630,6 +631,7 @@ public void testSimplifyCaseConditionsFoldWhenFalse() {
         assertThat(c.conditions().get(0).condition().toString(), startsWith("Equals[a{f}#"));
         assertThat(c.conditions().get(1).condition().toString(), startsWith("GreaterThan[a{f}#"));
         assertFalse(c.foldable());
+        assertEquals(TypeResolution.TYPE_RESOLVED, c.typeResolved());
     }
 
     public void testSimplifyCaseConditionsFoldWhenTrue() {
@@ -661,6 +663,7 @@ public void testSimplifyCaseConditionsFoldWhenTrue() {
         assertThat(c.conditions().get(0).condition().toString(), startsWith("Equals[a{f}#"));
         assertThat(c.conditions().get(1).condition().toString(), startsWith("Equals[=1,=1]#"));
         assertFalse(c.foldable());
+        assertEquals(TypeResolution.TYPE_RESOLVED, c.typeResolved());
     }
 
     public void testSimplifyCaseConditionsFoldCompletely() {
@@ -685,6 +688,7 @@ public void testSimplifyCaseConditionsFoldCompletely() {
         assertThat(c.conditions().get(0).condition().toString(), startsWith("Equals[=1,=1]#"));
         assertTrue(c.foldable());
         assertEquals("foo2", c.fold());
+        assertEquals(TypeResolution.TYPE_RESOLVED, c.typeResolved());
     }
 
     public void testSimplifyIif_ConditionTrue() {
@@ -696,6 +700,7 @@ public void testSimplifyIif_ConditionTrue() {
         assertEquals(1, iif.conditions().size());
         assertTrue(iif.foldable());
         assertEquals("foo", iif.fold());
+        assertEquals(TypeResolution.TYPE_RESOLVED, iif.typeResolved());
     }
 
     public void testSimplifyIif_ConditionFalse() {
@@ -707,6 +712,7 @@ public void testSimplifyIif_ConditionFalse() {
         assertEquals(0, iif.conditions().size());
         assertTrue(iif.foldable());
         assertEquals("bar", iif.fold());
+        assertEquals(TypeResolution.TYPE_RESOLVED, iif.typeResolved());
     }
 
     //