[CALCITE-7194] Simplify comparisons between function calls and literals to SEARCH

1. Generalize SargCollector in RexSimplify to handle comparisons with deterministic expressions.
2. Add Javadoc for accept variants in SargCollector
3. Prevent invalid SEARCH to interval/range transformations in DruidDateTimeUtils

Some plan changes in DruidAdapterIT/DruidAdapter2IT are due to the added restrictions in DruidDateTimeUtils.
When the SEARCH operand is not a plain column reference (RexInputRef) its generally unsafe to convert it to an interval; the entire DruidDateTimeUtils.createInterval was not meant to handle arbitrary complex expressions.

Author: zabetak

core/src/main/java/org/apache/calcite/rex/RexSimplify.java

@@ -3175,13 +3175,28 @@ static class SargCollector {
       this.negate = negate;
     }
 
+    /**
+     * Accepts an expression and converts it to a Sarg if possible.
+     *
+     * @param term the expression to convert
+     * @param newTerms the list holding the result of the conversion or the
+     *                 original term if it cannot be converted
+     */
     private void accept(RexNode term, List<RexNode> newTerms) {
       if (!accept_(term, newTerms)) {
         newTerms.add(term);
       }
       newTermsCount = newTerms.size();
     }
 
+    /**
+     * Accepts an expression and converts it to a Sarg if possible.
+     * Only certain kinds of expressions can be converted to Sargs.
+     *
+     * @param e the expression to convert
+     * @param newTerms the list to which the Sarg will be added if the expression is accepted
+     * @return true if the expression can be converted to a Sarg, false otherwise
+     */
     private boolean accept_(RexNode e, List<RexNode> newTerms) {
       switch (e.getKind()) {
       case LESS_THAN:
@@ -3204,31 +3219,25 @@ private boolean accept_(RexNode e, List<RexNode> newTerms) {
       }
     }
 
+    /**
+     * Accepts two operands from a binary comparison operator and converts them to a Sarg if
+     * possible. Only comparisons between a literal and a deterministic expression can be converted.
+     * Simplifications with non-deterministic expressions are generally avoided to ensure consistent
+     * results.
+     *
+     * @param left the left operand of the comparison
+     * @param right the right operand of the comparison
+     * @param kind the kind of comparison operator
+     * @param newTerms the list to which the Sarg will be added if accepted
+     * @return true if the operands can be converted to a Sarg, false otherwise
+     */
     private boolean accept2(RexNode left, RexNode right, SqlKind kind,
         List<RexNode> newTerms) {
-      switch (left.getKind()) {
-      case INPUT_REF:
-      case FIELD_ACCESS:
-      case CAST:
-        switch (right.getKind()) {
-        case LITERAL:
-          return accept2b(left, kind, (RexLiteral) right, newTerms);
-        default:
-          break;
-        }
-        return false;
-      case LITERAL:
-        switch (right.getKind()) {
-        case INPUT_REF:
-        case FIELD_ACCESS:
-        case CAST:
-          return accept2b(right, kind.reverse(), (RexLiteral) left, newTerms);
-        default:
-          break;
-        }
-        return false;
-      default:
-        break;
+      if (right.isA(SqlKind.LITERAL) && RexUtil.isDeterministic(left)) {
+        return accept2b(left, kind, (RexLiteral) right, newTerms);
+      }
+      if (left.isA(SqlKind.LITERAL) && RexUtil.isDeterministic(right)) {
+        return accept2b(right, kind.reverse(), (RexLiteral) left, newTerms);
       }
       return false;
     }
@@ -3238,7 +3247,14 @@ private static <E> E addFluent(List<? super E> list, E e) {
       return e;
     }
 
-    // always returns true
+    /**
+     * Accepts an operand from a null predicate and converts it to a Sarg.
+     *
+     * @param e the operand of the null predicate
+     * @param kind the kind of the null predicate
+     * @param newTerms the list to which the Sarg is added
+     * @return true since the operand is always converted to a Sarg
+     */
     private boolean accept1(RexNode e, SqlKind kind, List<RexNode> newTerms) {
       final RexSargBuilder b =
           map.computeIfAbsent(e, e2 ->
@@ -3256,6 +3272,17 @@ private boolean accept1(RexNode e, SqlKind kind, List<RexNode> newTerms) {
       }
     }
 
+    /**
+     * Accepts two operands from a binary comparison operator and converts them to a Sarg when the
+     * literal is not null. The conversion depends on the kind of comparison operator and only
+     * certain operators are supported.
+     *
+     * @param e any kind of deterministic expression
+     * @param kind the kind of comparison operator
+     * @param literal the literal operand of the comparison
+     * @param newTerms the list to which the Sarg is added if accepted
+     * @return false if the literal operand is null, true otherwise
+     */
     private boolean accept2b(RexNode e, SqlKind kind,
         RexLiteral literal, List<RexNode> newTerms) {
       if (literal.getValue() == null) {

core/src/test/java/org/apache/calcite/rex/RexProgramTest.java

@@ -2040,6 +2040,22 @@ private void checkExponentialCnf(int n) {
     checkSimplify(expr, simplified);
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7194">[CALCITE-7194]
+   * Simplify comparisons between function calls and literals to SEARCH</a>. */
+  @Test void testSimplifyComparisonsOfCallsWithLiteralsToSearch() {
+    RexNode v1 = input(tInt(), 1);
+    RexNode v2 = input(tInt(), 2);
+    RexNode plus = plus(v1, v2);
+    checkSimplify(and(gt(plus, literal(100000)), lt(plus, literal(200000))),
+        "SEARCH(+($1, $2), Sarg[(100000..200000)])");
+    checkSimplify(or(eq(plus, literal(100000)), eq(plus, literal(200000))),
+        "SEARCH(+($1, $2), Sarg[100000, 200000])");
+    RexNode ndc = rexBuilder.makeCall(getNoDeterministicOperator(), v1, v2);
+    checkSimplifyUnchanged(or(eq(ndc, literal(100000)), eq(ndc, literal(200000))));
+    checkSimplifyUnchanged(and(gt(ndc, literal(100000)), lt(ndc, literal(200000))));
+  }
+
   @Test void testSimplifySearchWithSinglePointSargToEquals() {
     Range<BigDecimal> r100 = Range.singleton(BigDecimal.valueOf(100));
     RangeSet<BigDecimal> rangeSet = ImmutableRangeSet.<BigDecimal>builder().add(r100).build();

core/src/test/java/org/apache/calcite/test/RelMetadataTest.java

@@ -4097,8 +4097,8 @@ private void assertExpressionLineage(
   @Test void testExpressionLineageBetweenExpressionWithJoin() {
     String sql = "select dept.deptno + empno between 1 and 2"
         + " from emp join dept on emp.deptno = dept.deptno";
-    String expected = "[AND(>=(+([CATALOG, SALES, DEPT].#0.$0, [CATALOG, SALES, EMP].#0.$0), 1),"
-        + " <=(+([CATALOG, SALES, DEPT].#0.$0, [CATALOG, SALES, EMP].#0.$0), 2))]";
+    String expected =
+        "[SEARCH(+([CATALOG, SALES, DEPT].#0.$0, [CATALOG, SALES, EMP].#0.$0), Sarg[[1..2]])]";
     String comment = "'empno' is column 0 in 'catalog.sales.emp', "
         + "'deptno' is column 0 in 'catalog.sales.dept', and "
         + "'dept.deptno + empno between 1 and 2' is translated into "

druid/src/main/java/org/apache/calcite/adapter/druid/DruidDateTimeUtils.java

@@ -237,6 +237,9 @@ && literalValue(call.getOperands().get(3)) != null) {
       return ranges.build();
 
     case SEARCH:
+      if (!canTransformSearchToRange(call)) {
+        return null;
+      }
       final RexLiteral right = (RexLiteral) call.operands.get(1);
       final Sarg<?> sarg = requireNonNull(right.getValueAs(Sarg.class));
       ranges = ImmutableList.builder();
@@ -255,6 +258,28 @@ && literalValue(call.getOperands().get(3)) != null) {
     }
   }
 
+  /**
+   * Returns whether the given SEARCH call can be transformed to a Druid range.
+   *
+   * @param call a SEARCH call
+   * @return whether the given SEARCH call can be transformed to a Druid range.
+   */
+  private static boolean canTransformSearchToRange(RexCall call) {
+    assert call.getKind() == SqlKind.SEARCH;
+    if (call.getOperands().get(0) instanceof RexInputRef) {
+      SqlTypeName literalType = call.operands.get(1).getType().getSqlTypeName();
+      switch (literalType) {
+      case TIMESTAMP:
+      case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+      case DATE:
+        return true;
+      default:
+        return false;
+      }
+    }
+    return false;
+  }
+
   private static Long toLong(Comparable comparable) {
     if (comparable instanceof TimestampString) {
       TimestampString timestampString = (TimestampString) comparable;

druid/src/test/java/org/apache/calcite/test/DruidAdapter2IT.java

@@ -1882,7 +1882,7 @@ private void checkGroupBySingleSortLimit(boolean approx) {
     final String plan = "PLAN=EnumerableInterpreter\n"
         + "  DruidQuery(table=[[foodmart, foodmart]], intervals=[[1900-01-09T00:00:00.000Z/"
         + "2992-01-10T00:00:00.000Z]], filter=[AND(=($2, 'Bird Call'), "
-        + "OR(=(EXTRACT(FLAG(WEEK), $0), 10), =(EXTRACT(FLAG(WEEK), $0), 11)))], "
+        + "SEARCH(EXTRACT(FLAG(WEEK), $0), Sarg[10L:BIGINT, 11L:BIGINT]:BIGINT))], "
         + "projects=[[$63, $90, $91]], "
         + "groups=[{0}], aggs=[[SUM($1), SUM($2)]], "
         + "post_projects=[[$0, 'Bird Call', -($1, $2)]])";
@@ -3054,7 +3054,8 @@ private void testCountWithApproxDistinct(boolean approx, String sql,
             + " CAST('1997-01-01' as DATE) GROUP BY  floor(\"timestamp\" to DAY) order by d limit 3";
     final String plan = "PLAN=EnumerableInterpreter\n"
         + "  DruidQuery(table=[[foodmart, foodmart]], "
-        + "intervals=[[1997-01-01T00:00:00.000Z/1997-02-01T00:00:00.000Z]], "
+        + "intervals=[[1900-01-09T00:00:00.000Z/2992-01-10T00:00:00.000Z]], "
+        + "filter=[=(FLOOR(CAST($0):DATE NOT NULL, FLAG(MONTH)), 1997-01-01)], "
         + "projects=[[FLOOR($0, FLAG(DAY))]], groups=[{0}], aggs=[[]], sort0=[0], "
         + "dir0=[ASC], fetch=[3])";
     sql(sql)

druid/src/test/java/org/apache/calcite/test/DruidAdapterIT.java

@@ -2188,8 +2188,9 @@ private void checkGroupBySingleSortLimit(boolean approx) {
     final String plan = "PLAN=EnumerableInterpreter\n"
         + "  DruidQuery(table=[[foodmart, foodmart]], "
         + "intervals=[[1900-01-09T00:00:00.000Z/2992-01-10T00:00:00.000Z]], "
-        + "filter=[AND(=($2, 'Bird Call'), OR(=(EXTRACT(FLAG(WEEK), $0), 10), "
-        + "=(EXTRACT(FLAG(WEEK), $0), 11)))], projects=[[$63, $90, $91]], "
+        + "filter=[AND(=($2, 'Bird Call'), "
+        + "SEARCH(EXTRACT(FLAG(WEEK), $0), Sarg[10L:BIGINT, 11L:BIGINT]:BIGINT))], "
+        + "projects=[[$63, $90, $91]], "
         + "groups=[{0}], aggs=[[SUM($1), SUM($2)]], post_projects=[[$0, 'Bird Call', -($1, $2)]])";
     sql(sql, FOODMART)
         .returnsOrdered("store_state=CA; brand_name=Bird Call; A=34.3646",
@@ -3696,7 +3697,8 @@ private void testCountWithApproxDistinct(boolean approx, String sql, String expe
             + " CAST('1997-01-01' as DATE) GROUP BY  floor(\"timestamp\" to DAY) order by d limit 3";
     final String plan = "PLAN=EnumerableInterpreter\n"
         + "  DruidQuery(table=[[foodmart, foodmart]], "
-        + "intervals=[[1997-01-01T00:00:00.000Z/1997-02-01T00:00:00.000Z]], "
+        + "intervals=[[1900-01-09T00:00:00.000Z/2992-01-10T00:00:00.000Z]], "
+        + "filter=[=(FLOOR(CAST($0):DATE NOT NULL, FLAG(MONTH)), 1997-01-01)], "
         + "projects=[[FLOOR($0, FLAG(DAY))]], groups=[{0}], aggs=[[]], "
         + "post_projects=[[CAST($0):TIMESTAMP(0) NOT NULL]], sort0=[0], dir0=[ASC], fetch=[3])";
     sql(sql, FOODMART)

geode/src/test/java/org/apache/calcite/adapter/geode/rel/GeodeBookstoreTest.java

@@ -486,7 +486,7 @@ private CalciteAssert.AssertThat calciteAssert() {
 
   @Test void testInSetFilterWithNestedStringField() {
     String expectedQuery = "SELECT primaryAddress.city AS city FROM /BookCustomer "
-        + "WHERE primaryAddress.city IN SET('Topeka', 'San Francisco')";
+        + "WHERE primaryAddress.city IN SET('San Francisco', 'Topeka')";
 
     calciteAssert()
         .query("SELECT primaryAddress['city'] AS city\n"