Fixing bug where Nested functions used in WHERE, GROUP BY, HAVING, an…

…d ORDER BY clauses don't fallback to legacy engine. (#1549)

Signed-off-by: forestmvey <[email protected]>

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

@@ -60,18 +60,21 @@
 import org.opensearch.sql.ast.tree.TableFunction;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
+import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.data.model.ExprMissingValue;
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.NamedExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.aggregation.Aggregator;
 import org.opensearch.sql.expression.aggregation.NamedAggregator;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.FunctionName;
 import org.opensearch.sql.expression.function.TableFunctionImplementation;
@@ -217,13 +220,36 @@ public LogicalPlan visitLimit(Limit node, AnalysisContext context) {
   public LogicalPlan visitFilter(Filter node, AnalysisContext context) {
     LogicalPlan child = node.getChild().get(0).accept(this, context);
     Expression condition = expressionAnalyzer.analyze(node.getCondition(), context);
+    verifySupportsCondition(condition);
 
     ExpressionReferenceOptimizer optimizer =
         new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child);
     Expression optimized = optimizer.optimize(condition, context);
     return new LogicalFilter(child, optimized);
   }
 
+  /**
+   * Ensure NESTED function is not used in WHERE, GROUP BY, and HAVING clauses.
+   * Fallback to legacy engine. Can remove when support is added for NESTED function in WHERE,
+   * GROUP BY, ORDER BY, and HAVING clauses.
+   * @param condition : Filter condition
+   */
+  private void verifySupportsCondition(Expression condition) {
+    if (condition instanceof FunctionExpression) {
+      if (((FunctionExpression) condition).getFunctionName().getFunctionName().equalsIgnoreCase(
+          BuiltinFunctionName.NESTED.name()
+      )) {
+        throw new SyntaxCheckException(
+            "Falling back to legacy engine. Nested function is not supported in WHERE,"
+                + " GROUP BY, and HAVING clauses."
+        );
+      }
+      ((FunctionExpression)condition).getArguments().stream()
+          .forEach(e -> verifySupportsCondition(e)
+      );
+    }
+  }
+
   /**
    * Build {@link LogicalRename}.
    */
@@ -273,7 +299,9 @@ public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
     }
 
     for (UnresolvedExpression expr : node.getGroupExprList()) {
-      groupbyBuilder.add(namedExpressionAnalyzer.analyze(expr, context));
+      NamedExpression resolvedExpr = namedExpressionAnalyzer.analyze(expr, context);
+      verifySupportsCondition(resolvedExpr.getDelegated());
+      groupbyBuilder.add(resolvedExpr);
     }
     ImmutableList<NamedExpression> groupBys = groupbyBuilder.build();
 

core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java

@@ -79,6 +79,7 @@
 import org.opensearch.sql.ast.tree.ML;
 import org.opensearch.sql.ast.tree.RareTopN.CommandType;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
@@ -1016,6 +1017,53 @@ public void select_all_from_subquery() {
     );
   }
 
+  /**
+   * Ensure Nested function falls back to legacy engine when used in GROUP BY clause.
+   * TODO Remove this test when support is added.
+   */
+  @Test
+  public void nested_group_by_clause_throws_syntax_exception() {
+    SyntaxCheckException exception = assertThrows(SyntaxCheckException.class,
+        () -> analyze(
+            AstDSL.project(
+                AstDSL.agg(
+                    AstDSL.relation("schema"),
+                    emptyList(),
+                    emptyList(),
+                    ImmutableList.of(alias("nested(message.info)",
+                        function("nested",
+                            qualifiedName("message", "info")))),
+                    emptyList()
+                )))
+    );
+    assertEquals("Falling back to legacy engine. Nested function is not supported in WHERE,"
+            + " GROUP BY, and HAVING clauses.",
+        exception.getMessage());
+  }
+
+  /**
+   * Ensure Nested function falls back to legacy engine when used in WHERE clause.
+   * TODO Remove this test when support is added.
+   */
+  @Test
+  public void nested_where_clause_throws_syntax_exception() {
+    SyntaxCheckException exception = assertThrows(SyntaxCheckException.class,
+        () -> analyze(
+            AstDSL.filter(
+                AstDSL.relation("schema"),
+                AstDSL.equalTo(
+                    AstDSL.function("nested", qualifiedName("message", "info")),
+                    AstDSL.stringLiteral("str")
+                )
+            )
+        )
+    );
+    assertEquals("Falling back to legacy engine. Nested function is not supported in WHERE,"
+            + " GROUP BY, and HAVING clauses.",
+        exception.getMessage());
+  }
+
+
   /**
    * SELECT name, AVG(age) FROM test GROUP BY name.
    */

integ-test/src/test/java/org/opensearch/sql/sql/NestedIT.java

@@ -17,7 +17,6 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
-
 import org.json.JSONArray;
 import org.json.JSONObject;
 import org.junit.Test;
@@ -171,6 +170,63 @@ public void nested_function_mixed_with_non_nested_type_test() {
         rows("zz", "a"));
   }
 
+  @Test
+  public void nested_function_with_where_clause() {
+    String query =
+        "SELECT nested(message.info) FROM " + TEST_INDEX_NESTED_TYPE + " WHERE nested(message.info) = 'a'";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertEquals(2, result.getInt("total"));
+    verifyDataRows(result,
+        rows("a"),
+        rows("a"));
+  }
+
+  @Test
+  public void nested_function_with_order_by_clause() {
+    String query =
+        "SELECT nested(message.info) FROM " + TEST_INDEX_NESTED_TYPE
+            + " ORDER BY nested(message.info)";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertEquals(6, result.getInt("total"));
+    verifyDataRows(result,
+        rows("a"),
+        rows("c"),
+        rows("a"),
+        rows("b"),
+        rows("c"),
+        rows("zz"));
+  }
+
+  // Nested function in GROUP BY clause is not yet implemented for JDBC format. This test ensures
+  // that the V2 engine falls back to legacy implementation.
+  // TODO Fix the test when NESTED is supported in GROUP BY in the V2 engine.
+  @Test
+  public void nested_function_with_group_by_clause() {
+    String query =
+        "SELECT count(*) FROM " + TEST_INDEX_NESTED_TYPE + " GROUP BY nested(message.info)";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertTrue(result.getJSONObject("error").get("details").toString().contains(
+        "Aggregation type nested is not yet implemented"
+    ));
+  }
+
+  // Nested function in HAVING clause is not yet implemented for JDBC format. This test ensures
+  // that the V2 engine falls back to legacy implementation.
+  // TODO Fix the test when NESTED is supported in HAVING in the V2 engine.
+  @Test
+  public void nested_function_with_having_clause() {
+    String query =
+        "SELECT count(*) FROM " + TEST_INDEX_NESTED_TYPE + " GROUP BY myNum HAVING nested(comment.likes) > 7";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertTrue(result.getJSONObject("error").get("details").toString().contains(
+        "For more details, please send request for Json format to see the raw response from OpenSearch engine."
+    ));
+  }
+
   @Test
   public void nested_function_mixed_with_non_nested_types_test() {
     String query =

sql/src/main/java/org/opensearch/sql/sql/parser/AstSortBuilder.java

@@ -17,12 +17,17 @@
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.ast.expression.Argument;
 import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.Function;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.NullOrder;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.Sort.SortOrder;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParserBaseVisitor;
 import org.opensearch.sql.sql.parser.context.QuerySpecification;
 
@@ -48,6 +53,15 @@ private List<Field> createSortFields() {
     List<UnresolvedExpression> items = querySpec.getOrderByItems();
     List<SortOption> options = querySpec.getOrderByOptions();
     for (int i = 0; i < items.size(); i++) {
+      // TODO remove me when Nested function is supported in ORDER BY clause.
+      if (items.get(i) instanceof Function
+          && ((Function)items.get(i)).getFuncName().equalsIgnoreCase(
+              BuiltinFunctionName.NESTED.name())
+      ) {
+        throw new SyntaxCheckException(
+            "Falling back to legacy engine. Nested function is not supported in ORDER BY clause."
+        );
+      }
       fields.add(
           new Field(
               querySpec.replaceIfAliasOrOrdinal(items.get(i)),

sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java

@@ -462,6 +462,21 @@ public void can_build_order_by_null_option() {
         buildAST("SELECT name FROM test ORDER BY name NULLS LAST"));
   }
 
+  /**
+   * Ensure Nested function falls back to legacy engine when used in an ORDER BY clause.
+   * TODO Remove this test when support is added.
+   */
+  @Test
+  public void nested_in_order_by_clause_throws_exception() {
+    SyntaxCheckException exception = assertThrows(SyntaxCheckException.class,
+        () -> buildAST("SELECT * FROM test ORDER BY nested(message.info)")
+    );
+
+    assertEquals(
+        "Falling back to legacy engine. Nested function is not supported in ORDER BY clause.",
+        exception.getMessage());
+  }
+
   @Test
   public void can_build_order_by_sort_order_keyword_insensitive() {
     assertEquals(