ES|QL: Skip CASE function from InferIsNotNull rule checks (#113123) (#113249)

* Skip CASE function from InferIsNotNull rule checks

(cherry picked from commit 2cccf13)

Author: astefan

docs/changelog/113123.yaml

@@ -0,0 +1,6 @@
+pr: 113123
+summary: "ES|QL: Skip CASE function from `InferIsNotNull` rule checks"
+area: ES|QL
+type: bug
+issues:
+ - 112704

x-pack/plugin/esql/qa/testFixtures/src/main/resources/conditional.csv-spec

@@ -261,3 +261,23 @@ warning:Line 1:38: java.lang.IllegalArgumentException: single-value function enc
 
 a:integer | b:integer | same:boolean
 ;
+
+caseOnTheValue_NotOnTheField
+required_capability: fixed_wrong_is_not_null_check_on_case
+
+FROM employees 
+| WHERE emp_no < 10022 AND emp_no > 10012 
+| KEEP languages, emp_no 
+| EVAL eval = CASE(languages == 1, null, languages == 2, "bilingual", languages > 2, "multilingual", languages IS NULL, "languages is null") 
+| SORT languages, emp_no 
+| WHERE eval IS NOT NULL;
+
+languages:integer| emp_no:integer|eval:keyword                    
+2                |10016          |bilingual        
+2                |10017          |bilingual        
+2                |10018          |bilingual        
+5                |10014          |multilingual     
+5                |10015          |multilingual     
+null             |10020          |languages is null
+null             |10021          |languages is null
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -309,7 +309,13 @@ public enum Cap {
         /**
          * Supported the text categorization function "CATEGORIZE".
          */
-        CATEGORIZE(true);
+        CATEGORIZE(true),
+
+        /**
+         * Don't optimize CASE IS NOT NULL function by not requiring the fields to be not null as well.
+         * https://github.com/elastic/elasticsearch/issues/112704
+         */
+        FIXED_WRONG_IS_NOT_NULL_CHECK_ON_CASE;
 
         private final boolean snapshotOnly;
         private final FeatureFlag featureFlag;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/InferIsNotNull.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.xpack.esql.core.expression.predicate.nulls.IsNotNull;
 import org.elasticsearch.xpack.esql.core.rule.Rule;
 import org.elasticsearch.xpack.esql.core.util.CollectionUtils;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
@@ -24,8 +25,7 @@
 import static java.util.Collections.emptySet;
 
 /**
- * Simplify IsNotNull targets by resolving the underlying expression to its root fields with unknown
- * nullability.
+ * Simplify IsNotNull targets by resolving the underlying expression to its root fields.
  * e.g.
  * (x + 1) / 2 IS NOT NULL --> x IS NOT NULL AND (x+1) / 2 IS NOT NULL
  * SUBSTRING(x, 3) > 4 IS NOT NULL --> x IS NOT NULL AND SUBSTRING(x, 3) > 4 IS NOT NULL
@@ -85,7 +85,7 @@ protected Set<Expression> resolveExpressionAsRootAttributes(Expression exp, Attr
 
     private boolean doResolve(Expression exp, AttributeMap<Expression> aliases, Set<Expression> resolvedExpressions) {
         boolean changed = false;
-        // check if the expression can be skipped or is not nullabe
+        // check if the expression can be skipped
         if (skipExpression(exp)) {
             resolvedExpressions.add(exp);
         } else {
@@ -106,6 +106,13 @@ private boolean doResolve(Expression exp, AttributeMap<Expression> aliases, Set<
     }
 
     private static boolean skipExpression(Expression e) {
-        return e instanceof Coalesce;
+        // These two functions can have a complex set of expressions as arguments that can mess up the simplification we are trying to add.
+        // If there is a "case(f is null, null, ...) is not null" expression,
+        // assuming that "case(f is null.....) is not null AND f is not null" (what this rule is doing) is a wrong assumption because
+        // the "case" function will want both null "f" and not null "f". Doing it like this contradicts the condition inside case, so we
+        // must avoid these cases.
+        // We could be smarter and look inside "case" and "coalesce" to see if there is any comparison of fields with "null" but,
+        // the complexity is too high to warrant an attempt _now_.
+        return e instanceof Coalesce || e instanceof Case;
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizerTests.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.StartsWith;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
@@ -58,8 +59,11 @@
 
 import static java.util.Collections.emptyMap;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.L;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.ONE;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_SEARCH_STATS;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_VERIFIER;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.THREE;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.TWO;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.getFieldAttribute;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOf;
@@ -81,10 +85,6 @@ public class LocalLogicalPlanOptimizerTests extends ESTestCase {
     private static LogicalPlanOptimizer logicalOptimizer;
     private static Map<String, EsField> mapping;
 
-    private static final Literal ONE = L(1);
-    private static final Literal TWO = L(2);
-    private static final Literal THREE = L(3);
-
     @BeforeClass
     public static void init() {
         parser = new EsqlParser();
@@ -386,38 +386,6 @@ public void testMissingFieldInFilterNoProjection() {
         );
     }
 
-    public void testIsNotNullOnCoalesce() {
-        var plan = localPlan("""
-              from test
-            | where coalesce(emp_no, salary) is not null
-            """);
-
-        var limit = as(plan, Limit.class);
-        var filter = as(limit.child(), Filter.class);
-        var inn = as(filter.condition(), IsNotNull.class);
-        var coalesce = as(inn.children().get(0), Coalesce.class);
-        assertThat(Expressions.names(coalesce.children()), contains("emp_no", "salary"));
-        var source = as(filter.child(), EsRelation.class);
-    }
-
-    public void testIsNotNullOnExpression() {
-        var plan = localPlan("""
-              from test
-            | eval x = emp_no + 1
-            | where x is not null
-            """);
-
-        var limit = as(plan, Limit.class);
-        var filter = as(limit.child(), Filter.class);
-        var inn = as(filter.condition(), IsNotNull.class);
-        assertThat(Expressions.names(inn.children()), contains("x"));
-        var eval = as(filter.child(), Eval.class);
-        filter = as(eval.child(), Filter.class);
-        inn = as(filter.condition(), IsNotNull.class);
-        assertThat(Expressions.names(inn.children()), contains("emp_no"));
-        var source = as(filter.child(), EsRelation.class);
-    }
-
     public void testSparseDocument() throws Exception {
         var query = """
             from large
@@ -516,6 +484,66 @@ public void testIsNotNullOnFunctionWithTwoFields() {
         assertEquals(expected, new InferIsNotNull().apply(f));
     }
 
+    public void testIsNotNullOnCoalesce() {
+        var plan = localPlan("""
+              from test
+            | where coalesce(emp_no, salary) is not null
+            """);
+
+        var limit = as(plan, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var inn = as(filter.condition(), IsNotNull.class);
+        var coalesce = as(inn.children().get(0), Coalesce.class);
+        assertThat(Expressions.names(coalesce.children()), contains("emp_no", "salary"));
+        var source = as(filter.child(), EsRelation.class);
+    }
+
+    public void testIsNotNullOnExpression() {
+        var plan = localPlan("""
+              from test
+            | eval x = emp_no + 1
+            | where x is not null
+            """);
+
+        var limit = as(plan, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var inn = as(filter.condition(), IsNotNull.class);
+        assertThat(Expressions.names(inn.children()), contains("x"));
+        var eval = as(filter.child(), Eval.class);
+        filter = as(eval.child(), Filter.class);
+        inn = as(filter.condition(), IsNotNull.class);
+        assertThat(Expressions.names(inn.children()), contains("emp_no"));
+        var source = as(filter.child(), EsRelation.class);
+    }
+
+    public void testIsNotNullOnCase() {
+        var plan = localPlan("""
+              from test
+            | where case(emp_no > 10000, "1", salary < 50000, "2", first_name) is not null
+            """);
+
+        var limit = as(plan, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var inn = as(filter.condition(), IsNotNull.class);
+        var caseF = as(inn.children().get(0), Case.class);
+        assertThat(Expressions.names(caseF.children()), contains("emp_no > 10000", "\"1\"", "salary < 50000", "\"2\"", "first_name"));
+        var source = as(filter.child(), EsRelation.class);
+    }
+
+    public void testIsNotNullOnCase_With_IS_NULL() {
+        var plan = localPlan("""
+              from test
+            | where case(emp_no IS NULL, "1", salary IS NOT NULL, "2", first_name) is not null
+            """);
+
+        var limit = as(plan, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var inn = as(filter.condition(), IsNotNull.class);
+        var caseF = as(inn.children().get(0), Case.class);
+        assertThat(Expressions.names(caseF.children()), contains("emp_no IS NULL", "\"1\"", "salary IS NOT NULL", "\"2\"", "first_name"));
+        var source = as(filter.child(), EsRelation.class);
+    }
+
     private IsNotNull isNotNull(Expression field) {
         return new IsNotNull(EMPTY, field);
     }