WIP: Implementing expand command

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

Author: yuancu

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

@@ -844,12 +844,10 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
 
     // 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
-    RelNode originalTable = relBuilder.peek();
     // 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);
@@ -858,24 +856,27 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
     Holder<RexCorrelVariable> correlVariable = Holder.empty();
     relBuilder.variable(correlVariable::set);
 
-    relBuilder.push(originalTable);
-
+    // Push a copy of the original table to the RelBuilder stack as right
+    // side of the join.
+    relBuilder.push(relBuilder.peek());
     RexNode correlArrayField =
         relBuilder.field(
-            context.rexBuilder.makeCorrel(originalTable.getRowType(), correlVariable.get().id),
+            context.rexBuilder.makeCorrel(relBuilder.peek().getRowType(), correlVariable.get().id),
             arrayFieldRex.getIndex());
 
-    relBuilder.project(
-        List.of(correlArrayField),
-        List.of(arrayField.getField().toString()),
-        false,
-        List.of(correlVariable.get().id));
+    // 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.
+    RexNode filterCondition = relBuilder.equals(correlArrayField, arrayFieldRex);
+    relBuilder.filter(filterCondition);
+    relBuilder.project(List.of(correlArrayField), List.of(arrayField.getField().toString()));
 
     // Alias is not supported in expand yet, we pass in an empty list
     relBuilder.uncollect(List.of(), false);
 
-    //    ImmutableBitSet requiredFields = originalTable.getRowType().getFieldList()
-    relBuilder.correlate(JoinRelType.INNER, correlVariable.get().id, arrayFieldRex);
+    // 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));
 
     return relBuilder.peek();
   }

core/src/main/java/org/opensearch/sql/exception/CalciteUnsupportedException.java

@@ -10,4 +10,8 @@ public class CalciteUnsupportedException extends QueryEngineException {
   public CalciteUnsupportedException(String message) {
     super(message);
   }
+
+  public CalciteUnsupportedException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -113,7 +113,7 @@ public void executeWithCalcite(
       } else {
         if (t instanceof Error) {
           // Calcite may throw AssertError during query execution.
-          listener.onFailure(new CalciteUnsupportedException(t.getMessage()));
+          listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
         } else {
           listener.onFailure((Exception) t);
         }

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

@@ -9,7 +9,7 @@
 
 import org.opensearch.sql.ppl.ExpandCommandIT;
 
-public class CalciteExpandIT extends ExpandCommandIT {
+public class CalciteExpandCommandIT extends ExpandCommandIT {
   @Override
   public void init() throws Exception {
     super.init();

integ-test/src/test/java/org/opensearch/sql/ppl/ExpandCommandIT.java

@@ -7,6 +7,7 @@
 
 package org.opensearch.sql.ppl;
 
+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;
@@ -21,21 +22,31 @@ public class ExpandCommandIT extends PPLIntegTestCase {
   public void init() throws Exception {
     super.init();
     loadIndex(Index.NESTED_SIMPLE);
+    loadIndex(Index.ARRAY);
   }
 
   @Test
-  public void testExpand() throws Exception {
+  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", "integer"),
         schema("id", "integer"),
-        schema("address", "object"));
+        schema("address", "struct"));
     verifyNumOfRows(response, 11);
   }
 
+  @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: double check if expand with alias is supported
   @Ignore
   @Test
@@ -48,7 +59,7 @@ public void testExpandWithAlias() throws Exception {
         schema("age", "integer"),
         schema("id", "integer"),
         schema("address", "array"),
-        schema("addr", "object"));
+        schema("addr", "struct"));
     verifyNumOfRows(response, 11);
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java

@@ -25,12 +25,14 @@
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.runtime.Hook;
 import org.apache.calcite.sql.SqlExplainLevel;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.opensearch.sql.ast.statement.Explain.ExplainFormat;
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners;
 import org.opensearch.sql.common.response.ResponseListener;
 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.executor.ExecutionContext;
 import org.opensearch.sql.executor.ExecutionEngine;
@@ -232,7 +234,18 @@ private void buildResultSet(
     for (int i = 1; i <= columnCount; ++i) {
       String columnName = metaData.getColumnName(i);
       RelDataType fieldType = fieldTypes.get(i - 1);
-      ExprType exprType = convertRelDataTypeToExprType(fieldType);
+      // The element type of struct and array is currently set to ANY.
+      // We set them using the runtime type as a workaround.
+      ExprType exprType;
+      if (fieldType.getSqlTypeName() == SqlTypeName.ANY) {
+        if (!values.isEmpty()) {
+          exprType = values.getFirst().tupleValue().get(columnName).type();
+        } else {
+          exprType = ExprCoreType.UNDEFINED;
+        }
+      } else {
+        exprType = convertRelDataTypeToExprType(fieldType);
+      }
       columns.add(new Column(columnName, null, exprType));
     }
     Schema schema = new Schema(columns);