Implement a minimal viable version of expand

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

Author: yuancu

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

@@ -50,6 +50,7 @@
 import org.opensearch.sql.ast.tree.CloseCursor;
 import org.opensearch.sql.ast.tree.Dedupe;
 import org.opensearch.sql.ast.tree.Eval;
+import org.opensearch.sql.ast.tree.Expand;
 import org.opensearch.sql.ast.tree.FetchCursor;
 import org.opensearch.sql.ast.tree.FillNull;
 import org.opensearch.sql.ast.tree.Filter;
@@ -623,6 +624,12 @@ public LogicalPlan visitML(ML node, AnalysisContext context) {
     return new LogicalML(child, node.getArguments());
   }
 
+  @Override
+  public LogicalPlan visitExpand(Expand expand, AnalysisContext context) {
+    throw new UnsupportedOperationException(
+        "Expand is supported only when " + CALCITE_ENGINE_ENABLED.getKeyValue() + "=true");
+  }
+
   /** Build {@link LogicalTrendline} for Trendline command. */
   @Override
   public LogicalPlan visitTrendline(Trendline node, AnalysisContext context) {

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

@@ -30,7 +30,7 @@ public Expand attach(UnresolvedPlan child) {
 
   @Override
   public List<UnresolvedPlan> getChild() {
-    return ImmutableList.of(child);
+    return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child);
   }
 
   @Override

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

@@ -840,44 +840,45 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
     // 1. Visit Children
     visitChildren(expand, context);
 
-    var relBuilder = context.relBuilder;
+    RelBuilder relBuilder = context.relBuilder;
 
     // 2. Get the field to expand
     Field arrayField = expand.getField();
-    // 3. Unnest the array field
-    // Analyze the array field to get its RexNode
     RexInputRef arrayFieldRex = (RexInputRef) rexVisitor.analyze(arrayField, context);
-    // Push the original table to the RelBuilder stack
-    // No alias is provided in the expand command, so we remove the original array field,
-    // then replace it with the unnest result.
-    //    relBuilder.projectExcept(arrayFieldRex);
 
-    // Capture the outer row in a CorrelationId
+    // 3. Capture the outer row in a CorrelationId
     Holder<RexCorrelVariable> correlVariable = Holder.empty();
     relBuilder.variable(correlVariable::set);
 
-    // Push a copy of the original table to the RelBuilder stack as right
-    // side of the join.
+    // 4. Push a copy of the original table to the RelBuilder stack as right
+    // side of the correlate (join).
     relBuilder.push(relBuilder.peek());
     RexNode correlArrayField =
         relBuilder.field(
             context.rexBuilder.makeCorrel(relBuilder.peek().getRowType(), correlVariable.get().id),
             arrayFieldRex.getIndex());
 
-    // Filter rows where the array field is the same as the left side
+    // 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.
-    // 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);
+
+    // 6. Project only the array field for the uncollect operation
     relBuilder.project(List.of(correlArrayField), List.of(arrayField.getField().toString()));
 
-    // Alias is not supported in expand yet, we pass in an empty list
+    // 7. Expand the array field using uncollect
     relBuilder.uncollect(List.of(), false);
 
+    // 8. Perform a nested-loop join (correlate) between the original table and the expanded
+    // array field.
     // The last parameter has to refer to the array to be expanded on the left side. It will
     // 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.
+    relBuilder.projectExcept(arrayFieldRex);
     return relBuilder.peek();
   }
 }

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

@@ -7,13 +7,65 @@
 
 package org.opensearch.sql.calcite.remote;
 
-import org.opensearch.sql.ppl.ExpandCommandIT;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ARRAY;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
+import static org.opensearch.sql.util.MatcherUtils.schema;
+import static org.opensearch.sql.util.MatcherUtils.verifyNumOfRows;
+import static org.opensearch.sql.util.MatcherUtils.verifySchema;
 
-public class CalciteExpandCommandIT extends ExpandCommandIT {
+import org.json.JSONObject;
+import org.junit.Ignore;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.ppl.PPLIntegTestCase;
+
+public class CalciteExpandCommandIT extends PPLIntegTestCase {
   @Override
   public void init() throws Exception {
     super.init();
+    loadIndex(Index.NESTED_SIMPLE);
+    loadIndex(Index.ARRAY);
     enableCalcite();
     disallowCalciteFallback();
   }
+
+  @Test
+  public void testExpandOnNested() throws Exception {
+    JSONObject response =
+        executeQuery(String.format("source=%s | expand address", TEST_INDEX_NESTED_SIMPLE));
+    verifySchema(
+        response,
+        schema("name", "string"),
+        schema("age", "bigint"),
+        schema("id", "bigint"),
+        schema("address", "struct"));
+    verifyNumOfRows(response, 11);
+  }
+
+  // TODO: confirm if expand on array (instead of nested) will be supported.
+  //  In Opensearch, a string field can store either a single string or an array of strings.
+  //  This makes it difficult to implement expand on array.
+  @Ignore
+  @Test
+  public void testExpandOnArray() throws Exception {
+    JSONObject response =
+        executeQuery(String.format("source=%s | expand strings", TEST_INDEX_ARRAY));
+    verifySchema(response, schema("numbers", "array"), schema("strings", "string"));
+    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("addr", "struct"));
+    verifyNumOfRows(response, 11);
+  }
 }