Fix unix test for flatten command

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

Author: yuancu

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

@@ -207,10 +207,41 @@ private void tryToRemoveNestedFields(CalcitePlanContext context) {
             .map(field -> (RexNode) context.relBuilder.field(field))
             .toList();
     if (!duplicatedNestedFields.isEmpty()) {
-      context.relBuilder.projectExcept(duplicatedNestedFields);
+      // This is a workaround to avoid the bug in Calcite:
+      // In {@link RelBuilder#project_(Iterable, Iterable, Iterable, boolean, Iterable)},
+      // the check `RexUtil.isIdentity(nodeList, inputRowType)` will pass when the input
+      // and the output nodeList refer to the same fields, even if the field name list
+      // is different. As a result, renaming operation will not be applied. This makes
+      // the logical plan for the flatten command incorrect, where the operation is
+      // equivalent to renaming the flattened sub-fields. E.g. emp.name -> name.
+      forceProjectExcept(context.relBuilder, duplicatedNestedFields);
     }
   }
 
+  /**
+   * Project except with force.
+   *
+   * <p>This method is copied from {@link RelBuilder#projectExcept(Iterable)} and modified with the
+   * force flag in project set to true. It is subject to future changes in Calcite.
+   *
+   * @param relBuilder RelBuilder
+   * @param expressions Expressions to exclude from the project
+   */
+  private static void forceProjectExcept(RelBuilder relBuilder, Iterable<RexNode> expressions) {
+    List<RexNode> allExpressions = new ArrayList<>(relBuilder.fields());
+    Set<RexNode> excludeExpressions = new HashSet<>();
+    for (RexNode excludeExp : expressions) {
+      if (!excludeExpressions.add(excludeExp)) {
+        throw new IllegalArgumentException(
+            "Input list contains duplicates. Expression " + excludeExp + " exists multiple times.");
+      }
+      if (!allExpressions.remove(excludeExp)) {
+        throw new IllegalArgumentException("Expression " + excludeExp.toString() + " not found.");
+      }
+    }
+    relBuilder.project(allExpressions, ImmutableList.of(), true);
+  }
+
   /**
    * Try to remove metadata fields in two cases:
    *

docs/user/ppl/cmd/flatten.rst

@@ -10,7 +10,6 @@ flatten
 
 Description
 ===========
-    From 3.1.0
 
 Use ``flatten`` command to flatten a nested struct / object field into separate
 fields in a document.
@@ -23,6 +22,10 @@ Note that ``flatten`` does not work on arrays. Please use ``expand`` command
 to expand an array field into multiple rows instead. If the field is an nested
 array of structs, only the first element of the array will be flattened.
 
+Version
+=======
+Since 3.1.0
+
 Syntax
 ======
 

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLFlattenTest.java

@@ -9,18 +9,22 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import lombok.RequiredArgsConstructor;
+import org.apache.calcite.DataContext;
 import org.apache.calcite.config.CalciteConnectionConfig;
+import org.apache.calcite.linq4j.Enumerable;
+import org.apache.calcite.linq4j.Linq4j;
 import org.apache.calcite.plan.RelTraitDef;
 import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rel.type.RelProtoDataType;
+import org.apache.calcite.schema.ScannableTable;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.SchemaPlus;
 import org.apache.calcite.schema.Statistic;
 import org.apache.calcite.schema.Statistics;
-import org.apache.calcite.schema.Table;
 import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.parser.SqlParser;
@@ -32,9 +36,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
-/**
- * Unit tests for {@code flatten} command in PPL.
- */
+/** Unit tests for {@code flatten} command in PPL. */
 public class CalcitePPLFlattenTest extends CalcitePPLAbstractTest {
   public CalcitePPLFlattenTest() {
     super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL);
@@ -45,7 +47,12 @@ protected Frameworks.ConfigBuilder config(CalciteAssert.SchemaSpec... schemaSpec
     final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
     final SchemaPlus schema = CalciteAssert.addSchema(rootSchema, schemaSpecs);
     // Add an empty table with name DEPT for test purpose
-    schema.add("DEPT", new TableWithStruct());
+    ImmutableList<Object[]> rows =
+        ImmutableList.of(
+            new Object[] {10, ImmutableList.of(7369, "ALLEN"), "SMITH", 7369},
+            new Object[] {20, ImmutableList.of(7499, "ALLEN"), "ALLEN", 7499},
+            new Object[] {30, ImmutableList.of(7521, "WARD"), "WARD", 7521});
+    schema.add("DEPT", new TableWithStruct(rows));
     return Frameworks.newConfigBuilder()
         .parserConfig(SqlParser.Config.DEFAULT)
         .defaultSchema(schema)
@@ -58,20 +65,54 @@ public void testFlatten() {
     String ppl = "source=DEPT | flatten EMP";
     RelNode root = getRelNode(ppl);
     // Regarded as an identity scan. See RelBuilder#L2801
-    String expectedLogical = "LogicalTableScan(table=[[scott, DEPT]])\n";
+    String expectedLogical =
+        "LogicalProject(DEPTNO=[$0], EMP=[$1], EMPNAME=[$2], EMPNO=[$3])\n"
+            + "  LogicalTableScan(table=[[scott, DEPT]])\n";
     verifyLogical(root, expectedLogical);
-    String expectedSparkSql = "SELECT *\nFROM `scott`.`DEPT`";
+    String expectedSparkSql =
+        "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME` `EMPNAME`, `EMP.EMPNO` `EMPNO`\nFROM `scott`.`DEPT`";
     verifyPPLToSparkSQL(root, expectedSparkSql);
+    String expectedResult =
+        "DEPTNO=10; EMP={7369, ALLEN}; EMPNAME=SMITH; EMPNO=7369\n"
+            + "DEPTNO=20; EMP={7499, ALLEN}; EMPNAME=ALLEN; EMPNO=7499\n"
+            + "DEPTNO=30; EMP={7521, WARD}; EMPNAME=WARD; EMPNO=7521\n";
+    verifyResult(root, expectedResult);
   }
 
   @Test
   public void testFlattenWithAliases() {
     String ppl = "source=DEPT | flatten EMP as name, number";
     RelNode root = getRelNode(ppl);
-    String expectedLogical = "LogicalTableScan(table=[[scott, DEPT]])\n";
+    String expectedLogical =
+        "LogicalProject(DEPTNO=[$0], EMP=[$1], name=[$2], number=[$3])\n"
+            + "  LogicalTableScan(table=[[scott, DEPT]])\n";
     verifyLogical(root, expectedLogical);
-    String expectedSparkSql = "SELECT *\nFROM `scott`.`DEPT`";
+    String expectedSparkSql =
+        "SELECT `DEPTNO`, `EMP`, `EMP.EMPNAME` `name`, `EMP.EMPNO` `number`\nFROM `scott`.`DEPT`";
     verifyPPLToSparkSQL(root, expectedSparkSql);
+    String expectedResult =
+        "DEPTNO=10; EMP={7369, ALLEN}; name=SMITH; number=7369\n"
+            + "DEPTNO=20; EMP={7499, ALLEN}; name=ALLEN; number=7499\n"
+            + "DEPTNO=30; EMP={7521, WARD}; name=WARD; number=7521\n";
+    verifyResult(root, expectedResult);
+  }
+
+  /**
+   * This validates that the created table is scannable and the nested fields are removed from the
+   * result.
+   */
+  @Test
+  public void testProject() {
+    String ppl = "source=DEPT";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalProject(DEPTNO=[$0], EMP=[$1])\n  LogicalTableScan(table=[[scott, DEPT]])\n";
+    verifyLogical(root, expectedLogical);
+    String expectedResult =
+        "DEPTNO=10; EMP={7369, ALLEN}\n"
+            + "DEPTNO=20; EMP={7499, ALLEN}\n"
+            + "DEPTNO=30; EMP={7521, WARD}\n";
+    verifyResult(root, expectedResult);
   }
 
   @Test
@@ -80,12 +121,15 @@ public void testFlattenWithMismatchedNumberOfAliasesShouldThrow() {
     Throwable t = Assert.assertThrows(IllegalArgumentException.class, () -> getRelNode(ppl));
     verifyErrorMessageContains(
         t,
-        "The number of aliases has to match the number of flattened fields. Expected 2 (EMP.EMPNO,"
-            + " EMP.EMPNAME), got 1 (name)");
+        "The number of aliases has to match the number of flattened fields. Expected 2"
+            + " (EMP.EMPNAME, EMP.EMPNO), got 1 (name)");
   }
 
   // There is no existing table with arrays. We create one for test purpose.
-  public static class TableWithStruct implements Table {
+  @RequiredArgsConstructor
+  public static class TableWithStruct implements ScannableTable {
+    private final ImmutableList<Object[]> rows;
+
     protected final RelProtoDataType protoRowType =
         factory ->
             factory
@@ -102,10 +146,15 @@ public static class TableWithStruct implements Table {
                 // E.g. struct emp will always hava emp.empno and emp.empname in its
                 // logical projection. We add these two fields to simulate this behavior
                 // in opensearch.
-                .add("EMP.EMPNO", SqlTypeName.INTEGER)
                 .add("EMP.EMPNAME", SqlTypeName.VARCHAR)
+                .add("EMP.EMPNO", SqlTypeName.INTEGER)
                 .build();
 
+    @Override
+    public Enumerable<@Nullable Object[]> scan(DataContext root) {
+      return Linq4j.asEnumerable(rows);
+    }
+
     @Override
     public RelDataType getRowType(RelDataTypeFactory typeFactory) {
       return protoRowType.apply(typeFactory);