Date Autocast for new engine

opendistro-for-elasticsearch · Jun 1, 2021 · f6d62d3 · f6d62d3
1 parent 7eeb805
commit f6d62d3
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 38 deletions.
diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/data/model/ExprStringValue.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/data/model/ExprStringValue.java
@@ -18,6 +18,7 @@
 import com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType;
 import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException;
+import java.time.Instant;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
@@ -80,6 +81,11 @@ public LocalTime timeValue() {
     }
   }
 
+  @Override
+  public Instant timestampValue() {
+    return new ExprTimestampValue(value).timestampValue();
+  }
+
   @Override
   public String toString() {
     return String.format("\"%s\"", value);

diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/data/type/ExprCoreType.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/data/type/ExprCoreType.java
@@ -131,4 +131,8 @@ public static List<ExprCoreType> coreTypes() {
   public static List<ExprType> numberTypes() {
     return ImmutableList.of(INTEGER, LONG, FLOAT, DOUBLE);
   }
+
+  public static List<ExprType> dateTypes() {
+    return ImmutableList.of(DATE, DATETIME, TIMESTAMP, TIME);
+  }
 }
diff --git a/...opendistroforelasticsearch/sql/expression/operator/predicate/BinaryPredicateOperator.java b/...opendistroforelasticsearch/sql/expression/operator/predicate/BinaryPredicateOperator.java
@@ -25,16 +25,26 @@
 
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprBooleanValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
 import com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.BuiltinFunctionName;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.BuiltinFunctionRepository;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionBuilder;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionDSL;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionName;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionResolver;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionSignature;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.SerializableFunction;
 import com.amazon.opendistroforelasticsearch.sql.utils.OperatorUtils;
 import com.google.common.collect.ImmutableTable;
 import com.google.common.collect.Table;
+import java.io.Serializable;
+import java.util.List;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
+import org.apache.commons.lang3.tuple.Pair;
+
 
 /**
  * The definition of binary predicate function
@@ -192,53 +202,44 @@ private static FunctionResolver notEqual() {
   }
 
   private static FunctionResolver less() {
-    return FunctionDSL
-        .define(BuiltinFunctionName.LESS.getName(), ExprCoreType.coreTypes().stream()
-            .map(type -> FunctionDSL
-                .impl(FunctionDSL
-                        .nullMissingHandling((v1, v2) -> ExprBooleanValue.of(v1.compareTo(v2) < 0)),
-                    BOOLEAN,
-                    type, type))
-            .collect(
-                Collectors.toList()));
+    return FunctionDSL.define(BuiltinFunctionName.LESS.getName(),
+        getComparisonFunctions((Function<Integer, Boolean> & Serializable) value -> value < 0));
   }
 
   private static FunctionResolver lte() {
-    return FunctionDSL
-        .define(BuiltinFunctionName.LTE.getName(), ExprCoreType.coreTypes().stream()
-            .map(type -> FunctionDSL
-                .impl(
-                    FunctionDSL
-                        .nullMissingHandling(
-                            (v1, v2) -> ExprBooleanValue.of(v1.compareTo(v2) <= 0)),
-                    BOOLEAN,
-                    type, type))
-            .collect(
-                Collectors.toList()));
+    return FunctionDSL.define(BuiltinFunctionName.LTE.getName(),
+        getComparisonFunctions((Function<Integer, Boolean> & Serializable) value -> value <= 0));
   }
 
   private static FunctionResolver greater() {
-    return FunctionDSL
-        .define(BuiltinFunctionName.GREATER.getName(), ExprCoreType.coreTypes().stream()
-            .map(type -> FunctionDSL
-                .impl(FunctionDSL
-                        .nullMissingHandling((v1, v2) -> ExprBooleanValue.of(v1.compareTo(v2) > 0)),
-                    BOOLEAN, type, type))
-            .collect(
-                Collectors.toList()));
+    return FunctionDSL.define(BuiltinFunctionName.GREATER.getName(),
+        getComparisonFunctions((Function<Integer, Boolean> & Serializable) value -> value > 0));
   }
 
   private static FunctionResolver gte() {
-    return FunctionDSL
-        .define(BuiltinFunctionName.GTE.getName(), ExprCoreType.coreTypes().stream()
-            .map(type -> FunctionDSL
-                .impl(
-                    FunctionDSL.nullMissingHandling(
-                        (v1, v2) -> ExprBooleanValue.of(v1.compareTo(v2) >= 0)),
-                    BOOLEAN,
-                    type, type))
-            .collect(
-                Collectors.toList()));
+    return FunctionDSL.define(BuiltinFunctionName.GTE.getName(),
+        getComparisonFunctions((Function<Integer, Boolean> & Serializable) value -> value >= 0));
+  }
+
+  private static List<SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>>
+      getComparisonFunctions(Function<Integer, Boolean> f) {
+    List<SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>>
+        comparisonFunctions = ExprCoreType.coreTypes()
+        .stream()
+        .map(type -> FunctionDSL.impl(FunctionDSL.nullMissingHandling(
+            (v1, v2) -> ExprBooleanValue.of(f.apply(v1.compareTo(v2)))), BOOLEAN, type, type))
+        .collect(Collectors.toList());
+
+    // all string to date conversions
+    comparisonFunctions.addAll(ExprCoreType.dateTypes()
+        .stream()
+        .map(type -> FunctionDSL.impl(FunctionDSL.nullMissingHandling((v1, v2) -> {
+          // casting V2 from string to whatever type v1 is
+          ExprValue v2Casted = ExprValueUtils.fromObjectValue(v2.value(), (ExprCoreType) v1.type());
+          return ExprBooleanValue.of(f.apply(v1.compareTo(v2Casted)));
+        }), BOOLEAN, type, STRING))
+        .collect(Collectors.toList()));
+    return comparisonFunctions;
   }
 
   private static FunctionResolver like() {

diff --git a/...distroforelasticsearch/sql/expression/operator/predicate/BinaryPredicateOperatorTest.java b/...distroforelasticsearch/sql/expression/operator/predicate/BinaryPredicateOperatorTest.java
@@ -40,12 +40,16 @@
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprBooleanValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprByteValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprCollectionValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDateValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDatetimeValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDoubleValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprFloatValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprIntegerValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprLongValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprShortValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprStringValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimeValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimestampValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTupleValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
@@ -148,6 +152,34 @@ private static Stream<Arguments> testCompareValueArguments() {
     builder.add(Arguments.of(new ExprShortValue(1), new ExprShortValue(1)));
     builder.add(Arguments.of(new ExprShortValue(1), new ExprShortValue(2)));
     builder.add(Arguments.of(new ExprShortValue(2), new ExprShortValue(1)));
+
+    builder.add(Arguments.of(new ExprDatetimeValue("2006-01-01 00:00:01"),
+        new ExprStringValue("2006-01-01 00:00:00")));
+    builder.add(Arguments.of(new ExprDatetimeValue("2006-01-01 00:00:00"),
+        new ExprStringValue("2006-01-01 00:00:01")));
+    builder.add(Arguments.of(new ExprDatetimeValue("2006-01-01 00:00:00"),
+        new ExprStringValue("2006-01-01 00:00:00")));
+
+    builder.add(Arguments.of(new ExprDateValue("2006-01-01"),
+        new ExprStringValue("2006-01-01")));
+    builder.add(Arguments.of(new ExprDateValue("2006-01-01"),
+        new ExprStringValue("2006-01-02")));
+    builder.add(Arguments.of(new ExprDateValue("2006-01-02"),
+        new ExprStringValue("2006-01-01")));
+
+    builder.add(Arguments.of(new ExprTimestampValue("2006-01-01 00:00:01"),
+        new ExprStringValue("2006-01-01 00:00:00")));
+    builder.add(Arguments.of(new ExprTimestampValue("2006-01-01 00:00:00"),
+        new ExprStringValue("2006-01-01 00:00:01")));
+    builder.add(Arguments.of(new ExprTimestampValue("2006-01-01 00:00:00"),
+        new ExprStringValue("2006-01-01 00:00:00")));
+
+    builder.add(Arguments.of(new ExprTimeValue("00:00:01"),
+        new ExprStringValue("00:00:00")));
+    builder.add(Arguments.of(new ExprTimeValue("00:00:00"),
+        new ExprStringValue("00:00:01")));
+    builder.add(Arguments.of(new ExprTimeValue("00:00:00"),
+        new ExprStringValue("00:00:00")));
     return builder.build();
   }
 

diff --git a/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/utils/ComparisonUtil.java b/core/src/test/java/com/amazon/opendistroforelasticsearch/sql/utils/ComparisonUtil.java
@@ -22,13 +22,19 @@
 import static com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils.getStringValue;
 
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprByteValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDateValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDatetimeValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDoubleValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprFloatValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprIntegerValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprLongValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprShortValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprStringValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimeValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprTimestampValue;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType;
 import com.amazon.opendistroforelasticsearch.sql.exception.ExpressionEvaluationException;
 
 public class ComparisonUtil {
@@ -57,6 +63,14 @@ public static int compare(ExprValue v1, ExprValue v2) {
       return getDoubleValue(v1).compareTo(getDoubleValue(v2));
     } else if (v1 instanceof ExprStringValue) {
       return getStringValue(v1).compareTo(getStringValue(v2));
+    } else if (v1 instanceof ExprDateValue) {
+      return v1.dateValue().compareTo(v2.dateValue());
+    } else if (v1 instanceof ExprDatetimeValue) {
+      return v1.datetimeValue().compareTo(v2.datetimeValue());
+    } else if (v1 instanceof ExprTimeValue) {
+      return v1.timeValue().compareTo(v2.timeValue());
+    } else if (v1 instanceof ExprTimestampValue) {
+      return v1.timestampValue().compareTo(v2.timestampValue());
     } else {
       throw new ExpressionEvaluationException(
           String.format("%s instances are not comparable", v1.getClass().getSimpleName()));

diff --git a/...istroforelasticsearch/sql/elasticsearch/storage/script/filter/FilterQueryBuilderTest.java b/...istroforelasticsearch/sql/elasticsearch/storage/script/filter/FilterQueryBuilderTest.java
@@ -16,8 +16,12 @@
 
 package com.amazon.opendistroforelasticsearch.sql.elasticsearch.storage.script.filter;
 
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.DATE;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.DATETIME;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRING;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIME;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIMESTAMP;
 import static com.amazon.opendistroforelasticsearch.sql.elasticsearch.data.type.ElasticsearchDataType.ES_TEXT_KEYWORD;
 import static com.amazon.opendistroforelasticsearch.sql.expression.DSL.literal;
 import static com.amazon.opendistroforelasticsearch.sql.expression.DSL.ref;
@@ -26,19 +30,36 @@
 import static org.mockito.Mockito.doAnswer;
 
 import com.amazon.opendistroforelasticsearch.sql.common.utils.StringUtils;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprByteValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprDoubleValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprFloatValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprIntegerValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprLongValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprShortValue;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.elasticsearch.storage.serialization.ExpressionSerializer;
 import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.FunctionExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.config.ExpressionConfig;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.BuiltinFunctionName;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Stream;
 import org.json.JSONObject;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayNameGeneration;
 import org.junit.jupiter.api.DisplayNameGenerator;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -99,6 +120,45 @@ void should_build_range_query_for_comparison_expression() {
             buildQuery(expr)));
   }
 
+  private static List<Arguments> validDateAutocast() {
+
+    return Arrays.asList(
+        Arguments.of(DATE, "2006-01-01"),
+        Arguments.of(TIME, "00:00:00"),
+        Arguments.of(TIMESTAMP, "2006-01-01 00:00:00"),
+        Arguments.of(DATETIME, "2006-01-01 00:00:00")
+    );
+  }
+
+  @ParameterizedTest
+  @MethodSource("validDateAutocast")
+  void shouldBuildRangeQueryForDateComparisonExpressionAutocast(ExprType exprType,
+      final String time) {
+    Expression[] params = {ref("date", exprType), literal(time)};
+
+    String serializedString = "\"" + time + "\"";
+    Map<Expression, Object[]> ranges = ImmutableMap.of(
+        dsl.less(params), new Object[]{null, serializedString, true, false},
+        dsl.greater(params), new Object[]{serializedString, null, false, true},
+        dsl.lte(params), new Object[]{null, serializedString, true, true},
+        dsl.gte(params), new Object[]{serializedString, null, true, true});
+
+    ranges.forEach((expr, range) ->
+        assertJsonEquals(
+            "{\n"
+                + "  \"range\" : {\n"
+                + "    \"date\" : {\n"
+                + "      \"from\" : " + range[0] + ",\n"
+                + "      \"to\" : " + range[1] + ",\n"
+                + "      \"include_lower\" : " + range[2] + ",\n"
+                + "      \"include_upper\" : " + range[3] + ",\n"
+                + "      \"boost\" : 1.0\n"
+                + "    }\n"
+                + "  }\n"
+                + "}",
+            buildQuery(expr)));
+  }
+
   @Test
   void should_build_wildcard_query_for_like_expression() {
     assertJsonEquals(