SQL: Optimisation fixes for conjunction merges (#50703)

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

@@ -19,7 +19,7 @@
  * In a SQL statement, an Expression is whatever a user specifies inside an
  * action, so for instance:
  *
- * {@code SELECT a, b, MAX(c, d) FROM i}
+ * {@code SELECT a, b, ABS(c) FROM i}
  *
  * a, b, ABS(c), and i are all Expressions, with ABS(c) being a Function
  * (which is a type of expression) with a single child, c.

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

@@ -1620,7 +1620,18 @@ private Expression combine(And and) {
 
             boolean changed = false;
 
-            for (Expression ex : Predicates.splitAnd(and)) {
+            List<Expression> andExps = Predicates.splitAnd(and);
+            // Ranges need to show up before BinaryComparisons in list, to allow the latter be optimized away into a Range, if possible
+            andExps.sort((o1, o2) -> {
+                if (o1 instanceof Range && o2 instanceof Range) {
+                    return 0; // keep ranges' order
+                } else if (o1 instanceof Range || o2 instanceof Range) {
+                    return o2 instanceof Range ? 1 : -1;
+                } else {
+                    return 0; // keep non-ranges' order
+                }
+            });
+            for (Expression ex : andExps) {
                 if (ex instanceof Range) {
                     Range r = (Range) ex;
                     if (findExistingRange(r, ranges, true)) {
@@ -1745,9 +1756,9 @@ private static boolean findExistingRange(Range main, List<Range> ranges, boolean
                             lowerEq = comp == 0 && main.includeLower() == other.includeLower();
                             // AND
                             if (conjunctive) {
-                                        // (2 < a < 3) AND (1 < a < 3) -> (1 < a < 3)
+                                // (2 < a < 3) AND (1 < a < 3) -> (2 < a < 3)
                                 lower = comp > 0 ||
-                                        // (2 < a < 3) AND (2 < a <= 3) -> (2 < a < 3)
+                                // (2 < a < 3) AND (2 <= a < 3) -> (2 < a < 3)
                                         (comp == 0 && !main.includeLower() && other.includeLower());
                             }
                             // OR
@@ -1846,7 +1857,7 @@ private boolean findConjunctiveComparisonInRange(BinaryComparison main, List<Ran
                                     ranges.remove(i);
                                     ranges.add(i,
                                             new Range(other.source(), other.value(),
-                                                    main.right(), lowerEq ? true : other.includeLower(),
+                                                    main.right(), lowerEq ? false : main instanceof GreaterThanOrEqual,
                                                     other.upper(), other.includeUpper()));
                                 }
 
@@ -1855,19 +1866,19 @@ private boolean findConjunctiveComparisonInRange(BinaryComparison main, List<Ran
                             }
                         }
                     } else if (main instanceof LessThan || main instanceof LessThanOrEqual) {
-                        if (other.lower().foldable()) {
-                            Integer comp = BinaryComparison.compare(value, other.lower().fold());
+                        if (other.upper().foldable()) {
+                            Integer comp = BinaryComparison.compare(value, other.upper().fold());
                             if (comp != null) {
                                 // a < 2 AND (1 < a <= 2) -> 1 < a < 2
                                 boolean upperEq = comp == 0 && other.includeUpper() && main instanceof LessThan;
                                 // a < 2 AND (1 < a < 3) -> 1 < a < 2
-                                boolean upper = comp > 0 || upperEq;
+                                boolean upper = comp < 0 || upperEq;
 
                                 if (upper) {
                                     ranges.remove(i);
                                     ranges.add(i, new Range(other.source(), other.value(),
                                             other.lower(), other.includeLower(),
-                                            main.right(), upperEq ? true : other.includeUpper()));
+                                            main.right(), upperEq ? false : main instanceof LessThanOrEqual));
                                 }
 
                                 // found a match

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

@@ -1058,6 +1058,57 @@ public void testCombineBinaryComparisonsInclude() {
         assertEquals(FIVE, r.right());
     }
 
+    // 2 < a AND (2 <= a < 3) -> 2 < a < 3
+    public void testCombineBinaryComparisonsAndRangeLower() {
+        FieldAttribute fa = getFieldAttribute();
+
+        GreaterThan gt = new GreaterThan(EMPTY, fa, TWO);
+        Range range = new Range(EMPTY, fa, TWO, true, THREE, false);
+
+        CombineBinaryComparisons rule = new CombineBinaryComparisons();
+        Expression exp = rule.rule(new And(EMPTY, gt, range));
+        assertEquals(Range.class, exp.getClass());
+        Range r = (Range)exp;
+        assertEquals(TWO, r.lower());
+        assertFalse(r.includeLower());
+        assertEquals(THREE, r.upper());
+        assertFalse(r.includeUpper());
+    }
+
+    // a < 4 AND (1 < a < 3) -> 1 < a < 3
+    public void testCombineBinaryComparisonsAndRangeUpper() {
+        FieldAttribute fa = getFieldAttribute();
+
+        LessThan lt = new LessThan(EMPTY, fa, FOUR);
+        Range range = new Range(EMPTY, fa, ONE, false, THREE, false);
+
+        CombineBinaryComparisons rule = new CombineBinaryComparisons();
+        Expression exp = rule.rule(new And(EMPTY, range, lt));
+        assertEquals(Range.class, exp.getClass());
+        Range r = (Range)exp;
+        assertEquals(ONE, r.lower());
+        assertFalse(r.includeLower());
+        assertEquals(THREE, r.upper());
+        assertFalse(r.includeUpper());
+    }
+
+    // a <= 2 AND (1 < a < 3) -> 1 < a <= 2
+    public void testCombineBinaryComparisonsAndRangeUpperEqual() {
+        FieldAttribute fa = getFieldAttribute();
+
+        LessThanOrEqual lte = new LessThanOrEqual(EMPTY, fa, TWO);
+        Range range = new Range(EMPTY, fa, ONE, false, THREE, false);
+
+        CombineBinaryComparisons rule = new CombineBinaryComparisons();
+        Expression exp = rule.rule(new And(EMPTY, lte, range));
+        assertEquals(Range.class, exp.getClass());
+        Range r = (Range)exp;
+        assertEquals(ONE, r.lower());
+        assertFalse(r.includeLower());
+        assertEquals(TWO, r.upper());
+        assertTrue(r.includeUpper());
+    }
+
     // 3 <= a AND 4 < a AND a <= 7 AND a < 6 -> 4 < a < 6
     public void testCombineMultipleBinaryComparisons() {
         FieldAttribute fa = getFieldAttribute();