SQL: Extend the optimisations for equalities (#50792) (#51098)

* Extend the optimizations for equalities

This commit supplements the optimisations of equalities in conjunctions
and disjunctions:
* for conjunctions, the existing optimizations with ranges are extended
with not-equalities and inequalities; these lead to a fast resolution,
the conjunction either being evaluate to a FALSE, or the non-equality
conditions being dropped as superfluous;
* optimisations for disjunctions are added to be applied against ranges,
inequalities and not-equalities; these lead to disjunction either
becoming TRUE or the equality being dropped, either as superfluous or
merged into a range/inequality.

* Adress review notes

* Fix the bug around wrongly optimizing 'a=2 OR a!=?', which only yields
TRUE for same values in equality and inequality.
* Var renamings, code style adjustments, comments corrections.

* Address further review comments. Extend optim.

- fix a few code comments;
- extend the Equals OR NotEquals optimitsation (a=2 OR a!=5 -> a!=5);
- extend the Equals OR Range optimisation on limits equality (a=2 OR
  2<=a<5 -> 2<=a<5);
- in case an equality is being removed in a conjunction, the rest of
  possible optimisations to test is now skipped.

* rename one var for better legiblity

- s/rmEqual/removeEquals

(cherry picked from commit 62e7c6a)

Author: bpintea

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

@@ -84,6 +84,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
@@ -814,7 +815,7 @@ private Expression simplifyAndOr(BinaryPredicate<?, ?, ?, ?> bc) {
                 }
 
                 //
-                // common factor extraction -> (a || b) && (a || c) => a && (b || c)
+                // common factor extraction -> (a || b) && (a || c) => a || (b && c)
                 //
                 List<Expression> leftSplit = splitOr(l);
                 List<Expression> rightSplit = splitOr(r);
@@ -852,7 +853,7 @@ private Expression simplifyAndOr(BinaryPredicate<?, ?, ?, ?> bc) {
                 }
 
                 //
-                // common factor extraction -> (a && b) || (a && c) => a || (b & c)
+                // common factor extraction -> (a && b) || (a && c) => a && (b || c)
                 //
                 List<Expression> leftSplit = splitAnd(l);
                 List<Expression> rightSplit = splitAnd(r);
@@ -958,9 +959,11 @@ private Expression literalToTheRight(BinaryOperator<?, ?, ?, ?> be) {
     }
 
     /**
-     * Propagate Equals to eliminate conjuncted Ranges.
-     * When encountering a different Equals or non-containing {@link Range}, the conjunction becomes false.
-     * When encountering a containing {@link Range}, the range gets eliminated by the equality.
+     * Propagate Equals to eliminate conjuncted Ranges or BinaryComparisons.
+     * When encountering a different Equals, non-containing {@link Range} or {@link BinaryComparison}, the conjunction becomes false.
+     * When encountering a containing {@link Range}, {@link BinaryComparison} or {@link NotEquals}, these get eliminated by the equality.
+     *
+     * Since this rule can eliminate Ranges and BinaryComparisons, it should be applied before {@link CombineBinaryComparisons}.
      *
      * This rule doesn't perform any promotion of {@link BinaryComparison}s, that is handled by
      * {@link CombineBinaryComparisons} on purpose as the resulting Range might be foldable
@@ -976,14 +979,20 @@ static class PropagateEquals extends OptimizerExpressionRule {
         protected Expression rule(Expression e) {
             if (e instanceof And) {
                 return propagate((And) e);
+            } else if (e instanceof Or) {
+                return propagate((Or) e);
             }
             return e;
         }
 
         // combine conjunction
         private Expression propagate(And and) {
             List<Range> ranges = new ArrayList<>();
+            // Only equalities, not-equalities and inequalities with a foldable .right are extracted separately;
+            // the others go into the general 'exps'.
             List<BinaryComparison> equals = new ArrayList<>();
+            List<NotEquals> notEquals = new ArrayList<>();
+            List<BinaryComparison> inequalities = new ArrayList<>();
             List<Expression> exps = new ArrayList<>();
 
             boolean changed = false;
@@ -996,35 +1005,42 @@ private Expression propagate(And and) {
                     // equals on different values evaluate to FALSE
                     if (otherEq.right().foldable()) {
                         for (BinaryComparison eq : equals) {
-                            // cannot evaluate equals so skip it
-                            if (!eq.right().foldable()) {
-                                continue;
-                            }
                             if (otherEq.left().semanticEquals(eq.left())) {
-                                if (eq.right().foldable() && otherEq.right().foldable()) {
-                                    Integer comp = BinaryComparison.compare(eq.right().fold(), otherEq.right().fold());
-                                    if (comp != null) {
-                                        // var cannot be equal to two different values at the same time
-                                        if (comp != 0) {
-                                            return FALSE;
-                                        }
+                                Integer comp = BinaryComparison.compare(eq.right().fold(), otherEq.right().fold());
+                                if (comp != null) {
+                                    // var cannot be equal to two different values at the same time
+                                    if (comp != 0) {
+                                        return FALSE;
                                     }
                                 }
                             }
                         }
+                        equals.add(otherEq);
+                    } else {
+                        exps.add(otherEq);
+                    }
+                } else if (ex instanceof GreaterThan || ex instanceof GreaterThanOrEqual ||
+                    ex instanceof LessThan || ex instanceof LessThanOrEqual) {
+                    BinaryComparison bc = (BinaryComparison) ex;
+                    if (bc.right().foldable()) {
+                        inequalities.add(bc);
+                    } else {
+                        exps.add(ex);
+                    }
+                } else if (ex instanceof NotEquals) {
+                    NotEquals otherNotEq = (NotEquals) ex;
+                    if (otherNotEq.right().foldable()) {
+                        notEquals.add(otherNotEq);
+                    } else {
+                        exps.add(ex);
                     }
-                    equals.add(otherEq);
                 } else {
                     exps.add(ex);
                 }
             }
 
             // check
             for (BinaryComparison eq : equals) {
-                // cannot evaluate equals so skip it
-                if (!eq.right().foldable()) {
-                    continue;
-                }
                 Object eqValue = eq.right().fold();
 
                 for (int i = 0; i < ranges.size(); i++) {
@@ -1060,9 +1076,192 @@ private Expression propagate(And and) {
                         changed = true;
                     }
                 }
+
+                // evaluate all NotEquals against the Equal
+                for (Iterator<NotEquals> iter = notEquals.iterator(); iter.hasNext(); ) {
+                    NotEquals neq = iter.next();
+                    if (eq.left().semanticEquals(neq.left())) {
+                        Integer comp = BinaryComparison.compare(eqValue, neq.right().fold());
+                        if (comp != null) {
+                            if (comp == 0) {
+                                return FALSE; // clashing and conflicting: a = 1 AND a != 1
+                            } else {
+                                iter.remove(); // clashing and redundant: a = 1 AND a != 2
+                                changed = true;
+                            }
+                        }
+                    }
+                }
+
+                // evaluate all inequalities against the Equal
+                for (Iterator<BinaryComparison> iter = inequalities.iterator(); iter.hasNext(); ) {
+                    BinaryComparison bc = iter.next();
+                    if (eq.left().semanticEquals(bc.left())) {
+                        Integer compare = BinaryComparison.compare(eqValue, bc.right().fold());
+                        if (compare != null) {
+                            if (bc instanceof LessThan || bc instanceof LessThanOrEqual) { // a = 2 AND a </<= ?
+                                if ((compare == 0 && bc instanceof LessThan) || // a = 2 AND a < 2
+                                    0 < compare) { // a = 2 AND a </<= 1
+                                    return FALSE;
+                                }
+                            } else if (bc instanceof GreaterThan || bc instanceof GreaterThanOrEqual) { // a = 2 AND a >/>= ?
+                                if ((compare == 0 && bc instanceof GreaterThan) || // a = 2 AND a > 2
+                                    compare < 0) { // a = 2 AND a >/>= 3
+                                    return FALSE;
+                                }
+                            }
+
+                            iter.remove();
+                            changed = true;
+                        }
+                    }
+                }
+            }
+
+            return changed ? Predicates.combineAnd(CollectionUtils.combine(exps, equals, notEquals, inequalities, ranges)) : and;
+        }
+
+        // combine disjunction:
+        // a = 2 OR a > 3 -> nop; a = 2 OR a > 1 -> a > 1
+        // a = 2 OR a < 3 -> a < 3; a = 2 OR a < 1 -> nop
+        // a = 2 OR 3 < a < 5 -> nop; a = 2 OR 1 < a < 3 -> 1 < a < 3; a = 2 OR 0 < a < 1 -> nop
+        // a = 2 OR a != 2 -> TRUE; a = 2 OR a = 5 -> nop; a = 2 OR a != 5 -> a != 5
+        private Expression propagate(Or or) {
+            List<Expression> exps = new ArrayList<>();
+            List<Equals> equals = new ArrayList<>(); // foldable right term Equals
+            List<NotEquals> notEquals = new ArrayList<>(); // foldable right term NotEquals
+            List<Range> ranges = new ArrayList<>();
+            List<BinaryComparison> inequalities = new ArrayList<>(); // foldable right term (=limit) BinaryComparision
+
+            // split expressions by type
+            for (Expression ex : Predicates.splitOr(or)) {
+                if (ex instanceof Equals) {
+                    Equals eq = (Equals) ex;
+                    if (eq.right().foldable()) {
+                        equals.add(eq);
+                    } else {
+                        exps.add(ex);
+                    }
+                } else if (ex instanceof NotEquals) {
+                    NotEquals neq = (NotEquals) ex;
+                    if (neq.right().foldable()) {
+                        notEquals.add(neq);
+                    } else {
+                        exps.add(ex);
+                    }
+                } else if (ex instanceof Range) {
+                    ranges.add((Range) ex);
+                } else if (ex instanceof BinaryComparison) {
+                    BinaryComparison bc = (BinaryComparison) ex;
+                    if (bc.right().foldable()) {
+                        inequalities.add(bc);
+                    } else {
+                        exps.add(ex);
+                    }
+                } else {
+                    exps.add(ex);
+                }
+            }
+
+            boolean updated = false; // has the expression been modified?
+
+            // evaluate the impact of each Equal over the different types of Expressions
+            for (Iterator<Equals> iterEq = equals.iterator(); iterEq.hasNext(); ) {
+                Equals eq = iterEq.next();
+                Object eqValue = eq.right().fold();
+                boolean removeEquals = false;
+
+                // Equals OR NotEquals
+                for (NotEquals neq : notEquals) {
+                    if (eq.left().semanticEquals(neq.left())) { // a = 2 OR a != ? -> ...
+                        Integer comp = BinaryComparison.compare(eqValue, neq.right().fold());
+                        if (comp != null) {
+                            if (comp == 0) { // a = 2 OR a != 2 -> TRUE
+                                return TRUE;
+                            } else { // a = 2 OR a != 5 -> a != 5
+                                removeEquals = true;
+                                break;
+                            }
+                        }
+                    }
+                }
+                if (removeEquals) {
+                    iterEq.remove();
+                    updated = true;
+                    continue;
+                }
+
+                // Equals OR Range
+                for (int i = 0; i < ranges.size(); i ++) { // might modify list, so use index loop
+                    Range range = ranges.get(i);
+                    if (eq.left().semanticEquals(range.value())) {
+                        Integer lowerComp = range.lower().foldable() ? BinaryComparison.compare(eqValue, range.lower().fold()) : null;
+                        Integer upperComp = range.upper().foldable() ? BinaryComparison.compare(eqValue, range.upper().fold()) : null;
+
+                        if (lowerComp != null && lowerComp == 0) {
+                            if (!range.includeLower()) { // a = 2 OR 2 < a < ? -> 2 <= a < ?
+                                ranges.set(i, new Range(range.source(), range.value(), range.lower(), true,
+                                    range.upper(), range.includeUpper()));
+                            } // else : a = 2 OR 2 <= a < ? -> 2 <= a < ?
+                            removeEquals = true; // update range with lower equality instead or simply superfluous
+                            break;
+                        } else if (upperComp != null && upperComp == 0) {
+                            if (!range.includeUpper()) { // a = 2 OR ? < a < 2 -> ? < a <= 2
+                                ranges.set(i, new Range(range.source(), range.value(), range.lower(), range.includeLower(),
+                                    range.upper(), true));
+                            } // else : a = 2 OR ? < a <= 2 -> ? < a <= 2
+                            removeEquals = true; // update range with upper equality instead
+                            break;
+                        } else if (lowerComp != null && upperComp != null) {
+                            if (0 < lowerComp && upperComp < 0) { // a = 2 OR 1 < a < 3
+                                removeEquals = true; // equality is superfluous
+                                break;
+                            }
+                        }
+                    }
+                }
+                if (removeEquals) {
+                    iterEq.remove();
+                    updated = true;
+                    continue;
+                }
+
+                // Equals OR Inequality
+                for (int i = 0; i < inequalities.size(); i ++) {
+                    BinaryComparison bc = inequalities.get(i);
+                    if (eq.left().semanticEquals(bc.left())) {
+                        Integer comp = BinaryComparison.compare(eqValue, bc.right().fold());
+                        if (comp != null) {
+                            if (bc instanceof GreaterThan || bc instanceof GreaterThanOrEqual) {
+                                if (comp < 0) { // a = 1 OR a > 2 -> nop
+                                    continue;
+                                } else if (comp == 0 && bc instanceof GreaterThan) { // a = 2 OR a > 2 -> a >= 2
+                                    inequalities.set(i, new GreaterThanOrEqual(bc.source(), bc.left(), bc.right()));
+                                } // else (0 < comp || bc instanceof GreaterThanOrEqual) :
+                                // a = 3 OR a > 2 -> a > 2; a = 2 OR a => 2 -> a => 2
+
+                                removeEquals = true; // update range with equality instead or simply superfluous
+                                break;
+                            } else if (bc instanceof LessThan || bc instanceof LessThanOrEqual) {
+                                if (comp > 0) { // a = 2 OR a < 1 -> nop
+                                    continue;
+                                }
+                                if (comp == 0 && bc instanceof LessThan) { // a = 2 OR a < 2 -> a <= 2
+                                    inequalities.set(i, new LessThanOrEqual(bc.source(), bc.left(), bc.right()));
+                                } // else (comp < 0 || bc instanceof LessThanOrEqual) : a = 2 OR a < 3 -> a < 3; a = 2 OR a <= 2 -> a <= 2
+                                removeEquals = true; // update range with equality instead or simply superfluous
+                                break;
+                            }
+                        }
+                    }
+                }
+                if (removeEquals) {
+                    iterEq.remove();
+                    updated = true;
+                }
             }
 
-            return changed ? Predicates.combineAnd(CollectionUtils.combine(exps, equals, ranges)) : and;
+            return updated ? Predicates.combineOr(CollectionUtils.combine(exps, equals, notEquals, inequalities, ranges)) : or;
         }
     }