Optimized type converting in DSL filters (#272)

* optimized cast in filter queries Signed-off-by: chloe-zh <[email protected]> * update Signed-off-by: chloe-zh <[email protected]> * update Signed-off-by: chloe-zh <[email protected]> * update Signed-off-by: chloe-zh <[email protected]> * added unit tests Signed-off-by: chloe-zh <[email protected]> * update Signed-off-by: chloe-zh <[email protected]>
opensearch-project · Nov 9, 2021 · 65f449b · 65f449b
1 parent 8c4b8f2
commit 65f449b
Show file tree

Hide file tree

Showing 6 changed files with 516 additions and 9 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/data/model/ExprTimestampValue.java b/core/src/main/java/org/opensearch/sql/data/model/ExprTimestampValue.java
@@ -56,7 +56,7 @@ public class ExprTimestampValue extends AbstractExprValue {
   /**
    * todo. only support timestamp in format yyyy-MM-dd HH:mm:ss.
    */
-  private static final DateTimeFormatter FORMATTER_WITNOUT_NANO = DateTimeFormatter
+  private static final DateTimeFormatter FORMATTER_WITHOUT_NANO = DateTimeFormatter
       .ofPattern("yyyy-MM-dd HH:mm:ss");
   private final Instant timestamp;
 
@@ -92,7 +92,7 @@ public ExprTimestampValue(String timestamp) {
 
   @Override
   public String value() {
-    return timestamp.getNano() == 0 ? FORMATTER_WITNOUT_NANO.withZone(ZONE)
+    return timestamp.getNano() == 0 ? FORMATTER_WITHOUT_NANO.withZone(ZONE)
         .format(timestamp.truncatedTo(ChronoUnit.SECONDS))
         : FORMATTER_VARIABLE_MICROS.withZone(ZONE).format(timestamp);
   }

diff --git a/core/src/main/java/org/opensearch/sql/expression/operator/convert/TypeCastOperator.java b/core/src/main/java/org/opensearch/sql/expression/operator/convert/TypeCastOperator.java
@@ -101,9 +101,9 @@ private static FunctionResolver castToString() {
   private static FunctionResolver castToByte() {
     return FunctionDSL.define(BuiltinFunctionName.CAST_TO_BYTE.getName(),
         impl(nullMissingHandling(
-            (v) -> new ExprByteValue(Short.valueOf(v.stringValue()))), BYTE, STRING),
+            (v) -> new ExprByteValue(Byte.valueOf(v.stringValue()))), BYTE, STRING),
         impl(nullMissingHandling(
-            (v) -> new ExprByteValue(v.shortValue())), BYTE, DOUBLE),
+            (v) -> new ExprByteValue(v.byteValue())), BYTE, DOUBLE),
         impl(nullMissingHandling(
             (v) -> new ExprByteValue(v.booleanValue() ? 1 : 0)), BYTE, BOOLEAN)
     );

diff --git a/...src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/LuceneQuery.java b/...src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/LuceneQuery.java
@@ -29,14 +29,32 @@
 
 import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.OPENSEARCH_TEXT_KEYWORD;
 
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+import java.util.function.Function;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.sql.data.model.ExprBooleanValue;
+import org.opensearch.sql.data.model.ExprByteValue;
+import org.opensearch.sql.data.model.ExprDateValue;
+import org.opensearch.sql.data.model.ExprDatetimeValue;
+import org.opensearch.sql.data.model.ExprDoubleValue;
+import org.opensearch.sql.data.model.ExprFloatValue;
+import org.opensearch.sql.data.model.ExprIntegerValue;
+import org.opensearch.sql.data.model.ExprLongValue;
+import org.opensearch.sql.data.model.ExprShortValue;
+import org.opensearch.sql.data.model.ExprStringValue;
+import org.opensearch.sql.data.model.ExprTimeValue;
+import org.opensearch.sql.data.model.ExprTimestampValue;
 import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.FunctionName;
 
 /**
  * Lucene query abstraction that builds Lucene query from function expression.
@@ -55,7 +73,8 @@ public abstract class LuceneQuery {
   public boolean canSupport(FunctionExpression func) {
     return (func.getArguments().size() == 2)
         && (func.getArguments().get(0) instanceof ReferenceExpression)
-        && (func.getArguments().get(1) instanceof LiteralExpression)
+        && (func.getArguments().get(1) instanceof LiteralExpression
+        || literalExpressionWrappedByCast(func))
         || isMultiParameterQuery(func);
   }
 
@@ -74,18 +93,144 @@ private boolean isMultiParameterQuery(FunctionExpression func) {
     return true;
   }
 
+  /**
+   * Check if the second argument of the function is a literal expression wrapped by cast function.
+   */
+  private boolean literalExpressionWrappedByCast(FunctionExpression func) {
+    if (func.getArguments().get(1) instanceof FunctionExpression) {
+      FunctionExpression expr = (FunctionExpression) func.getArguments().get(1);
+      return castMap.containsKey(expr.getFunctionName())
+          && expr.getArguments().get(0) instanceof LiteralExpression;
+    }
+    return false;
+  }
+
   /**
    * Build Lucene query from function expression.
+   * The cast function is converted to literal expressions before generating DSL.
    *
    * @param func  function
    * @return      query
    */
   public QueryBuilder build(FunctionExpression func) {
     ReferenceExpression ref = (ReferenceExpression) func.getArguments().get(0);
-    LiteralExpression literal = (LiteralExpression) func.getArguments().get(1);
-    return doBuild(ref.getAttr(), ref.type(), literal.valueOf(null));
+    Expression expr = func.getArguments().get(1);
+    ExprValue literalValue = expr instanceof LiteralExpression ? expr
+        .valueOf(null) : cast((FunctionExpression) expr);
+    return doBuild(ref.getAttr(), ref.type(), literalValue);
   }
 
+  private ExprValue cast(FunctionExpression castFunction) {
+    return castMap.get(castFunction.getFunctionName()).apply(
+        (LiteralExpression) castFunction.getArguments().get(0));
+  }
+
+  /**
+   * Type converting map.
+   */
+  private final Map<FunctionName, Function<LiteralExpression, ExprValue>> castMap = ImmutableMap
+      .<FunctionName, Function<LiteralExpression, ExprValue>>builder()
+      .put(BuiltinFunctionName.CAST_TO_STRING.getName(), expr -> {
+        if (!expr.type().equals(ExprCoreType.STRING)) {
+          return new ExprStringValue(String.valueOf(expr.valueOf(null).value()));
+        } else {
+          return expr.valueOf(null);
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_BYTE.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprByteValue(expr.valueOf(null).byteValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprByteValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprByteValue(Byte.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_SHORT.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprShortValue(expr.valueOf(null).shortValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprShortValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprShortValue(Short.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_INT.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprIntegerValue(expr.valueOf(null).integerValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprIntegerValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprIntegerValue(Integer.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_LONG.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprLongValue(expr.valueOf(null).longValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprLongValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprLongValue(Long.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_FLOAT.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprFloatValue(expr.valueOf(null).floatValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprFloatValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprFloatValue(Float.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_DOUBLE.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return new ExprDoubleValue(expr.valueOf(null).doubleValue());
+        } else if (expr.type().equals(ExprCoreType.BOOLEAN)) {
+          return new ExprDoubleValue(expr.valueOf(null).booleanValue() ? 1 : 0);
+        } else {
+          return new ExprDoubleValue(Double.valueOf(expr.valueOf(null).stringValue()));
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_BOOLEAN.getName(), expr -> {
+        if (ExprCoreType.numberTypes().contains(expr.type())) {
+          return expr.valueOf(null).doubleValue() == 1
+              ? ExprBooleanValue.of(true) : ExprBooleanValue.of(false);
+        } else if (expr.type().equals(ExprCoreType.STRING)) {
+          return ExprBooleanValue.of(Boolean.valueOf(expr.valueOf(null).stringValue()));
+        } else {
+          return expr.valueOf(null);
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_DATE.getName(), expr -> {
+        if (expr.type().equals(ExprCoreType.STRING)) {
+          return new ExprDateValue(expr.valueOf(null).stringValue());
+        } else {
+          return new ExprDateValue(expr.valueOf(null).dateValue());
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_TIME.getName(), expr -> {
+        if (expr.type().equals(ExprCoreType.STRING)) {
+          return new ExprTimeValue(expr.valueOf(null).stringValue());
+        } else {
+          return new ExprTimeValue(expr.valueOf(null).timeValue());
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_DATETIME.getName(), expr -> {
+        if (expr.type().equals(ExprCoreType.STRING)) {
+          return new ExprDatetimeValue(expr.valueOf(null).stringValue());
+        } else {
+          return new ExprDatetimeValue(expr.valueOf(null).datetimeValue());
+        }
+      })
+      .put(BuiltinFunctionName.CAST_TO_TIMESTAMP.getName(), expr -> {
+        if (expr.type().equals(ExprCoreType.STRING)) {
+          return new ExprTimestampValue(expr.valueOf(null).stringValue());
+        } else {
+          return new ExprTimestampValue(expr.valueOf(null).timestampValue());
+        }
+      })
+      .build();
+
   /**
    * Build method that subclass implements by default which is to build query
    * from reference and literal in function arguments.

diff --git a/.../src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/RangeQuery.java b/.../src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/RangeQuery.java
@@ -32,6 +32,7 @@
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.RangeQueryBuilder;
 import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 
 /**
@@ -51,7 +52,7 @@ public enum Comparison {
 
   @Override
   protected QueryBuilder doBuild(String fieldName, ExprType fieldType, ExprValue literal) {
-    Object value = literal.value();
+    Object value = value(literal);
 
     RangeQueryBuilder query = QueryBuilders.rangeQuery(fieldName);
     switch (comparison) {
@@ -68,4 +69,12 @@ protected QueryBuilder doBuild(String fieldName, ExprType fieldType, ExprValue l
     }
   }
 
+  private Object value(ExprValue literal) {
+    if (literal.type().equals(ExprCoreType.TIMESTAMP)) {
+      return literal.timestampValue().toEpochMilli();
+    } else {
+      return literal.value();
+    }
+  }
+
 }
diff --git a/...h/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/TermQuery.java b/...h/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/TermQuery.java
@@ -30,6 +30,7 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 
 /**
@@ -40,7 +41,15 @@ public class TermQuery extends LuceneQuery {
   @Override
   protected QueryBuilder doBuild(String fieldName, ExprType fieldType, ExprValue literal) {
     fieldName = convertTextToKeyword(fieldName, fieldType);
-    return QueryBuilders.termQuery(fieldName, literal.value());
+    return QueryBuilders.termQuery(fieldName, value(literal));
+  }
+
+  private Object value(ExprValue literal) {
+    if (literal.type().equals(ExprCoreType.TIMESTAMP)) {
+      return literal.timestampValue().toEpochMilli();
+    } else {
+      return literal.value();
+    }
   }
 
 }