Add position() to V2 engine (#177)

Signed-off-by: Margarit Hakobyan <[email protected]> Rbased
Bit-Quill · Dec 5, 2022 · fc1bccf · fc1bccf
1 parent 3786f04
commit fc1bccf
Show file tree

Hide file tree

Showing 7 changed files with 76 additions and 3 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java b/core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java
@@ -34,6 +34,7 @@
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
+import org.opensearch.sql.ast.expression.PositionFunction;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.Span;
@@ -201,6 +202,13 @@ public Expression visitHighlightFunction(HighlightFunction node, AnalysisContext
     return new HighlightExpression(expr);
   }
 
+  @Override
+  public Expression visitPositionFunction(PositionFunction node, AnalysisContext context) {
+    Expression stringPatternExpr = node.getStringPatternExpr().accept(this, context);
+    Expression searchStringExpr = node.getSearchStringExpr().accept(this, context);
+    return DSL.position(stringPatternExpr, searchStringExpr);
+  }
+
   @Override
   public Expression visitIn(In node, AnalysisContext context) {
     return visitIn(node.getField(), node.getValueList(), context);

diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java
@@ -26,6 +26,7 @@
 import org.opensearch.sql.ast.expression.Map;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
+import org.opensearch.sql.ast.expression.PositionFunction;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.Span;
@@ -273,6 +274,10 @@ public T visitHighlightFunction(HighlightFunction node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitPositionFunction(PositionFunction node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitStatement(Statement node, C context) {
     return visit(node, context);
   }

diff --git a/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java b/core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java
@@ -32,6 +32,7 @@
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.ParseMethod;
+import org.opensearch.sql.ast.expression.PositionFunction;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.SpanUnit;
@@ -283,6 +284,11 @@ public UnresolvedExpression highlight(UnresolvedExpression fieldName,
     return new HighlightFunction(fieldName, arguments);
   }
 
+  public UnresolvedExpression position(UnresolvedExpression stringPatternExpr,
+                                       UnresolvedExpression searchStringExpr) {
+    return new PositionFunction(stringPatternExpr, searchStringExpr);
+  }
+
   public UnresolvedExpression window(UnresolvedExpression function,
                                      List<UnresolvedExpression> partitionByList,
                                      List<Pair<SortOption, UnresolvedExpression>> sortList) {

diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/PositionFunction.java b/core/src/main/java/org/opensearch/sql/ast/expression/PositionFunction.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.Arrays;
+import java.util.List;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+
+/**
+ * Expression node of Position function.
+ */
+@AllArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@Getter
+@ToString
+public class PositionFunction extends UnresolvedExpression {
+  @Getter
+  private UnresolvedExpression stringPatternExpr;
+  @Getter
+  private UnresolvedExpression searchStringExpr;
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return Arrays.asList(stringPatternExpr, searchStringExpr);
+  }
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitPositionFunction(this, context);
+  }
+}
diff --git a/core/src/main/java/org/opensearch/sql/expression/DSL.java b/core/src/main/java/org/opensearch/sql/expression/DSL.java
@@ -234,6 +234,10 @@ public static FunctionExpression position(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.POSITION, expressions);
   }
 
+  public static FunctionExpression position(Expression... expressions) {
+    return compile(BuiltinFunctionName.POSITION, expressions);
+  }
+
   public static FunctionExpression truncate(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.TRUNCATE, expressions);
   }

diff --git a/core/src/test/java/org/opensearch/sql/analysis/NamedExpressionAnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/NamedExpressionAnalyzerTest.java
@@ -16,6 +16,7 @@
 import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.Literal;
+import org.opensearch.sql.ast.expression.PositionFunction;
 import org.opensearch.sql.expression.NamedExpression;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.test.context.ContextConfiguration;
@@ -48,4 +49,14 @@ void visit_highlight() {
     NamedExpression analyze = analyzer.analyze(alias, analysisContext);
     assertEquals("highlight(fieldA)", analyze.getNameOrAlias());
   }
+
+  @Test
+  void visit_position() {
+    Alias alias = AstDSL.alias("position(fieldA IN fieldB)",
+        new PositionFunction(AstDSL.stringLiteral("fieldA"), AstDSL.stringLiteral("fieldB")));
+    NamedExpressionAnalyzer analyzer = new NamedExpressionAnalyzer(expressionAnalyzer);
+
+    NamedExpression analyze = analyzer.analyze(alias, analysisContext);
+    assertEquals("position(fieldA IN fieldB)", analyze.getNameOrAlias());
+  }
 }
diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4
@@ -321,9 +321,9 @@ highlightFunction
     : HIGHLIGHT LR_BRACKET relevanceField (COMMA highlightArg)* RR_BRACKET
     ;
 
-    positionFunction
-        : POSITION LR_BRACKET functionArg IN functionArg RR_BRACKET
-        ;
+positionFunction
+    : POSITION LR_BRACKET functionArg IN functionArg RR_BRACKET
+    ;
 
 scalarFunctionName
     : mathematicalFunctionName