SQL: Optimizer rule for folding nullable expressions (#35080)

Add optimization for folding nullable expressions with a NULL 
argument. This is a variant of folding for the NULL case.

Fix #34826

(cherry picked from commit 6abddd8)

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Expression.java

@@ -78,6 +78,7 @@ public Object fold() {
         throw new SqlIllegalArgumentException("Should not fold expression");
     }
 
+    // whether the expression becomes null if at least one param/input is null
     public abstract boolean nullable();
 
     // the references/inputs/leaves of the expression tree

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Literal.java

@@ -161,7 +161,11 @@ public static Literal of(String name, Expression foldable) {
         if (name == null) {
             name = foldable instanceof NamedExpression ? ((NamedExpression) foldable).name() : String.valueOf(fold);
         }
-
         return new Literal(foldable.location(), name, fold, foldable.dataType());
     }
-}
+
+    public static Literal of(Expression source, Object value) {
+        String name = source instanceof NamedExpression ? ((NamedExpression) source).name() : String.valueOf(value);
+        return new Literal(source.location(), name, value, source.dataType());
+    }
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/Function.java

@@ -46,7 +46,7 @@ public String name() {
 
     @Override
     public boolean nullable() {
-        return false;
+        return Expressions.nullable(children());
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/string/Concat.java

@@ -49,6 +49,11 @@ protected Pipe makePipe() {
         return new ConcatFunctionPipe(location(), this, Expressions.pipe(left()), Expressions.pipe(right()));
     }
 
+    @Override
+    public boolean nullable() {
+        return left().nullable() && right().nullable();
+    }
+
     @Override
     public boolean foldable() {
         return left().foldable() && right().foldable();
@@ -80,4 +85,4 @@ public ScriptTemplate scriptWithField(FieldAttribute field) {
     public DataType dataType() {
         return DataType.KEYWORD;
     }
-}
+}

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/predicate/logical/BinaryLogic.java

@@ -33,4 +33,9 @@ protected TypeResolution resolveInputType(Expression e, Expressions.ParamOrdinal
     protected Pipe makePipe() {
         return new BinaryLogicPipe(location(), this, Expressions.pipe(left()), Expressions.pipe(right()), function());
     }
+
+    @Override
+    public boolean nullable() {
+        return left().nullable() && right().nullable();
+    }
 }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

@@ -39,6 +39,8 @@
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryOperator;
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryOperator.Negateable;
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryPredicate;
+import org.elasticsearch.xpack.sql.expression.predicate.In;
+import org.elasticsearch.xpack.sql.expression.predicate.IsNotNull;
 import org.elasticsearch.xpack.sql.expression.predicate.Predicates;
 import org.elasticsearch.xpack.sql.expression.predicate.Range;
 import org.elasticsearch.xpack.sql.expression.predicate.logical.And;
@@ -63,6 +65,7 @@
 import org.elasticsearch.xpack.sql.rule.RuleExecutor;
 import org.elasticsearch.xpack.sql.session.EmptyExecutable;
 import org.elasticsearch.xpack.sql.session.SingletonExecutable;
+import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.CollectionUtils;
 
 import java.util.ArrayList;
@@ -122,6 +125,7 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 new CombineProjections(),
                 // folding
                 new ReplaceFoldableAttributes(),
+                new FoldNull(),
                 new ConstantFolding(),
                 // boolean
                 new BooleanSimplification(),
@@ -682,8 +686,7 @@ protected LogicalPlan rule(Filter filter) {
                 if (TRUE.equals(filter.condition())) {
                     return filter.child();
                 }
-                // TODO: add comparison with null as well
-                if (FALSE.equals(filter.condition())) {
+                if (FALSE.equals(filter.condition()) || FoldNull.isNull(filter.condition())) {
                     return new LocalRelation(filter.location(), new EmptyExecutable(filter.output()));
                 }
             }
@@ -1112,6 +1115,41 @@ private boolean canPropagateFoldable(LogicalPlan p) {
         }
     }
 
+    static class FoldNull extends OptimizerExpressionRule {
+
+        FoldNull() {
+            super(TransformDirection.UP);
+        }
+
+        private static boolean isNull(Expression ex) {
+            return DataType.NULL == ex.dataType() || (ex.foldable() && ex.fold() == null);
+        }
+
+        @Override
+        protected Expression rule(Expression e) {
+            if (e instanceof IsNotNull) {
+                if (((IsNotNull) e).field().nullable() == false) {
+                    return new Literal(e.location(), Expressions.name(e), Boolean.TRUE, DataType.BOOLEAN);
+                }
+            }
+            // see https://github.com/elastic/elasticsearch/issues/34876
+            // similar for IsNull once it gets introduced
+
+            if (e instanceof In) {
+                In in = (In) e;
+                if (isNull(in.value())) {
+                    return Literal.of(in, null);
+                }
+            }
+
+            if (e.nullable() && Expressions.anyMatch(e.children(), FoldNull::isNull)) {
+                return Literal.of(e, null);
+            }
+
+            return e;
+        }
+    }
+
     static class ConstantFolding extends OptimizerExpressionRule {
 
         ConstantFolding() {

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.sql.expression.function.aggregate.AggregateFunction;
 import org.elasticsearch.xpack.sql.expression.function.aggregate.Count;
 import org.elasticsearch.xpack.sql.expression.function.scalar.Cast;
+import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DayName;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DayOfMonth;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DayOfYear;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.MonthOfYear;
@@ -28,8 +29,11 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ASin;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ATan;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.Abs;
+import org.elasticsearch.xpack.sql.expression.function.scalar.math.Cos;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.E;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.Floor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Ascii;
+import org.elasticsearch.xpack.sql.expression.function.scalar.string.Repeat;
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryOperator;
 import org.elasticsearch.xpack.sql.expression.predicate.In;
 import org.elasticsearch.xpack.sql.expression.predicate.IsNotNull;
@@ -56,6 +60,7 @@
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.CombineBinaryComparisons;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.CombineProjections;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.ConstantFolding;
+import org.elasticsearch.xpack.sql.optimizer.Optimizer.FoldNull;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.PropagateEquals;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.PruneDuplicateFunctions;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.PruneSubqueryAliases;
@@ -374,10 +379,36 @@ private static Object foldOperator(BinaryOperator<?, ?, ?, ?> b) {
         return ((Literal) new ConstantFolding().rule(b)).value();
     }
 
+    public void testNullFoldingIsNotNull() {
+        assertEquals(Literal.TRUE, new FoldNull().rule(new IsNotNull(EMPTY, Literal.TRUE)));
+    }
+
+    public void testGenericNullableExpression() {
+        FoldNull rule = new FoldNull();
+        // date-time
+        assertNullLiteral(rule.rule(new DayName(EMPTY, Literal.NULL, randomTimeZone())));
+        // math function
+        assertNullLiteral(rule.rule(new Cos(EMPTY, Literal.NULL)));
+        // string function
+        assertNullLiteral(rule.rule(new Ascii(EMPTY, Literal.NULL)));
+        assertNullLiteral(rule.rule(new Repeat(EMPTY, getFieldAttribute(), Literal.NULL)));
+        // arithmetic
+        assertNullLiteral(rule.rule(new Add(EMPTY, getFieldAttribute(), Literal.NULL)));
+        // comparison
+        assertNullLiteral(rule.rule(new GreaterThan(EMPTY, getFieldAttribute(), Literal.NULL)));
+        // regex
+        assertNullLiteral(rule.rule(new RLike(EMPTY, getFieldAttribute(), Literal.NULL)));
+    }
+
     //
     // Logical simplifications
     //
 
+    private void assertNullLiteral(Expression expression) {
+        assertEquals(Literal.class, expression.getClass());
+        assertNull(((Literal) expression).fold());
+    }
+
     public void testBinaryComparisonSimplification() {
         assertEquals(Literal.TRUE, new BinaryComparisonSimplification().rule(new Equals(EMPTY, FIVE, FIVE)));
         assertEquals(Literal.TRUE, new BinaryComparisonSimplification().rule(new GreaterThanOrEqual(EMPTY, FIVE, FIVE)));

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.xpack.sql.expression.function.FunctionRegistry;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer;
 import org.elasticsearch.xpack.sql.parser.SqlParser;
+import org.elasticsearch.xpack.sql.plan.physical.EsQueryExec;
 import org.elasticsearch.xpack.sql.plan.physical.LocalExec;
 import org.elasticsearch.xpack.sql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.sql.session.EmptyExecutable;
@@ -64,6 +65,24 @@ public void testFoldingToLocalExecWithProject() {
         assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
     }
 
+    public void testFoldingOfIsNotNull() {
+        PhysicalPlan p = plan("SELECT keyword FROM test WHERE (keyword IS NULL) IS NOT NULL");
+        assertEquals(EsQueryExec.class, p.getClass());
+        EsQueryExec ee = (EsQueryExec) p;
+        assertEquals(1, ee.output().size());
+        assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
+    }
+
+    public void testFoldingToLocalExecWithNullFilter() {
+        PhysicalPlan p = plan("SELECT keyword FROM test WHERE null IN (1, 2)");
+        assertEquals(LocalExec.class, p.getClass());
+        LocalExec le = (LocalExec) p;
+        assertEquals(EmptyExecutable.class, le.executable().getClass());
+        EmptyExecutable ee = (EmptyExecutable) le.executable();
+        assertEquals(1, ee.output().size());
+        assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
+    }
+
     public void testFoldingToLocalExecWithProject_FoldableIn() {
         PhysicalPlan p = plan("SELECT keyword FROM test WHERE int IN (null, null)");
         assertEquals(LocalExec.class, p.getClass());

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

@@ -3,7 +3,7 @@
 //
 
 nullDate
-SELECT YEAR(CAST(NULL AS DATE)) d;
+SELECT YEAR(CAST(NULL AS DATE)) AS d;
 
 d:i
 null