Union query result should work in JDBC format

Signed-off-by: Lantao Jin <ltjin@amazon.com>

integ-test/src/test/java/org/opensearch/sql/legacy/CsvFormatResponseIT.java

@@ -706,6 +706,103 @@ public void selectFunctionAsFieldTest() throws IOException {
     Assert.assertEquals(1, headers.size());
   }
 
+  @Test
+  public void unionTest() throws IOException {
+    String query =
+        String.format(
+            Locale.ROOT,
+            "SELECT firstname, lastname FROM %s LIMIT 3 "
+                + "UNION ALL SELECT firstname, lastname FROM %s LIMIT 3",
+            TestsConstants.TEST_INDEX_ACCOUNT,
+            TestsConstants.TEST_INDEX_ACCOUNT);
+    CSVResult result = executeCsvRequest(query, false, false, false);
+    List<String> headers = result.getHeaders();
+    Assert.assertEquals(2, headers.size());
+    Assert.assertTrue(headers.contains("firstname"));
+    Assert.assertTrue(headers.contains("lastname"));
+
+    List<String> lines = result.getLines();
+    Assert.assertEquals(6, lines.size());
+    assertEquals(lines.get(0), "Amber,Duke");
+    assertEquals(lines.get(1), "Hattie,Bond");
+    assertEquals(lines.get(2), "Nanette,Bates");
+    assertEquals(lines.get(0), "Amber,Duke");
+    assertEquals(lines.get(1), "Hattie,Bond");
+    assertEquals(lines.get(2), "Nanette,Bates");
+  }
+
+  @Test
+  public void unionWithAliasLeftTest() throws IOException {
+    String query =
+        String.format(
+            Locale.ROOT,
+            "SELECT lastname AS firstname FROM %s LIMIT 3 "
+                + "UNION ALL SELECT firstname FROM %s LIMIT 3",
+            TestsConstants.TEST_INDEX_ACCOUNT,
+            TestsConstants.TEST_INDEX_ACCOUNT);
+    CSVResult result = executeCsvRequest(query, false, false, false);
+    List<String> headers = result.getHeaders();
+    Assert.assertEquals(1, headers.size());
+    Assert.assertTrue(headers.contains("firstname"));
+
+    List<String> lines = result.getLines();
+    Assert.assertEquals(6, lines.size());
+    assertEquals(lines.get(0), "Duke");
+    assertEquals(lines.get(1), "Bond");
+    assertEquals(lines.get(2), "Bates");
+    assertEquals(lines.get(3), "Amber");
+    assertEquals(lines.get(4), "Hattie");
+    assertEquals(lines.get(5), "Nanette");
+  }
+
+  @Test
+  public void unionWithAliasRightTest() throws IOException {
+    String query =
+        String.format(
+            Locale.ROOT,
+            "SELECT firstname FROM %s LIMIT 3 "
+                + "UNION ALL SELECT lastname AS firstname FROM %s LIMIT 3",
+            TestsConstants.TEST_INDEX_ACCOUNT,
+            TestsConstants.TEST_INDEX_ACCOUNT);
+    CSVResult result = executeCsvRequest(query, false, false, false);
+    List<String> headers = result.getHeaders();
+    Assert.assertEquals(1, headers.size());
+    Assert.assertTrue(headers.contains("firstname"));
+
+    List<String> lines = result.getLines();
+    Assert.assertEquals(6, lines.size());
+    assertEquals(lines.get(0), "Amber");
+    assertEquals(lines.get(1), "Hattie");
+    assertEquals(lines.get(2), "Nanette");
+    assertEquals(lines.get(3), "Duke");
+    assertEquals(lines.get(4), "Bond");
+    assertEquals(lines.get(5), "Bates");
+  }
+
+  @Test
+  public void unionWithAliasBothSideTest() throws IOException {
+    String query =
+        String.format(
+            Locale.ROOT,
+            "SELECT firstname AS name FROM %s LIMIT 3 "
+                + "UNION ALL SELECT lastname AS name FROM %s LIMIT 3",
+            TestsConstants.TEST_INDEX_ACCOUNT,
+            TestsConstants.TEST_INDEX_ACCOUNT);
+    CSVResult result = executeCsvRequest(query, false, false, false);
+    List<String> headers = result.getHeaders();
+    Assert.assertEquals(1, headers.size());
+    Assert.assertTrue(headers.contains("name"));
+
+    List<String> lines = result.getLines();
+    Assert.assertEquals(6, lines.size());
+    assertEquals(lines.get(0), "Amber");
+    assertEquals(lines.get(1), "Hattie");
+    assertEquals(lines.get(2), "Nanette");
+    assertEquals(lines.get(3), "Duke");
+    assertEquals(lines.get(4), "Bond");
+    assertEquals(lines.get(5), "Bates");
+  }
+
   private void verifyFieldOrder(final String[] expectedFields) throws IOException {
 
     final String fields = String.join(", ", expectedFields);

integ-test/src/test/java/org/opensearch/sql/legacy/MultiQueryIT.java

@@ -58,7 +58,31 @@ public void unionAllSameRequestOnlyOneRecordTwice() throws IOException {
   }
 
   @Test
-  public void unionAllOnlyOneRecordEachWithAlias() throws IOException {
+  public void unionAllOnlyOneRecordWithAliasLeft() throws IOException {
+    String query =
+        String.format(
+            "SELECT firstname as dog_name FROM %s WHERE firstname = 'Amber' "
+                + "UNION ALL "
+                + "SELECT dog_name FROM %s WHERE dog_name = 'rex'",
+            TestsConstants.TEST_INDEX_ACCOUNT, TestsConstants.TEST_INDEX_DOG);
+
+    JSONObject response = executeQuery(query);
+    assertThat(getHits(response).length(), equalTo(2));
+
+    Set<String> names = new HashSet<>();
+    JSONArray hits = getHits(response);
+    for (int i = 0; i < hits.length(); i++) {
+      JSONObject hit = hits.getJSONObject(i);
+      JSONObject source = getSource(hit);
+
+      names.add(source.getString("dog_name"));
+    }
+
+    assertThat(names, hasItems("Amber", "rex"));
+  }
+
+  @Test
+  public void unionAllOnlyOneRecordWithAliasRight() throws IOException {
     String query =
         String.format(
             "SELECT firstname FROM %s WHERE firstname = 'Amber' "
@@ -81,6 +105,30 @@ public void unionAllOnlyOneRecordEachWithAlias() throws IOException {
     assertThat(names, hasItems("Amber", "rex"));
   }
 
+  @Test
+  public void unionAllOnlyOneRecordWithAliasBothSide() throws IOException {
+    String query =
+        String.format(
+            "SELECT firstname AS name FROM %s WHERE firstname = 'Amber' "
+                + "UNION ALL "
+                + "SELECT dog_name AS name FROM %s WHERE dog_name = 'rex'",
+            TestsConstants.TEST_INDEX_ACCOUNT, TestsConstants.TEST_INDEX_DOG);
+
+    JSONObject response = executeQuery(query);
+    assertThat(getHits(response).length(), equalTo(2));
+
+    Set<String> names = new HashSet<>();
+    JSONArray hits = getHits(response);
+    for (int i = 0; i < hits.length(); i++) {
+      JSONObject hit = hits.getJSONObject(i);
+      JSONObject source = getSource(hit);
+
+      names.add(source.getString("name"));
+    }
+
+    assertThat(names, hasItems("Amber", "rex"));
+  }
+
   @Test
   public void unionAllOnlyOneRecordEachWithComplexAlias() throws IOException {
     String query =

integ-test/src/test/java/org/opensearch/sql/legacy/PrettyFormatResponseIT.java

@@ -578,6 +578,80 @@ public void fieldOrderOther() throws IOException {
     testFieldOrder(expectedFields, expectedValues);
   }
 
+  @Test
+  public void unionQuery() throws IOException {
+    JSONObject response =
+        executeQuery(
+            String.format(
+                Locale.ROOT,
+                "SELECT firstname, lastname FROM %s "
+                    + "UNION ALL SELECT firstname, lastname FROM %s",
+                TestsConstants.TEST_INDEX_ACCOUNT,
+                TestsConstants.TEST_INDEX_ACCOUNT));
+
+    List<String> fields = Arrays.asList("firstname", "lastname");
+    JSONArray dataRows = getDataRows(response);
+    assertContainsColumns(getSchema(response), fields);
+    assertContainsData(dataRows, fields);
+  }
+
+  @Test
+  public void unionQueryWithAliasLeft() throws IOException {
+    JSONObject response =
+        executeQuery(
+            String.format(
+                Locale.ROOT,
+                "SELECT lastname AS firstname FROM %s UNION ALL SELECT firstname FROM %s",
+                TestsConstants.TEST_INDEX_ACCOUNT,
+                TestsConstants.TEST_INDEX_ACCOUNT));
+    List<String> fields = List.of("lastname");
+    Map<String, String> aliases = new HashMap<>();
+    aliases.put("lastname", "firstname");
+    JSONArray schema = getSchema(response);
+    JSONArray dataRows = getDataRows(response);
+    assertContainsColumns(schema, fields);
+    assertContainsAliases(schema, aliases);
+    assertContainsData(dataRows, fields);
+  }
+
+  @Test
+  public void unionQueryWithAliasRight() throws IOException {
+    JSONObject response =
+        executeQuery(
+            String.format(
+                Locale.ROOT,
+                "SELECT firstname FROM %s UNION ALL SELECT lastname AS firstname FROM %s",
+                TestsConstants.TEST_INDEX_ACCOUNT,
+                TestsConstants.TEST_INDEX_ACCOUNT));
+    List<String> fields = List.of("firstname");
+    JSONArray schema = getSchema(response);
+    JSONArray dataRows = getDataRows(response);
+    assertContainsColumns(schema, fields);
+    // Query schema uses first subquery schema, so alias in second subquery doesn't count in.
+    assertNoAlias(schema);
+    assertContainsData(dataRows, fields);
+  }
+
+  @Test
+  public void unionQueryWithAliasBothSide() throws IOException {
+    JSONObject response =
+        executeQuery(
+            String.format(
+                Locale.ROOT,
+                "SELECT firstname AS name FROM %s UNION ALL SELECT lastname AS name FROM %s",
+                TestsConstants.TEST_INDEX_ACCOUNT,
+                TestsConstants.TEST_INDEX_ACCOUNT));
+    List<String> fields = List.of("firstname");
+    Map<String, String> aliases = new HashMap<>();
+    aliases.put("firstname", "name");
+    aliases.put("lastname", "name");
+    JSONArray schema = getSchema(response);
+    JSONArray dataRows = getDataRows(response);
+    assertContainsColumns(schema, fields);
+    assertContainsAliases(schema, aliases);
+    assertContainsData(dataRows, fields);
+  }
+
   private void testFieldOrder(final String[] expectedFields, final Object[] expectedValues)
       throws IOException {
 
@@ -644,6 +718,13 @@ private void assertContainsAliases(JSONArray schema, Map<String, String> aliases
     }
   }
 
+  private void assertNoAlias(JSONArray schema) {
+    for (int i = 0; i < schema.length(); i++) {
+      JSONObject column = schema.getJSONObject(i);
+      assertFalse(column.has("alias"));
+    }
+  }
+
   private void assertContainsData(JSONArray dataRows, Collection<String> fields) {
     assertThat(dataRows.length(), greaterThan(0));
     JSONArray row = dataRows.getJSONArray(0);

legacy/src/main/java/org/opensearch/sql/legacy/domain/Select.java

@@ -47,6 +47,7 @@ public class Select extends Query {
   private List<SubQueryExpression> subQueries;
   private boolean selectAll = false;
   private JoinType nestedJoinType = JoinType.COMMA;
+  private boolean partOfUnion = false;
 
   public boolean isQuery = false;
   public boolean isAggregate = false;
@@ -187,4 +188,13 @@ public boolean isOrderdSelect() {
   public boolean isSelectAll() {
     return selectAll;
   }
+
+  public void setPartOfUnion(boolean partOfUnion) {
+    this.partOfUnion = partOfUnion;
+  }
+
+  /** Return true is this SELECT is used in UNION */
+  public boolean isPartOfUnion() {
+    return partOfUnion;
+  }
 }

legacy/src/main/java/org/opensearch/sql/legacy/domain/Union.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy.domain;
+
+import lombok.Getter;
+
+@Getter
+public class Union extends Query {
+  private final Select firstTable;
+  private final Select secondTable;
+
+  public Union(Select firstTable, Select secondTable) {
+    this.firstTable = firstTable;
+    this.secondTable = secondTable;
+    this.firstTable.setPartOfUnion(true);
+    this.secondTable.setPartOfUnion(true);
+  }
+}

legacy/src/main/java/org/opensearch/sql/legacy/executor/format/DateFieldFormatter.java

@@ -105,7 +105,7 @@ private List<String> getFormatsForColumn(String columnName) {
   private Set<String> getDateColumns(List<Schema.Column> columns) {
     return columns.stream()
         .filter(column -> column.getType().equals(Schema.Type.DATE.nameLowerCase()))
-        .map(Schema.Column::getName)
+        .map(Schema.Column::getIdentifier)
         .collect(Collectors.toSet());
   }
 

legacy/src/main/java/org/opensearch/sql/legacy/executor/format/Schema.java

@@ -15,7 +15,6 @@
 import org.opensearch.sql.legacy.domain.IndexStatement;
 
 public class Schema implements Iterable<Schema.Column> {
-
   private String indexName;
   private List<Column> columns;
 
@@ -44,7 +43,7 @@ public String getIndexName() {
   }
 
   public List<String> getHeaders() {
-    return columns.stream().map(column -> column.getName()).collect(Collectors.toList());
+    return columns.stream().map(Column::getIdentifier).collect(Collectors.toList());
   }
 
   public List<Column> getColumns() {
@@ -166,5 +165,21 @@ public String getIdentifier() {
     public Type getEnumType() {
       return type;
     }
+
+    @Override
+    public String toString() {
+      return "Column{"
+          + "name='"
+          + name
+          + '\''
+          + ", alias='"
+          + alias
+          + '\''
+          + ", type="
+          + type
+          + ", identifiedByAlias="
+          + identifiedByAlias
+          + '}';
+    }
   }
 }