Support expand with alias

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

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

@@ -97,8 +97,8 @@ public static Eval eval(UnresolvedPlan input, Let... projectList) {
     return new Eval(Arrays.asList(projectList)).attach(input);
   }
 
-  public Expand expand(UnresolvedPlan input, Field field) {
-    return new Expand(field).attach(input);
+  public Expand expand(UnresolvedPlan input, Field field, String alias) {
+    return new Expand(field, alias).attach(input);
   }
 
   public static UnresolvedPlan projectWithArg(

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

@@ -7,6 +7,7 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import javax.annotation.Nullable;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
@@ -21,6 +22,7 @@ public class Expand extends UnresolvedPlan {
 
   private UnresolvedPlan child;
   @Getter private final Field field;
+  @Getter @Nullable private final String alias;
 
   @Override
   public Expand attach(UnresolvedPlan child) {

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

@@ -842,9 +842,10 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
 
     RelBuilder relBuilder = context.relBuilder;
 
-    // 2. Get the field to expand
+    // 2. Get the field to expand and an optional alias.
     Field arrayField = expand.getField();
     RexInputRef arrayFieldRex = (RexInputRef) rexVisitor.analyze(arrayField, context);
+    String alias = expand.getAlias();
 
     // 3. Capture the outer row in a CorrelationId
     Holder<RexCorrelVariable> correlVariable = Holder.empty();
@@ -860,7 +861,7 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
 
     // 5. Filter rows where the array field is the same as the left side
     // TODO: This is not a standard way to use correlate and uncollect together.
-    // A filter should not be necessary. Correct it in the future.
+    //  A filter should not be necessary. Correct it in the future.
     RexNode filterCondition = relBuilder.equals(correlArrayField, arrayFieldRex);
     relBuilder.filter(filterCondition);
 
@@ -876,9 +877,17 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
     // be used by the right side to correlate with the left side.
     relBuilder.correlate(JoinRelType.INNER, correlVariable.get().id, List.of(arrayFieldRex));
 
-    // 8. Remove the original array field from the output. No alias is currently supported in the
-    // expand command, so it can be safely deleted. Its name is re-used for the expanded element.
+    // 8. Remove the original array field from the output.
+    // TODO: RFC: should we keep the original array field when alias is present?
     relBuilder.projectExcept(arrayFieldRex);
+    if (alias != null) {
+      // Sub-nested fields cannot be removed after renaming the nested field.
+      tryToRemoveNestedFields(context);
+      RexInputRef expandedField = relBuilder.field(arrayField.getField().toString());
+      List<String> names = new ArrayList<>(relBuilder.peek().getRowType().getFieldNames());
+      names.set(expandedField.getIndex(), alias);
+      relBuilder.rename(names);
+    }
     return relBuilder.peek();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExpandCommandIT.java

@@ -53,18 +53,15 @@ public void testExpandOnArray() throws Exception {
     verifyNumOfRows(response, 5);
   }
 
-  // TODO: confirm if expand with alias will be supported
-  @Ignore
   @Test
   public void testExpandWithAlias() throws Exception {
     JSONObject response =
         executeQuery(String.format("source=%s | expand address as addr", TEST_INDEX_NESTED_SIMPLE));
     verifySchema(
         response,
         schema("name", "string"),
-        schema("age", "integer"),
-        schema("id", "integer"),
-        schema("address", "array"),
+        schema("age", "bigint"),
+        schema("id", "bigint"),
         schema("addr", "struct"));
     verifyNumOfRows(response, 11);
   }

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -240,7 +240,7 @@ trendlineType
    ;
 
 expandCommand
-    : EXPAND fieldExpression
+    : EXPAND fieldExpression (AS alias = qualifiedName)?
     ;
 
 kmeansCommand

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java

@@ -419,7 +419,8 @@ public UnresolvedPlan visitTopCommand(OpenSearchPPLParser.TopCommandContext ctx)
   @Override
   public UnresolvedPlan visitExpandCommand(OpenSearchPPLParser.ExpandCommandContext ctx) {
     Field fieldExpression = (Field) internalVisitExpression(ctx.fieldExpression());
-    return new Expand(fieldExpression);
+    String alias = ctx.alias != null ? internalVisitExpression(ctx.alias).toString() : null;
+    return new Expand(fieldExpression, alias);
   }
 
   @Override