Search Command Revamp (#4152)

Signed-off-by: Vamsi Manohar <[email protected]>

Author: vamsimanohar

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -86,6 +86,7 @@
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Rex;
+import org.opensearch.sql.ast.tree.Search;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.SubqueryAlias;
@@ -278,6 +279,18 @@ public LogicalPlan visitLimit(Limit node, AnalysisContext context) {
     return new LogicalLimit(child, node.getLimit(), node.getOffset());
   }
 
+  @Override
+  public LogicalPlan visitSearch(Search node, AnalysisContext context) {
+    LogicalPlan child = node.getChild().get(0).accept(this, context);
+    Function queryStringFunc =
+        AstDSL.function(
+            "query_string",
+            AstDSL.unresolvedArg("query", AstDSL.stringLiteral(node.getQueryString())));
+
+    Expression analyzed = expressionAnalyzer.analyze(queryStringFunc, context);
+    return new LogicalFilter(child, analyzed);
+  }
+
   @Override
   public LogicalPlan visitFilter(Filter node, AnalysisContext context) {
     LogicalPlan child = node.getChild().get(0).accept(this, context);

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -75,6 +75,7 @@
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.SPath;
+import org.opensearch.sql.ast.tree.Search;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.SubqueryAlias;
 import org.opensearch.sql.ast.tree.TableFunction;
@@ -131,6 +132,10 @@ public T visitTableFunction(TableFunction node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitSearch(Search node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitFilter(Filter node, C context) {
     return visitChildren(node, context);
   }

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -71,6 +71,7 @@
 import org.opensearch.sql.ast.tree.RelationSubquery;
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.SPath;
+import org.opensearch.sql.ast.tree.Search;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.SpanBin;
@@ -109,6 +110,10 @@ public UnresolvedPlan describe(String tableName) {
     return new DescribeRelation(qualifiedName(tableName));
   }
 
+  public static UnresolvedPlan search(UnresolvedPlan input, String queryString) {
+    return new Search(input, queryString);
+  }
+
   public UnresolvedPlan subqueryAlias(UnresolvedPlan child, String alias) {
     return new SubqueryAlias(child, alias);
   }

core/src/main/java/org/opensearch/sql/ast/expression/SearchAnd.java

@@ -0,0 +1,34 @@
+/*
+ * 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.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/** Search expression for AND operator. */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchAnd extends SearchExpression {
+
+  private final SearchExpression left;
+  private final SearchExpression right;
+
+  @Override
+  public String toQueryString() {
+    return left.toQueryString() + " AND " + right.toQueryString();
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return Arrays.asList(left, right);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchComparison.java

@@ -0,0 +1,70 @@
+/*
+ * 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.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.utils.QueryStringUtils;
+
+/** Search expression for field comparisons. */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchComparison extends SearchExpression {
+
+  public enum Operator {
+    EQUALS("="),
+    NOT_EQUALS("!="),
+    LESS_THAN("<"),
+    LESS_OR_EQUAL("<="),
+    GREATER_THAN(">"),
+    GREATER_OR_EQUAL(">=");
+
+    private final String symbol;
+
+    Operator(String symbol) {
+      this.symbol = symbol;
+    }
+
+    public String getSymbol() {
+      return symbol;
+    }
+  }
+
+  private final Field field;
+  private final Operator operator;
+  private final SearchLiteral value;
+
+  @Override
+  public String toQueryString() {
+    String fieldName = QueryStringUtils.escapeFieldName(field.getField().toString());
+    String valueStr = value.toQueryString();
+    switch (operator) {
+      case NOT_EQUALS:
+        return "( _exists_:" + fieldName + " AND NOT " + fieldName + ":" + valueStr + " )";
+      case GREATER_THAN:
+        return fieldName + ":>" + valueStr;
+      case GREATER_OR_EQUAL:
+        return fieldName + ":>=" + valueStr;
+      case LESS_THAN:
+        return fieldName + ":<" + valueStr;
+      case LESS_OR_EQUAL:
+        return fieldName + ":<=" + valueStr;
+      default:
+        return fieldName + ":" + valueStr;
+    }
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return Arrays.asList(field, value);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchExpression.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+/** Base class for search expressions that get converted to query_string syntax. */
+public abstract class SearchExpression extends UnresolvedExpression {
+
+  /**
+   * Convert this search expression to query_string syntax.
+   *
+   * @return the query string representation
+   */
+  public abstract String toQueryString();
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitChildren(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchGroup.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.Collections;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/** Search expression for grouped expressions (parentheses). */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchGroup extends SearchExpression {
+
+  private final SearchExpression expression;
+
+  @Override
+  public String toQueryString() {
+    return "(" + expression.toQueryString() + ")";
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return Collections.singletonList(expression);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchIn.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.utils.QueryStringUtils;
+
+/** Search expression for IN operator. */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchIn extends SearchExpression {
+
+  private final Field field;
+  private final List<SearchLiteral> values;
+
+  @Override
+  public String toQueryString() {
+    String fieldName = QueryStringUtils.escapeFieldName(field.getField().toString());
+    String valueList =
+        values.stream().map(SearchLiteral::toQueryString).collect(Collectors.joining(" OR "));
+
+    return fieldName + ":( " + valueList + " )";
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    List<UnresolvedExpression> children = new ArrayList<>();
+    children.add(field);
+    // SearchLiteral extends SearchExpression which extends UnresolvedExpression
+    children.addAll(values);
+    return children;
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchLiteral.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.Collections;
+import java.util.List;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.utils.QueryStringUtils;
+
+/** Search expression for standalone literals. */
+@Getter
+@AllArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchLiteral extends SearchExpression {
+
+  private final UnresolvedExpression literal;
+  private final boolean isPhrase;
+
+  @Override
+  public String toQueryString() {
+    if (literal instanceof Literal) {
+      Literal lit = (Literal) literal;
+      Object val = lit.getValue();
+
+      // Numbers don't need escaping
+      if (val instanceof Number) {
+        return val.toString();
+      }
+
+      // Strings
+      if (val instanceof String) {
+        String str = (String) val;
+
+        // Phrase search - preserve quotes
+        if (isPhrase) {
+          // Escape special chars inside the phrase
+          str = QueryStringUtils.escapeLuceneSpecialCharacters(str);
+          return "\"" + str + "\"";
+        }
+
+        // Regular string - escape special characters
+        return QueryStringUtils.escapeLuceneSpecialCharacters(str);
+      }
+    }
+
+    // Default: escape the text representation
+    String text = literal.toString();
+    return QueryStringUtils.escapeLuceneSpecialCharacters(text);
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return Collections.singletonList(literal);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/SearchNot.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.Collections;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/** Search expression for NOT operator. */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchNot extends SearchExpression {
+
+  private final SearchExpression expression;
+
+  @Override
+  public String toQueryString() {
+    return "NOT(" + expression.toQueryString() + ")";
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return Collections.singletonList(expression);
+  }
+}