Support PPL format command

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -30,6 +30,7 @@
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.QueryEngineException;
 import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler;
 
 /**
  * {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the
@@ -137,13 +138,15 @@ public RelNode plan(String query) throws Exception {
   private static class CustomVisitorStrategy implements PlanningStrategy {
     private final UnifiedQueryContext context;
     private final UnifiedQueryParser<UnresolvedPlan> parser;
-    private final CalciteRelNodeVisitor relNodeVisitor =
-        new CalciteRelNodeVisitor(new EmptyDataSourceService());
+    private final CalciteRelNodeVisitor relNodeVisitor;
 
     @SuppressWarnings("unchecked")
     CustomVisitorStrategy(UnifiedQueryContext context) {
       this.context = context;
       this.parser = (UnifiedQueryParser<UnresolvedPlan>) context.getParser();
+      this.relNodeVisitor =
+          new CalciteRelNodeVisitor(
+              new EmptyDataSourceService(), PPLSearchPredicateCompiler.INSTANCE);
     }
 
     @Override

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

@@ -75,6 +75,7 @@
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Flatten;
 import org.opensearch.sql.ast.tree.Foreach;
+import org.opensearch.sql.ast.tree.Format;
 import org.opensearch.sql.ast.tree.GraphLookup;
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Join;
@@ -299,6 +300,9 @@ public LogicalPlan visitLimit(Limit node, AnalysisContext context) {
 
   @Override
   public LogicalPlan visitSearch(Search node, AnalysisContext context) {
+    if (node.hasImplicitSubquery()) {
+      throw getOnlyForCalciteException("Implicit format subsearch");
+    }
     LogicalPlan child = node.getChild().get(0).accept(this, context);
     Function queryStringFunc =
         AstDSL.function(
@@ -833,6 +837,11 @@ public LogicalPlan visitReverse(Reverse node, AnalysisContext context) {
     throw getOnlyForCalciteException("Reverse");
   }
 
+  @Override
+  public LogicalPlan visitFormat(Format node, AnalysisContext context) {
+    throw getOnlyForCalciteException("Format");
+  }
+
   @Override
   public LogicalPlan visitSpath(SPath node, AnalysisContext context) {
     throw getOnlyForCalciteException("Spath");

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

@@ -64,6 +64,7 @@
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Flatten;
 import org.opensearch.sql.ast.tree.Foreach;
+import org.opensearch.sql.ast.tree.Format;
 import org.opensearch.sql.ast.tree.GraphLookup;
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Join;
@@ -166,6 +167,10 @@ public T visitFlatten(Flatten node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitFormat(Format node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitTrendline(Trendline node, C context) {
     return visitChildren(node, context);
   }

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

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+
+/** A subsearch used as a predicate term by the parent {@code search} command. */
+@Getter
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class SearchSubquery extends SearchExpression {
+
+  private final UnresolvedPlan query;
+
+  @Override
+  public String toQueryString() {
+    throw new IllegalStateException(
+        "An implicit search subquery must be planned as a correlated runtime input");
+  }
+
+  @Override
+  public String toAnonymizedString() {
+    return "[ subsearch ]";
+  }
+
+  @Override
+  public List<? extends UnresolvedExpression> getChild() {
+    return List.of();
+  }
+}

core/src/main/java/org/opensearch/sql/ast/tree/Format.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+/** AST node that collapses input rows into a formatted search expression. */
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class Format extends UnresolvedPlan {
+
+  public static final String DEFAULT_MV_SEPARATOR = "OR";
+  public static final int DEFAULT_MAX_RESULTS = 0;
+  public static final String DEFAULT_ROW_PREFIX = "(";
+  public static final String DEFAULT_COLUMN_PREFIX = "(";
+  public static final String DEFAULT_COLUMN_SEPARATOR = "AND";
+  public static final String DEFAULT_COLUMN_END = ")";
+  public static final String DEFAULT_ROW_SEPARATOR = "OR";
+  public static final String DEFAULT_ROW_END = ")";
+  public static final String DEFAULT_EMPTY_STRING = "NOT ()";
+
+  private final String mvSeparator;
+  private final int maxResults;
+  private final String rowPrefix;
+  private final String columnPrefix;
+  private final String columnSeparator;
+  private final String columnEnd;
+  private final String rowSeparator;
+  private final String rowEnd;
+  private final String emptyString;
+
+  /** Whether the result is consumed by an enclosing search rather than returned to the user. */
+  private boolean implicit;
+
+  private UnresolvedPlan child;
+
+  @Override
+  public Format attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return child == null ? ImmutableList.of() : ImmutableList.of(child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitFormat(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/tree/Search.java

@@ -14,6 +14,7 @@
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.SearchExpression;
+import org.opensearch.sql.ast.expression.SearchSubquery;
 
 /**
  * Logical plan node for Search operation. Represents search expressions that get converted to
@@ -26,7 +27,7 @@
 public class Search extends UnresolvedPlan {
 
   @EqualsAndHashCode.Include private final UnresolvedPlan child;
-  @EqualsAndHashCode.Include private final String queryString;
+  @EqualsAndHashCode.Include private final @Nullable String queryString;
 
   // Currently it's only for anonymizer
   private final @Nullable SearchExpression originalExpression;
@@ -35,6 +36,28 @@ public Search(UnresolvedPlan child, String queryString) {
     this(child, queryString, null);
   }
 
+  /** Creates a search that may contain an implicit subquery. */
+  public static Search fromExpression(UnresolvedPlan child, SearchExpression originalExpression) {
+    return new Search(
+        child,
+        containsSubquery(originalExpression) ? null : originalExpression.toQueryString(),
+        originalExpression);
+  }
+
+  public boolean hasImplicitSubquery() {
+    return originalExpression != null && containsSubquery(originalExpression);
+  }
+
+  private static boolean containsSubquery(SearchExpression expression) {
+    if (expression instanceof SearchSubquery) {
+      return true;
+    }
+    return expression.getChild().stream()
+        .filter(SearchExpression.class::isInstance)
+        .map(SearchExpression.class::cast)
+        .anyMatch(Search::containsSubquery);
+  }
+
   @Override
   public List<UnresolvedPlan> getChild() {
     return ImmutableList.of(child);

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -9,10 +9,13 @@
 
 import java.sql.Connection;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.Stack;
 import java.util.function.BiFunction;
 import lombok.Getter;
@@ -21,6 +24,7 @@
 import org.apache.calcite.rex.RexCorrelVariable;
 import org.apache.calcite.rex.RexLambdaRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexSubQuery;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.ast.expression.AggregateFunction;
@@ -83,6 +87,10 @@ public class CalcitePlanContext {
   private final Stack<RexCorrelVariable> correlVar = new Stack<>();
   private final Stack<List<RexNode>> windowPartitions = new Stack<>();
 
+  /** Identity marker for scalar subqueries produced by an implicit format inside search. */
+  private final Set<RexSubQuery> implicitFormatSubqueries =
+      Collections.newSetFromMap(new IdentityHashMap<>());
+
   @Getter public Map<String, RexLambdaRef> rexLambdaRefMap;
 
   /**
@@ -205,6 +213,14 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
     }
   }
 
+  void registerImplicitFormatSubquery(RexSubQuery subquery) {
+    implicitFormatSubqueries.add(subquery);
+  }
+
+  boolean isImplicitFormatSubquery(RexSubQuery subquery) {
+    return implicitFormatSubqueries.contains(subquery);
+  }
+
   /**
    * Creates a clone of this context that shares the relBuilder with the parent. This allows lambda
    * expressions to reference fields from the current row while having their own lambda variable