Add Nested Function Support in SELECT Clause

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

@@ -359,6 +359,14 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
     List<NamedExpression> namedExpressions =
         selectExpressionAnalyzer.analyze(node.getProjectList(), context,
             new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child));
+
+    for (UnresolvedExpression expr : node.getProjectList()) {
+      NestedAnalyzer nestedAnalyzer = new NestedAnalyzer(
+          namedExpressions, expressionAnalyzer, child
+      );
+      child = nestedAnalyzer.analyze(expr, context);
+    }
+
     // new context
     context.push();
     TypeEnvironment newEnv = context.peek();

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

@@ -325,17 +325,6 @@ private Expression visitIdentifier(String ident, AnalysisContext context) {
     ReferenceExpression ref = DSL.ref(ident,
         typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, ident)));
 
-    // Fall back to old engine too if type is not supported semantically
-    if (isTypeNotSupported(ref.type())) {
-      throw new SyntaxCheckException(String.format(
-          "Identifier [%s] of type [%s] is not supported yet", ident, ref.type()));
-    }
     return ref;
   }
-
-  // Array type is not supporte yet.
-  private boolean isTypeNotSupported(ExprType type) {
-    return "array".equalsIgnoreCase(type.typeName());
-  }
-
 }

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

@@ -0,0 +1,70 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+
+import java.util.List;
+import java.util.Map;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Alias;
+import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.expression.NamedExpression;
+import org.opensearch.sql.expression.ReferenceExpression;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.planner.logical.LogicalNested;
+import org.opensearch.sql.planner.logical.LogicalPlan;
+
+/**
+ * Analyze the Nested Function in the {@link AnalysisContext} to construct the {@link
+ * LogicalPlan}.
+ */
+@RequiredArgsConstructor
+public class NestedAnalyzer extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
+  private final List<NamedExpression> namedExpressions;
+  private final ExpressionAnalyzer expressionAnalyzer;
+  private final LogicalPlan child;
+
+  public LogicalPlan analyze(UnresolvedExpression projectItem, AnalysisContext context) {
+    LogicalPlan nested = projectItem.accept(this, context);
+    return (nested == null) ? child : nested;
+  }
+
+  @Override
+  public LogicalPlan visitAlias(Alias node, AnalysisContext context) {
+    return node.getDelegated().accept(this, context);
+  }
+
+  @Override
+  public LogicalPlan visitFunction(Function node, AnalysisContext context) {
+    if (node.getFuncName().equalsIgnoreCase(BuiltinFunctionName.NESTED.name())) {
+
+      List<UnresolvedExpression> expressions = node.getFuncArgs();
+      ReferenceExpression nestedField =
+          (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(0), context);
+      Map<String, ReferenceExpression> args;
+      if (expressions.size() == 2) {
+        args = Map.of(
+            "field", nestedField,
+            "path", (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(1), context)
+        );
+      } else {
+        args = Map.of(
+            "field", (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(0), context),
+            "path", generatePath(nestedField.toString())
+        );
+      }
+      return new LogicalNested(child, List.of(args), namedExpressions);
+    }
+    return null;
+  }
+
+  private ReferenceExpression generatePath(String field) {
+    return new ReferenceExpression(field.substring(0, field.lastIndexOf(".")), STRING);
+  }
+}

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

@@ -60,7 +60,6 @@
 import org.opensearch.sql.ast.tree.TableFunction;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
-import org.opensearch.sql.expression.function.BuiltinFunctionName;
 
 /**
  * Class of static methods to create specific node instances.

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -624,6 +624,10 @@ public static FunctionExpression xor(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.XOR, expressions);
   }
 
+  public static FunctionExpression nested(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.NESTED, expressions);
+  }
+
   public static FunctionExpression not(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.NOT, expressions);
   }

core/src/main/java/org/opensearch/sql/expression/ReferenceExpression.java

@@ -15,6 +15,7 @@
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.data.model.ExprTupleValue;
 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.env.Environment;
 
@@ -100,7 +101,12 @@ public ExprValue resolve(ExprTupleValue value) {
   }
 
   private ExprValue resolve(ExprValue value, List<String> paths) {
-    final ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, paths));
+    ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, paths));
+    // For array types only first index currently supported.
+    if (value.type().equals(ExprCoreType.ARRAY)) {
+      wholePathValue = value.collectionValue().get(0).keyValue(paths.get(0));
+    }
+
     if (!wholePathValue.isMissing() || paths.size() == 1) {
       return wholePathValue;
     } else {

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -187,6 +187,8 @@ public enum BuiltinFunctionName {
   STDDEV_POP(FunctionName.of("stddev_pop")),
   // take top documents from aggregation bucket.
   TAKE(FunctionName.of("take")),
+  // Not always an aggregation query
+  NESTED(FunctionName.of("nested")),
 
   /**
    * Text Functions.

core/src/main/java/org/opensearch/sql/expression/function/NestedFunctionResolver.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.expression.nested.NestedFunction;
+
+@RequiredArgsConstructor
+public class NestedFunctionResolver
+    implements FunctionResolver {
+
+  @Getter
+  private final FunctionName functionName;
+
+  @Override
+  public Pair<FunctionSignature, FunctionBuilder> resolve(FunctionSignature unresolvedSignature) {
+    if (!unresolvedSignature.getFunctionName().equals(functionName)) {
+      throw new SemanticCheckException(String.format("Expected '%s' but got '%s'",
+          functionName.getFunctionName(), unresolvedSignature.getFunctionName().getFunctionName()));
+    }
+
+    FunctionBuilder buildFunction = (functionProperties, args)
+        -> new NestedFunction(functionName, args);
+    return Pair.of(unresolvedSignature, buildFunction);
+  }
+}

core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java

@@ -40,6 +40,8 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(match_phrase_prefix());
     repository.register(wildcard_query(BuiltinFunctionName.WILDCARD_QUERY));
     repository.register(wildcard_query(BuiltinFunctionName.WILDCARDQUERY));
+    // Functions supported in SELECT clause
+    repository.register(nested(BuiltinFunctionName.NESTED));
   }
 
   private static FunctionResolver match_bool_prefix() {
@@ -86,6 +88,11 @@ private static FunctionResolver wildcard_query(BuiltinFunctionName wildcardQuery
     return new RelevanceFunctionResolver(funcName);
   }
 
+  private static FunctionResolver nested(BuiltinFunctionName nested) {
+    FunctionName funcName = nested.getName();
+    return new NestedFunctionResolver(funcName);
+  }
+
   public static class OpenSearchFunction extends FunctionExpression {
     private final FunctionName functionName;
     private final List<Expression> arguments;

core/src/main/java/org/opensearch/sql/expression/nested/NestedFunction.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.nested;
+
+import java.util.List;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.env.Environment;
+import org.opensearch.sql.expression.function.FunctionName;
+import org.opensearch.sql.expression.function.OpenSearchFunctions;
+
+public class NestedFunction extends OpenSearchFunctions.OpenSearchFunction {
+  private final List<Expression> arguments;
+
+  /**
+   * Required argument constructor.
+   * @param functionName name of the function
+   * @param arguments a list of expressions
+   */
+  public NestedFunction(FunctionName functionName, List<Expression> arguments) {
+    super(functionName, arguments);
+    this.arguments = arguments;
+  }
+
+  @Override
+  public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+    // Only need to resolve field argument which is always first index in member variable.
+    return valueEnv.resolve(this.arguments.get(0));
+  }
+
+  @Override
+  public ExprType type() {
+    return this.arguments.get(0).type();
+  }
+}

core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java

@@ -11,6 +11,7 @@
 import org.opensearch.sql.planner.logical.LogicalEval;
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalLimit;
+import org.opensearch.sql.planner.logical.LogicalNested;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
 import org.opensearch.sql.planner.logical.LogicalProject;
@@ -32,6 +33,7 @@
 import org.opensearch.sql.planner.physical.RemoveOperator;
 import org.opensearch.sql.planner.physical.RenameOperator;
 import org.opensearch.sql.planner.physical.SortOperator;
+import org.opensearch.sql.planner.physical.UnnestOperator;
 import org.opensearch.sql.planner.physical.ValuesOperator;
 import org.opensearch.sql.planner.physical.WindowOperator;
 import org.opensearch.sql.storage.read.TableScanBuilder;
@@ -94,6 +96,11 @@ public PhysicalPlan visitEval(LogicalEval node, C context) {
     return new EvalOperator(visitChild(node, context), node.getExpressions());
   }
 
+  @Override
+  public PhysicalPlan visitUnnest(LogicalNested node, C context) {
+    return new UnnestOperator(visitChild(node, context), node.getFields());
+  }
+
   @Override
   public PhysicalPlan visitSort(LogicalSort node, C context) {
     return new SortOperator(visitChild(node, context), node.getSortList());

core/src/main/java/org/opensearch/sql/planner/logical/LogicalNested.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.logical;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.expression.NamedExpression;
+import org.opensearch.sql.expression.ReferenceExpression;
+
+@EqualsAndHashCode(callSuper = true)
+@Getter
+@ToString
+public class LogicalNested extends LogicalPlan {
+  private final List<Map<String, ReferenceExpression>> fields;
+  private final List<NamedExpression> projectList;
+
+  /**
+   * Constructor of LogicalNested.
+   *
+   */
+  public LogicalNested(
+      LogicalPlan childPlan,
+      List<Map<String, ReferenceExpression>> fields,
+      List<NamedExpression> projectList
+  ) {
+    super(Collections.singletonList(childPlan));
+    this.fields = fields;
+    this.projectList = projectList;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitUnnest(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java

@@ -74,6 +74,14 @@ public LogicalPlan highlight(LogicalPlan input, Expression field,
     return new LogicalHighlight(input, field, arguments);
   }
 
+
+  public static LogicalPlan nested(
+      LogicalPlan input,
+      List<Map<String, ReferenceExpression>> nestedArgs,
+      List<NamedExpression> projectList) {
+    return new LogicalNested(input, nestedArgs, projectList);
+  }
+
   public static LogicalPlan remove(LogicalPlan input, ReferenceExpression... fields) {
     return new LogicalRemove(input, ImmutableSet.copyOf(fields));
   }

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java

@@ -73,6 +73,10 @@ public R visitEval(LogicalEval plan, C context) {
     return visitNode(plan, context);
   }
 
+  public R visitUnnest(LogicalNested plan, C context) {
+    return visitNode(plan, context);
+  }
+
   public R visitSort(LogicalSort plan, C context) {
     return visitNode(plan, context);
   }

core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java

@@ -14,6 +14,7 @@
 import java.util.stream.Collectors;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter;
+import org.opensearch.sql.planner.optimizer.rule.MergeNestedAndNested;
 import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
 import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
 import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown;
@@ -51,11 +52,13 @@ public static LogicalPlanOptimizer create() {
          * Phase 2: Transformations that rely on data source push down capability
          */
         new CreateTableScanBuilder(),
+        new MergeNestedAndNested(),
         TableScanPushDown.PUSH_DOWN_FILTER,
         TableScanPushDown.PUSH_DOWN_AGGREGATION,
         TableScanPushDown.PUSH_DOWN_SORT,
         TableScanPushDown.PUSH_DOWN_LIMIT,
         TableScanPushDown.PUSH_DOWN_HIGHLIGHT,
+        TableScanPushDown.PUSH_DOWN_NESTED,
         TableScanPushDown.PUSH_DOWN_PROJECT,
         new CreateTableWriteBuilder()));
   }

core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java

@@ -16,6 +16,7 @@
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalHighlight;
 import org.opensearch.sql.planner.logical.LogicalLimit;
+import org.opensearch.sql.planner.logical.LogicalNested;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalProject;
 import org.opensearch.sql.planner.logical.LogicalRelation;
@@ -65,6 +66,13 @@ public static <T extends LogicalPlan> Pattern<LogicalHighlight> highlight(Patter
     return Pattern.typeOf(LogicalHighlight.class).with(source(pattern));
   }
 
+  /**
+   * Logical nested operator with a given pattern on inner field.
+   */
+  public static <T extends LogicalPlan> Pattern<LogicalNested> nested(Pattern<T> pattern) {
+    return Pattern.typeOf(LogicalNested.class).with(source(pattern));
+  }
+
   /**
    * Logical project operator with a given pattern on inner field.
    */