Add Support for Field Star in Nested Function.

Signed-off-by: forestmvey <[email protected]>

Author: forestmvey

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

@@ -29,6 +29,7 @@
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.commons.math3.analysis.function.Exp;
 import org.opensearch.sql.DataSourceSchemaName;
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
@@ -469,8 +470,13 @@ public LogicalPlan visitSort(Sort node, AnalysisContext context) {
         node.getSortList().stream()
             .map(
                 sortField -> {
-                  Expression expression = optimizer.optimize(
-                      expressionAnalyzer.analyze(sortField.getField(), context), context);
+                  var analyzed = expressionAnalyzer.analyze(sortField.getField(), context);
+                  if (analyzed == null) {
+                    throw new UnsupportedOperationException(
+                        String.format("Invalid use of expression %s", sortField.getField())
+                    );
+                  }
+                  Expression expression = optimizer.optimize(analyzed, context);
                   return ImmutablePair.of(analyzeSortOption(sortField.getFieldArgs()), expression);
                 })
             .collect(Collectors.toList());

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

@@ -186,7 +186,16 @@ public Expression visitFunction(Function node, AnalysisContext context) {
     FunctionName functionName = FunctionName.of(node.getFuncName());
     List<Expression> arguments =
         node.getFuncArgs().stream()
-            .map(unresolvedExpression -> analyze(unresolvedExpression, context))
+            .map(unresolvedExpression -> {
+              var ret = analyze(unresolvedExpression, context);
+              if (ret == null) {
+                throw new UnsupportedOperationException(
+                    String.format("Invalid use of expression %s", unresolvedExpression)
+                );
+              } else {
+                return ret;
+              }
+            })
             .collect(Collectors.toList());
     return (Expression) repository.compile(context.getFunctionProperties(),
         functionName, arguments);

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

@@ -15,8 +15,10 @@
 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.NestedAllTupleFields;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.NamedExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
@@ -43,6 +45,32 @@ public LogicalPlan visitAlias(Alias node, AnalysisContext context) {
     return node.getDelegated().accept(this, context);
   }
 
+  @Override
+  public LogicalPlan visitNestedAllTupleFields(NestedAllTupleFields node, AnalysisContext context) {
+    List<Map<String, ReferenceExpression>> args = new ArrayList<>();
+    for (NamedExpression namedExpr : namedExpressions) {
+
+      if (namedExpr.getDelegated() instanceof FunctionExpression
+          && ((FunctionExpression) namedExpr.getDelegated()).getFunctionName()
+          .getFunctionName().equalsIgnoreCase(BuiltinFunctionName.NESTED.name())) {
+
+        ReferenceExpression field =
+            (ReferenceExpression) ((FunctionExpression) namedExpr.getDelegated())
+                .getArguments().get(0);
+
+        // If path is same as NestedAllTupleFields path
+        if (field.getAttr().substring(0, field.getAttr().lastIndexOf("."))
+            .equalsIgnoreCase(node.getPath())) {
+          args.add(Map.of(
+              "field", field,
+              "path", new ReferenceExpression(node.getPath(), STRING)));
+        }
+      }
+    }
+
+    return mergeChildIfLogicalNested(args);
+  }
+
   @Override
   public LogicalPlan visitFunction(Function node, AnalysisContext context) {
     if (node.getFuncName().equalsIgnoreCase(BuiltinFunctionName.NESTED.name())) {
@@ -52,6 +80,8 @@ public LogicalPlan visitFunction(Function node, AnalysisContext context) {
       ReferenceExpression nestedField =
           (ReferenceExpression)expressionAnalyzer.analyze(expressions.get(0), context);
       Map<String, ReferenceExpression> args;
+
+      // Path parameter is supplied
       if (expressions.size() == 2) {
         args = Map.of(
             "field", nestedField,
@@ -63,16 +93,28 @@ public LogicalPlan visitFunction(Function node, AnalysisContext context) {
             "path", generatePath(nestedField.toString())
         );
       }
-      if (child instanceof LogicalNested) {
-        ((LogicalNested)child).addFields(args);
-        return child;
-      } else {
-        return new LogicalNested(child, new ArrayList<>(Arrays.asList(args)), namedExpressions);
-      }
+
+      return mergeChildIfLogicalNested(new ArrayList<>(Arrays.asList(args)));
     }
     return null;
   }
 
+  /**
+   * NestedAnalyzer visits all functions in SELECT clause, creates logical plans for each and
+   * merges them. This is to avoid another merge rule in LogicalPlanOptimizer:create().
+   * @param args field and path params to add to logical plan.
+   * @return child of logical nested with added args, or new LogicalNested.
+   */
+  private LogicalPlan mergeChildIfLogicalNested(List<Map<String, ReferenceExpression>> args) {
+    if (child instanceof LogicalNested) {
+      for (var arg : args) {
+        ((LogicalNested) child).addFields(arg);
+      }
+      return child;
+    }
+    return new LogicalNested(child, args, namedExpressions);
+  }
+
   /**
    * Validate each parameter used in nested function in SELECT clause. Any supplied parameter
    * for a nested function in a SELECT statement must be a valid qualified name, and the field

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

@@ -10,13 +10,17 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.analysis.symbol.Namespace;
+import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.NestedAllTupleFields;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.data.type.ExprType;
@@ -58,6 +62,11 @@ public List<NamedExpression> visitField(Field node, AnalysisContext context) {
 
   @Override
   public List<NamedExpression> visitAlias(Alias node, AnalysisContext context) {
+    // Expand all nested fields if used in SELECT clause
+    if (node.getDelegated() instanceof NestedAllTupleFields) {
+      return node.getDelegated().accept(this, context);
+    }
+
     Expression expr = referenceIfSymbolDefined(node, context);
     return Collections.singletonList(DSL.named(
         unqualifiedNameIfFieldOnly(node, context),
@@ -100,6 +109,29 @@ public List<NamedExpression> visitAllFields(AllFields node,
         new ReferenceExpression(entry.getKey(), entry.getValue()))).collect(Collectors.toList());
   }
 
+  @Override
+  public List<NamedExpression> visitNestedAllTupleFields(NestedAllTupleFields node,
+                                                   AnalysisContext context) {
+    TypeEnvironment environment = context.peek();
+    Map<String, ExprType> lookupAllTupleFields =
+        environment.lookupAllTupleFields(Namespace.FIELD_NAME);
+    environment.resolve(new Symbol(Namespace.FIELD_NAME, node.getPath()));
+
+    // Match all fields with same path as used in nested function.
+    Pattern p = Pattern.compile(node.getPath() + "\\.[^\\.]+$");
+    return lookupAllTupleFields.entrySet().stream()
+        .filter(field -> p.matcher(field.getKey()).find())
+        .map(entry -> {
+          Expression nestedFunc = new Function(
+              "nested",
+              List.of(
+                  new QualifiedName(List.of(entry.getKey().split("\\."))))
+          ).accept(expressionAnalyzer, context);
+          return DSL.named("nested(" + entry.getKey() + ")", nestedFunc);
+        })
+        .collect(Collectors.toList());
+  }
+
   /**
    * Get unqualified name if select item is just a field. For example, suppose an index
    * named "accounts", return "age" for "SELECT accounts.age". But do nothing for expression

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

@@ -85,6 +85,17 @@ public Map<String, ExprType> lookupAllFields(Namespace namespace) {
     return result;
   }
 
+  /**
+   * Resolve all fields in the current environment.
+   * @param namespace     a namespace
+   * @return              all symbols in the namespace
+   */
+  public Map<String, ExprType> lookupAllTupleFields(Namespace namespace) {
+    Map<String, ExprType> result = new LinkedHashMap<>();
+    symbolTable.lookupAllTupleFields(namespace).forEach(result::putIfAbsent);
+    return result;
+  }
+
   /**
    * Define symbol with the type.
    *

core/src/main/java/org/opensearch/sql/analysis/symbol/SymbolTable.java

@@ -128,6 +128,21 @@ public Map<String, ExprType> lookupAllFields(Namespace namespace) {
     return results;
   }
 
+  /**
+   * Look up all top level symbols in the namespace.
+   *
+   * @param namespace     a namespace
+   * @return              all symbols in the namespace map
+   */
+  public Map<String, ExprType> lookupAllTupleFields(Namespace namespace) {
+    final LinkedHashMap<String, ExprType> allSymbols =
+        orderedTable.getOrDefault(namespace, new LinkedHashMap<>());
+    final LinkedHashMap<String, ExprType> result = new LinkedHashMap<>();
+    allSymbols.entrySet().stream()
+        .forEach(entry -> result.put(entry.getKey(), entry.getValue()));
+    return result;
+  }
+
   /**
    * Check if namespace map in empty (none definition).
    *

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

@@ -25,6 +25,7 @@
 import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.Map;
+import org.opensearch.sql.ast.expression.NestedAllTupleFields;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.QualifiedName;
@@ -238,6 +239,10 @@ public T visitAllFields(AllFields node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitNestedAllTupleFields(NestedAllTupleFields node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitInterval(Interval node, C context) {
     return visitChildren(node, context);
   }

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

@@ -30,6 +30,7 @@
 import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.Map;
+import org.opensearch.sql.ast.expression.NestedAllTupleFields;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.ParseMethod;
@@ -377,6 +378,10 @@ public Alias alias(String name, UnresolvedExpression expr, String alias) {
     return new Alias(name, expr, alias);
   }
 
+  public NestedAllTupleFields nestedAllTupleFields(String path) {
+    return new NestedAllTupleFields(path);
+  }
+
   public static List<UnresolvedExpression> exprList(UnresolvedExpression... exprList) {
     return Arrays.asList(exprList);
   }

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

@@ -0,0 +1,42 @@
+/*
+ * 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;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
+
+/**
+ * Represents all tuple fields used in nested function.
+ */
+@ToString
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+public class NestedAllTupleFields extends UnresolvedExpression {
+  @Getter
+  private final String path;
+
+  @Override
+  public List<? extends Node> getChild() {
+    return Collections.emptyList();
+  }
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitNestedAllTupleFields(this, context);
+  }
+
+  @Override
+  public String toString() {
+    return String.format("nested(%s.*)", path);
+  }
+}