Expression pushdown optimization (#663)

* Test serializer

* Add JDK serializer impl

* Add UT

* Support boolean literal

* Add UT for comparison expression

* Add UT for function expression

* Add UT for multiple fields

* Use expr value factory

* Add comments

* Test coverage

* Refactor package and class

* Add UT for script factory and leaf factory

* Add support for date doc value

* More UT

* Refactor edge case UT

* Add UT for value factory

* Add UT for serializer error case

* Fix checkstyle

* Add expression visitor and UT

* Use expression visitor in script engine

* Add UT for visitor

* Push down query

* Handle text keyword

* Cast long/double doc value to int/float

* Don't push down if illegal state exception thrown

* Fix broken IT due to field type change

* Prepare PR

* Address PR comments

* Address PR comments

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/Expression.java

@@ -34,4 +34,15 @@ public interface Expression extends Serializable {
    * The type of the expression.
    */
   ExprType type();
+
+  /**
+   * Accept a visitor to visit current expression node.
+   * @param visitor   visitor
+   * @param context   context
+   * @param <T>       result type
+   * @param <C>       context type
+   * @return          result accumulated by visitor when visiting
+   */
+  <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context);
+
 }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/ExpressionNodeVisitor.java

@@ -0,0 +1,77 @@
+/*
+ *    Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License").
+ *    You may not use this file except in compliance with the License.
+ *    A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    or in the "license" file accompanying this file. This file is distributed
+ *    on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *    express or implied. See the License for the specific language governing
+ *    permissions and limitations under the License.
+ *
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.expression;
+
+import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
+import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionImplementation;
+
+/**
+ * Abstract visitor for expression tree nodes.
+ * @param <T>   type of return value to accumulate when visiting.
+ * @param <C>   type of context.
+ */
+public abstract class ExpressionNodeVisitor<T, C> {
+
+  public T visitNode(Expression node, C context) {
+    return null;
+  }
+
+  /**
+   * Visit children nodes in function arguments.
+   * @param node      function node
+   * @param context   context
+   * @return          result
+   */
+  public T visitChildren(FunctionImplementation node, C context) {
+    T result = defaultResult();
+
+    for (Expression child : node.getArguments()) {
+      T childResult = child.accept(this, context);
+      result = aggregateResult(result, childResult);
+    }
+    return result;
+  }
+
+  private T defaultResult() {
+    return null;
+  }
+
+  private T aggregateResult(T aggregate, T nextResult) {
+    return nextResult;
+  }
+
+  public T visitLiteral(LiteralExpression node, C context) {
+    return visitNode(node, context);
+  }
+
+  public T visitNamed(NamedExpression node, C context) {
+    return visitNode(node, context);
+  }
+
+  public T visitReference(ReferenceExpression node, C context) {
+    return visitNode(node, context);
+  }
+
+  public T visitFunction(FunctionExpression node, C context) {
+    return visitChildren(node, context);
+  }
+
+  public T visitAggregator(Aggregator<?> node, C context) {
+    return visitChildren(node, context);
+  }
+
+}

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/FunctionExpression.java

@@ -35,4 +35,10 @@ public abstract class FunctionExpression implements Expression, FunctionImplemen
 
   @Getter
   private final List<Expression> arguments;
+
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitFunction(this, context);
+  }
+
 }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/LiteralExpression.java

@@ -39,6 +39,11 @@ public ExprType type() {
     return exprValue.type();
   }
 
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitLiteral(this, context);
+  }
+
   @Override
   public String toString() {
     return exprValue.toString();

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/NamedExpression.java

@@ -71,4 +71,9 @@ public String getName() {
     return Strings.isNullOrEmpty(alias) ? name : alias;
   }
 
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitNamed(this, context);
+  }
+
 }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/ReferenceExpression.java

@@ -40,6 +40,11 @@ public ExprType type() {
     return type;
   }
 
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitReference(this, context);
+  }
+
   @Override
   public String toString() {
     return attr;

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/aggregation/Aggregator.java

@@ -21,6 +21,7 @@
 import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.exception.ExpressionEvaluationException;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
+import com.amazon.opendistroforelasticsearch.sql.expression.ExpressionNodeVisitor;
 import com.amazon.opendistroforelasticsearch.sql.expression.env.Environment;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionImplementation;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionName;
@@ -70,4 +71,10 @@ public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
   public ExprType type() {
     return returnType;
   }
+
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitAggregator(this, context);
+  }
+
 }

core/src/test/java/com/amazon/opendistroforelasticsearch/sql/expression/ExpressionNodeVisitorTest.java

@@ -0,0 +1,89 @@
+/*
+ *    Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License").
+ *    You may not use this file except in compliance with the License.
+ *    A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    or in the "license" file accompanying this file. This file is distributed
+ *    on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *    express or implied. See the License for the specific language governing
+ *    permissions and limitations under the License.
+ *
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.expression;
+
+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.expression.DSL.literal;
+import static com.amazon.opendistroforelasticsearch.sql.expression.DSL.named;
+import static com.amazon.opendistroforelasticsearch.sql.expression.DSL.ref;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
+import com.amazon.opendistroforelasticsearch.sql.expression.config.ExpressionConfig;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.DisplayNameGeneration;
+import org.junit.jupiter.api.DisplayNameGenerator;
+import org.junit.jupiter.api.Test;
+
+@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
+class ExpressionNodeVisitorTest {
+
+  private final DSL dsl = new ExpressionConfig().dsl(new ExpressionConfig().functionRepository());
+
+  @Test
+  void should_return_null_by_default() {
+    ExpressionNodeVisitor<Object, Object> visitor = new ExpressionNodeVisitor<Object, Object>(){};
+    assertNull(literal(10).accept(visitor, null));
+    assertNull(ref("name", STRING).accept(visitor, null));
+    assertNull(named("bool", literal(true)).accept(visitor, null));
+    assertNull(dsl.abs(literal(-10)).accept(visitor, null));
+    assertNull(dsl.sum(literal(10)).accept(visitor, null));
+  }
+
+  @Test
+  void can_visit_all_types_of_expression_node() {
+    Expression expr =
+        dsl.sum(
+            dsl.add(
+                ref("balance", INTEGER),
+                literal(10)));
+
+    Expression actual = expr.accept(new ExpressionNodeVisitor<Expression, Object>() {
+      @Override
+      public Expression visitLiteral(LiteralExpression node, Object context) {
+        return node;
+      }
+
+      @Override
+      public Expression visitReference(ReferenceExpression node, Object context) {
+        return node;
+      }
+
+      @Override
+      public Expression visitFunction(FunctionExpression node, Object context) {
+        return dsl.add(visitArguments(node.getArguments(), context));
+      }
+
+      @Override
+      public Expression visitAggregator(Aggregator<?> node, Object context) {
+        return dsl.sum(visitArguments(node.getArguments(), context));
+      }
+
+      private Expression[] visitArguments(List<Expression> arguments, Object context) {
+        return arguments.stream()
+                        .map(arg -> arg.accept(this, context))
+                        .toArray(Expression[]::new);
+      }
+    }, null);
+
+    assertEquals(expr, actual);
+  }
+
+}

elasticsearch/src/main/java/com/amazon/opendistroforelasticsearch/sql/elasticsearch/data/value/ElasticsearchExprTextKeywordValue.java

@@ -0,0 +1,40 @@
+/*
+ *    Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License").
+ *    You may not use this file except in compliance with the License.
+ *    A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    or in the "license" file accompanying this file. This file is distributed
+ *    on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *    express or implied. See the License for the specific language governing
+ *    permissions and limitations under the License.
+ *
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.elasticsearch.data.value;
+
+import static com.amazon.opendistroforelasticsearch.sql.elasticsearch.data.type.ElasticsearchDataType.ES_TEXT_KEYWORD;
+
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprStringValue;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
+
+/**
+ * Expression Text Keyword Value, it is an extension of the ExprValue by Elasticsearch.
+ * This mostly represents a multi-field in Elasticsearch which has a text field and a
+ * keyword field inside to preserve the original text.
+ */
+public class ElasticsearchExprTextKeywordValue extends ExprStringValue {
+
+  public ElasticsearchExprTextKeywordValue(String value) {
+    super(value);
+  }
+
+  @Override
+  public ExprType type() {
+    return ES_TEXT_KEYWORD;
+  }
+
+}