Compare temporal IN/NOT IN values in the field's timestamp domain instead of string-collapsing them

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java

@@ -228,27 +228,30 @@ public RexNode visitIn(In node, CalcitePlanContext context) {
     final RexNode field = analyze(node.getField(), context);
     final List<RexNode> valueList =
         node.getValueList().stream().map(value -> analyze(value, context)).toList();
+    // When the field is a temporal type, do NOT use leastRestrictive + rexBuilder.makeIn. For a
+    // temporal field tested against string/date literals, leastRestrictive collapses the common
+    // type to VARCHAR (the EXPR_DATE / EXPR_TIMESTAMP UDTs are VARCHAR-backed), so makeIn casts the
+    // field DOWN to VARCHAR and string-compares mismatched renderings (e.g. '2018-06-23 00:00:00'
+    // against '2018-06-23') — silently matching nothing. Rewrite the membership test as an OR of
+    // PPL `=` comparisons, the same temporal-aware comparison path visitCompare takes for `=`, so
+    // each value is coerced to the field's timestamp domain before comparison.
+    ExprType fieldExprType = OpenSearchTypeFactory.convertRelDataTypeToExprType(field.getType());
+    if (TEMPORAL_TYPES.contains(fieldExprType)) {
+      List<RexNode> equalities =
+          valueList.stream()
+              .map(value -> PPLFuncImpTable.INSTANCE.resolve(context.rexBuilder, "=", field, value))
+              .toList();
+      return context.relBuilder.or(equalities);
+    }
     final List<RelDataType> dataTypes =
-        new java.util.ArrayList<>(valueList.stream().map(RexNode::getType).toList());
+        new ArrayList<>(valueList.stream().map(RexNode::getType).toList());
     dataTypes.add(field.getType());
     RelDataType commonType = context.rexBuilder.getTypeFactory().leastRestrictive(dataTypes);
     if (commonType != null) {
       List<RexNode> newValueList =
           valueList.stream().map(value -> context.rexBuilder.makeCast(commonType, value)).toList();
       return context.rexBuilder.makeIn(field, newValueList);
     }
-    // leastRestrictive() has no common type for mixed temporal representations — e.g. a standard
-    // Calcite TIMESTAMP field tested against EXPR_DATE UDT values (`ts in (date('...'), ...)`).
-    // Mirror visitBetween: widen all-temporal operands through CoercionUtils so the comparison
-    // runs, while leaving genuinely incompatible mixes to raise SemanticCheckException below.
-    List<RexNode> operands = new ArrayList<>(valueList);
-    operands.add(field);
-    List<RexNode> widened = widenTemporalOperands(context, operands);
-    if (widened != null) {
-      RexNode widenedField = widened.get(widened.size() - 1);
-      List<RexNode> widenedValues = widened.subList(0, widened.size() - 1);
-      return context.rexBuilder.makeIn(widenedField, widenedValues);
-    }
     List<ExprType> exprTypes =
         dataTypes.stream().map(OpenSearchTypeFactory::convertRelDataTypeToExprType).toList();
     throw new SemanticCheckException(
@@ -257,6 +260,24 @@ public RexNode visitIn(In node, CalcitePlanContext context) {
             exprTypes.getLast(), exprTypes.subList(0, exprTypes.size() - 1)));
   }
 
+  private static final Set<ExprType> TEMPORAL_TYPES =
+      Set.of(ExprCoreType.DATE, ExprCoreType.TIME, ExprCoreType.TIMESTAMP);
+
+  @Override
+  public RexNode visitCompare(Compare node, CalcitePlanContext context) {
+    RexNode left = analyze(node.getLeft(), context);
+    RexNode right = analyze(node.getRight(), context);
+    String op = node.getOperator();
+    // Handle boolean_field != literal -> IS_NOT_TRUE/IS_NOT_FALSE
+    if ("!=".equals(op) || "<>".equals(op)) {
+      RexNode result = tryMakeBooleanNotEquals(left, right, context);
+      if (result != null) {
+        return result;
+      }
+    }
+    return PPLFuncImpTable.INSTANCE.resolve(context.rexBuilder, op, left, right);
+  }
+
   /**
    * Widens a set of operands to a common temporal type when, and only when, every operand is a
    * temporal type (DATE / TIME / TIMESTAMP), including the EXPR_DATE / EXPR_TIME / EXPR_TIMESTAMP
@@ -276,24 +297,6 @@ public RexNode visitIn(In node, CalcitePlanContext context) {
     return CoercionUtils.widenArguments(context.rexBuilder, operands);
   }
 
-  private static final Set<ExprType> TEMPORAL_TYPES =
-      Set.of(ExprCoreType.DATE, ExprCoreType.TIME, ExprCoreType.TIMESTAMP);
-
-  @Override
-  public RexNode visitCompare(Compare node, CalcitePlanContext context) {
-    RexNode left = analyze(node.getLeft(), context);
-    RexNode right = analyze(node.getRight(), context);
-    String op = node.getOperator();
-    // Handle boolean_field != literal -> IS_NOT_TRUE/IS_NOT_FALSE
-    if ("!=".equals(op) || "<>".equals(op)) {
-      RexNode result = tryMakeBooleanNotEquals(left, right, context);
-      if (result != null) {
-        return result;
-      }
-    }
-    return PPLFuncImpTable.INSTANCE.resolve(context.rexBuilder, op, left, right);
-  }
-
   @Override
   public RexNode visitBetween(Between node, CalcitePlanContext context) {
     RexNode value = analyze(node.getValue(), context);

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLBasicIT.java

@@ -517,8 +517,10 @@ public void testDateBetween() throws IOException {
 
   @Test
   public void testDateIn() throws IOException {
-    // birthdate is a TIMESTAMP-typed field; the IN values are DATE literals — same mixed-temporal
-    // coercion path as testDateBetween, exercised through visitIn.
+    // birthdate is a TIMESTAMP-typed field; the IN values are DATE literals. visitIn must compare
+    // each value in the field's temporal domain (via PPL `=`) rather than letting leastRestrictive
+    // collapse the common type to VARCHAR and string-compare mismatched renderings — which silently
+    // matched nothing without pushdown. See visitIn in CalciteRexNodeVisitor.
     JSONObject actual =
         executeQuery(
             String.format(
@@ -530,6 +532,38 @@ public void testDateIn() throws IOException {
         actual, rows("Nanette", "2018-06-23 00:00:00"), rows("Elinor", "2018-06-27 00:00:00"));
   }
 
+  @Test
+  public void testDateNotIn() throws IOException {
+    // Complement of testDateIn: NOT IN over a temporal field. Exercises the complemented-points
+    // pushdown branch, which must also format timestamp values rather than emit a flat terms query.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | where birthdate not in (DATE '2018-06-23', DATE '2018-06-27') |"
+                    + " fields firstname | sort firstname",
+                TEST_INDEX_BANK));
+    verifySchema(actual, schema("firstname", "string"));
+    verifyDataRows(
+        actual,
+        rows("Amber JOHnny"),
+        rows("Dale"),
+        rows("Dillard"),
+        rows("Hattie"),
+        rows("Virginia"));
+  }
+
+  @Test
+  public void testIn() throws IOException {
+    // Non-temporal IN keeps the leastRestrictive + makeIn path; guards against the temporal-field
+    // special-case in visitIn regressing membership tests over ordinary columns.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | where age in (32, 36) | fields firstname, age", TEST_INDEX_BANK));
+    verifySchema(actual, schema("firstname", "string"), schema("age", "int"));
+    verifyDataRows(actual, rows("Amber JOHnny", 32), rows("Hattie", 36), rows("Elinor", 36));
+  }
+
   @Test
   public void testXor() throws IOException {
     JSONObject result =

opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java

@@ -784,16 +784,29 @@ private QueryExpression like(RexCall call) {
 
     private static QueryExpression constructQueryExpressionForSearch(
         RexCall call, SwapResult pair) {
+      boolean isTimeStamp =
+          (pair.getKey() instanceof NamedFieldExpression namedField)
+              && namedField.isTimeStampType();
       if (isSearchWithComplementedPoints(call)) {
+        // A NOT IN over a timestamp field cannot use a flat terms query: timestamp values must be
+        // normalized + formatted (see in()/termsQuery, which does neither). Decompose into the
+        // complement of an OR of formatted per-point range queries, mirroring the range branch.
+        if (isTimeStamp) {
+          return pointsAsTimestampOr(call, pair).not();
+        }
         return QueryExpression.create(pair.getKey()).notIn(pair.getValue());
       } else if (isSearchWithPoints(call)) {
+        // IN over a timestamp field: the flat terms query in in() emits unformatted values like
+        // '2018-06-23 00:00:00' that the date field cannot parse. Decompose into an OR of formatted
+        // per-point range queries (equals(point, true) builds gte/lte + format="date_time"), the
+        // same path the range branch below and visitCompare's `=` already take.
+        if (isTimeStamp) {
+          return pointsAsTimestampOr(call, pair);
+        }
         return QueryExpression.create(pair.getKey()).in(pair.getValue());
       } else {
         Sarg<?> sarg = pair.getValue().literal.getValueAs(Sarg.class);
         Set<? extends Range<?>> rangeSet = requireNonNull(sarg).rangeSet.asRanges();
-        boolean isTimeStamp =
-            (pair.getKey() instanceof NamedFieldExpression namedField)
-                && namedField.isTimeStampType();
         List<QueryExpression> queryExpressions =
             rangeSet.stream()
                 .map(
@@ -811,6 +824,30 @@ private static QueryExpression constructQueryExpressionForSearch(
       }
     }
 
+    /**
+     * Builds an OR of per-point {@code equals(point, true)} query expressions for a points-only
+     * Sarg over a timestamp field. Each point becomes a {@code gte/lte} range query with {@code
+     * format="date_time"}, so the timestamp value is normalized to the date field's accepted format
+     * — unlike a flat terms query, which emits unformatted values the date field rejects.
+     */
+    private static QueryExpression pointsAsTimestampOr(RexCall call, SwapResult pair) {
+      RexLiteral literal = (RexLiteral) call.getOperands().get(1);
+      Sarg<?> sarg = requireNonNull(literal.getValueAs(Sarg.class), "Sarg");
+      Set<? extends Range<?>> ranges =
+          (sarg.isComplementedPoints() ? sarg.negate() : sarg).rangeSet.asRanges();
+      List<QueryExpression> queryExpressions =
+          ranges.stream()
+              .map(
+                  range ->
+                      QueryExpression.create(pair.getKey())
+                          .equals(sargPointValue(range.lowerEndpoint()), true))
+              .toList();
+      if (queryExpressions.size() == 1) {
+        return queryExpressions.getFirst();
+      }
+      return CompoundQueryExpression.or(queryExpressions.toArray(new QueryExpression[0]));
+    }
+
     private QueryExpression andOr(RexCall call) {
       QueryExpression[] expressions = new QueryExpression[call.getOperands().size()];
       PredicateAnalyzerException firstError = null;