[analytics-engine] Strip AE-unsupported fields from test datasets at load

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -14,6 +14,7 @@
 import org.apache.calcite.rel.RelRoot;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.runtime.CalciteException;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.util.SqlVisitor;
@@ -25,6 +26,7 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.common.error.ErrorReport;
 import org.opensearch.sql.exception.QueryEngineException;
 import org.opensearch.sql.exception.SemanticCheckException;
 
@@ -74,13 +76,18 @@ public RelNode plan(String query) {
     } catch (SyntaxCheckException
         | QueryEngineException
         | UnsupportedOperationException
-        | IllegalArgumentException e) {
-      LOG.error("Failed to plan query: {}", e.getMessage());
+        | IllegalArgumentException
+        | ErrorReport e) {
+      LOG.warn("Failed to plan query: {}", e.getMessage());
       throw e;
+    } catch (CalciteException e) {
+      // Calcite validation errors (e.g. table not found) indicate an invalid query.
+      LOG.warn("Failed to plan query, invalid query: {}", e.getMessage());
+      throw new SemanticCheckException(e.getMessage(), e);
     } catch (AssertionError e) {
       // Calcite throws assertion error directly when building bad RelNode
       String message = "Failed to plan query: invalid plan structure";
-      LOG.error(message, e);
+      LOG.warn(message, e);
       throw new SemanticCheckException(message, e);
     } catch (Exception e) {
       String message = "Failed to plan query: unexpected error";

api/src/main/java/org/opensearch/sql/api/parser/SqlV2QueryParser.java

@@ -8,8 +8,11 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.existsSubquery;
 import static org.opensearch.sql.ast.dsl.AstDSL.inSubquery;
 import static org.opensearch.sql.ast.dsl.AstDSL.join;
+import static org.opensearch.sql.ast.dsl.AstDSL.union;
 
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
@@ -22,6 +25,7 @@
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.ExistsSubqueryExpressionAtomContext;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InSubqueryPredicateContext;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.JoinClauseContext;
+import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.UnionSelectContext;
 import org.opensearch.sql.sql.parser.AstBuilder;
 import org.opensearch.sql.sql.parser.AstExpressionBuilder;
 import org.opensearch.sql.sql.parser.AstStatementBuilder;
@@ -81,6 +85,13 @@ private JoinType toJoinType(JoinClauseContext ctx) {
       };
     }
 
+    @Override
+    public UnresolvedPlan visitUnionSelect(UnionSelectContext ctx) {
+      List<UnresolvedPlan> datasets =
+          ctx.querySpecification().stream().map(this::visit).collect(Collectors.toList());
+      return union(datasets);
+    }
+
     /**
      * Expression builder with IN/EXISTS subquery support. Accesses the enclosing AstBuilder to
      * visit subquery plan nodes. Must be created via {@link #createExpressionBuilder()} because the

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java

@@ -200,6 +200,42 @@ WHERE age NOT IN (SELECT age FROM catalog.departments WHERE dept_name = 'Enginee
             """);
   }
 
+  @Test
+  public void testUnionAll() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees UNION ALL SELECT dept_name FROM catalog.departments
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[true])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+              LogicalProject(dept_name=[$1])
+                LogicalTableScan(table=[[catalog, departments]])
+            """);
+  }
+
+  @Test
+  public void testMultiWayUnion() {
+    givenQuery(
+            """
+            SELECT name FROM catalog.employees
+            UNION ALL SELECT dept_name FROM catalog.departments
+            UNION ALL SELECT name FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalUnion(all=[true])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+              LogicalProject(dept_name=[$1])
+                LogicalTableScan(table=[[catalog, departments]])
+              LogicalProject(name=[$1])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
   @Test
   public void testNotExistsSubquery() {
     givenQuery(
@@ -241,6 +277,57 @@ public void selectExpressionWithoutFrom() {
             """);
   }
 
+  @Test
+  public void testGroupByAggregateAlias() {
+    givenQuery(
+            """
+            SELECT department, SUM(age) AS total FROM catalog.employees GROUP BY department
+            """)
+        .assertPlan(
+            """
+            LogicalProject(department=[$0], total=[$1])
+              LogicalAggregate(group=[{0}], SUM(age)=[SUM($1)])
+                LogicalProject(department=[$3], age=[$2])
+                  LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testOrderByAggregateAlias() {
+    givenQuery(
+            """
+            SELECT department, COUNT(*) AS cnt FROM catalog.employees
+              GROUP BY department ORDER BY cnt DESC LIMIT 3
+            """)
+        .assertPlan(
+            """
+            LogicalSort(sort0=[$1], dir0=[DESC-nulls-last])
+              LogicalProject(department=[$1], cnt=[$0])
+                LogicalSort(sort0=[$0], dir0=[DESC-nulls-last], fetch=[3])
+                  LogicalProject(COUNT(*)=[$1], department=[$0])
+                    LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
+                      LogicalProject(department=[$3])
+                        LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testAliasPreservedInOutputSchema() {
+    givenQuery("SELECT COUNT(*) AS cnt FROM catalog.employees").assertFields("cnt");
+
+    givenQuery("SELECT department, COUNT(*) AS cnt FROM catalog.employees GROUP BY department")
+        .assertFields("department", "cnt");
+
+    givenQuery("SELECT department, COUNT(*) FROM catalog.employees GROUP BY department")
+        .assertFields("department", "COUNT(*)");
+
+    givenQuery("SELECT MAX(age) + MIN(age) AS range_sum FROM catalog.employees")
+        .assertFields("range_sum");
+
+    givenQuery("SELECT id, name, age AS years, department FROM catalog.employees")
+        .assertFields("id", "name", "years", "department");
+  }
+
   @Test
   public void testHavingMaxCol() {
     givenQuery(
@@ -259,6 +346,47 @@ GROUP BY department HAVING MAX(age) > 30
             """);
   }
 
+  @Test
+  public void testCountStarWithFilter() {
+    givenQuery("SELECT COUNT(*) FILTER(WHERE age > 30) FROM catalog.employees")
+        .assertPlan(
+            """
+            LogicalAggregate(group=[{}], COUNT(*) FILTER(WHERE age > 30)=[COUNT() FILTER $0])
+              LogicalProject($f1=[>($2, 30)])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testFilteredAggregateWithGroupBy() {
+    givenQuery(
+            """
+            SELECT department, SUM(age) FILTER(WHERE age > 30) FROM catalog.employees
+              GROUP BY department
+            """)
+        .assertPlan(
+            """
+            LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[SUM($1) FILTER $2])
+              LogicalProject(department=[$3], age=[$2], $f3=[>($2, 30)])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testMultipleFilteredAggregates() {
+    givenQuery(
+            """
+            SELECT MAX(age) FILTER(WHERE age > 30), MIN(age) FILTER(WHERE age < 50)
+              FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalAggregate(group=[{}], MAX(age) FILTER(WHERE age > 30)=[MAX($0) FILTER $1], MIN(age) FILTER(WHERE age < 50)=[MIN($0) FILTER $2])
+              LogicalProject(age=[$2], $f4=[>($2, 30)], $f5=[<($2, 50)])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
   @Test
   public void testScalarFnOverAggregate() {
     givenQuery("SELECT ABS(MAX(age)) FROM catalog.employees")
@@ -337,6 +465,33 @@ GROUP BY department HAVING MAX(age) > 30 AND MIN(age) < 50
             """);
   }
 
+  @Test
+  public void testCountDistinctWindowWithOrderBy() {
+    // No frame printed: RANGE .. CURRENT ROW is Calcite's default for ORDER BY.
+    givenQuery(
+            """
+            SELECT department, COUNT(DISTINCT name) OVER(ORDER BY department) FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalProject(department=[$3], COUNT(DISTINCT name) OVER(ORDER BY department)=[COUNT(DISTINCT $1) OVER (ORDER BY $3 NULLS FIRST)])
+              LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSumWindowWithPartitionAndOrderBy() {
+    givenQuery(
+            """
+            SELECT name, SUM(age) OVER(PARTITION BY department ORDER BY age) FROM catalog.employees
+            """)
+        .assertPlan(
+            """
+            LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
+              LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
   @Test
   public void testWindowOrderByDefaultsNullsFirst() {
     // Window function ORDER BY without explicit NULLS FIRST/LAST defaults to NULLS FIRST,

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerTest.java

@@ -10,10 +10,12 @@
 
 import java.util.Map;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.runtime.CalciteException;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.impl.AbstractSchema;
 import org.junit.Test;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.common.error.ErrorReport;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.executor.QueryType;
 
@@ -72,7 +74,7 @@ public void testPPLQueryPlanningWithDefaultNamespaceMultiLevel() {
 
     // This is valid in SparkSQL, but Calcite requires "catalog" as the default root schema to
     // resolve it
-    assertThrows(IllegalStateException.class, () -> planner.plan("source = opensearch.employees"));
+    assertThrows(SemanticCheckException.class, () -> planner.plan("source = opensearch.employees"));
   }
 
   @Test
@@ -131,6 +133,20 @@ public void semanticErrorIsRethrownAsSemanticCheckException() {
         .assertErrorMessageEquals("Source and target patterns have different wildcard counts");
   }
 
+  @Test
+  public void fieldNotFoundIsRethrownAsErrorReport() {
+    givenInvalidQuery("source = catalog.employees | where unknown_field = 1")
+        .assertErrorType(ErrorReport.class)
+        .assertErrorMessageContains("Field [unknown_field] not found");
+  }
+
+  @Test
+  public void invalidTableIsRethrownAsSemanticCheckException() {
+    givenInvalidQuery("source = catalog.nonexistent_table")
+        .assertErrorType(SemanticCheckException.class)
+        .assertCauseType(CalciteException.class);
+  }
+
   @Test
   public void assertionErrorIsWrappedAsSemanticCheckException() {
     // Remove when the underlying Calcite assertion is fixed.

build.gradle

@@ -6,7 +6,7 @@
 
 buildscript {
     ext {
-        opensearch_version = System.getProperty("opensearch.version", "3.7.0-SNAPSHOT")
+        opensearch_version = System.getProperty("opensearch.version", "3.8.0-SNAPSHOT")
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
         buildVersionQualifier = System.getProperty("build.version_qualifier", "")
         version_tokens = opensearch_version.tokenize('-')

core/build.gradle

@@ -65,11 +65,11 @@ dependencies {
     }
     api 'org.apache.calcite:calcite-linq4j:1.41.0'
     api project(':common')
-    compileOnly 'org.opensearch.sandbox:analytics-api:3.7.0-SNAPSHOT'
+    compileOnly 'org.opensearch.sandbox:analytics-api:3.8.0-SNAPSHOT'
     // Needed because analytics-api's QueryPlanExecutor signature uses
     // org.opensearch.core.action.ActionListener; AnalyticsExecutionEngine references that type.
     compileOnly group: 'org.opensearch', name: 'opensearch-core', version: "${opensearch_version}"
-    testImplementation 'org.opensearch.sandbox:analytics-api:3.7.0-SNAPSHOT'
+    testImplementation 'org.opensearch.sandbox:analytics-api:3.8.0-SNAPSHOT'
     testImplementation group: 'org.opensearch', name: 'opensearch-core', version: "${opensearch_version}"
     implementation "com.github.seancfoley:ipaddress:5.4.2"
     implementation "com.jayway.jsonpath:json-path:2.9.0"

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

@@ -84,6 +84,7 @@
 import org.opensearch.sql.ast.tree.SubqueryAlias;
 import org.opensearch.sql.ast.tree.TableFunction;
 import org.opensearch.sql.ast.tree.Trendline;
+import org.opensearch.sql.ast.tree.Union;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
 import org.opensearch.sql.calcite.plan.OpenSearchConstants;
@@ -781,4 +782,8 @@ public static InSubquery inSubquery(UnresolvedPlan query, UnresolvedExpression..
   public static ExistsSubquery existsSubquery(UnresolvedPlan query) {
     return new ExistsSubquery(query);
   }
+
+  public static Union union(List<UnresolvedPlan> datasets) {
+    return new Union(datasets);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/expression/WindowBound.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.sql.ast.expression;
 
+import lombok.EqualsAndHashCode;
 import lombok.Getter;
 
 public abstract class WindowBound {
@@ -25,6 +26,7 @@ public boolean isPreceding() {
     }
   }
 
+  @EqualsAndHashCode(callSuper = false)
   public static class CurrentRowWindowBound extends WindowBound {
     CurrentRowWindowBound() {}
 

core/src/main/java/org/opensearch/sql/ast/expression/WindowFrame.java

@@ -40,6 +40,13 @@ public static WindowFrame toCurrentRow() {
         AstDSL.stringLiteral("CURRENT ROW"));
   }
 
+  public static WindowFrame rangeToCurrentRow() {
+    return WindowFrame.of(
+        FrameType.RANGE,
+        AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
+        AstDSL.stringLiteral("CURRENT ROW"));
+  }
+
   public static WindowFrame of(FrameType type, String lower, String upper) {
     return WindowFrame.of(type, AstDSL.stringLiteral(lower), AstDSL.stringLiteral(upper));
   }

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

@@ -21,15 +21,25 @@
 @RequiredArgsConstructor
 @AllArgsConstructor
 public class Union extends UnresolvedPlan {
+  /** Input subplans (operands) combined by this UNION ALL. */
   private final List<UnresolvedPlan> datasets;
 
+  /** Whether inputs are unified to a common schema by name (PPL) vs combined positionally (SQL). */
+  private boolean unifySchema;
+
+  /** Optional cap on output rows (PPL {@code maxout}); {@code null} if unbounded. */
   private Integer maxout;
 
+  /** PPL constructor: UNION ALL with schema unification. */
+  public Union(List<UnresolvedPlan> datasets, Integer maxout) {
+    this(datasets, true, maxout);
+  }
+
   @Override
   public UnresolvedPlan attach(UnresolvedPlan child) {
     List<UnresolvedPlan> newDatasets =
         ImmutableList.<UnresolvedPlan>builder().add(child).addAll(datasets).build();
-    return new Union(newDatasets, maxout);
+    return new Union(newDatasets, unifySchema, maxout);
   }
 
   @Override