Adding test coverage for nested in ORDER BY clause.

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

docs/user/dql/functions.rst

@@ -4466,6 +4466,17 @@ Example with ``field`` and ``path`` parameters in the SELECT and WHERE clause::
     | b                               |
     +---------------------------------+
 
+Example with ``field`` and ``path`` parameters in the SELECT and ORDER BY clause::
+
+    os> SELECT nested(message.info, message) FROM nested ORDER BY nested(message.info, message) DESC;
+    fetched rows / total rows = 2/2
+    +---------------------------------+
+    | nested(message.info, message)   |
+    |---------------------------------|
+    | b                               |
+    | a                               |
+    +---------------------------------+
+
 
 System Functions
 ================

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

@@ -187,6 +187,40 @@ public void nested_function_with_order_by_clause() {
         rows("zz"));
   }
 
+  @Test
+  public void nested_function_with_order_by_clause_desc() {
+    String query =
+        "SELECT nested(message.info) FROM " + TEST_INDEX_NESTED_TYPE
+            + " ORDER BY nested(message.info, message) DESC";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertEquals(6, result.getInt("total"));
+    verifyDataRows(result,
+        rows("zz"),
+        rows("c"),
+        rows("c"),
+        rows("a"),
+        rows("b"),
+        rows("a"));
+  }
+
+  @Test
+  public void nested_function_and_field_with_order_by_clause() {
+    String query =
+        "SELECT nested(message.info), myNum FROM " + TEST_INDEX_NESTED_TYPE
+            + " ORDER BY nested(message.info, message), myNum";
+    JSONObject result = executeJdbcRequest(query);
+
+    assertEquals(6, result.getInt("total"));
+    verifyDataRows(result,
+        rows("a", 1),
+        rows("c", 4),
+        rows("a", 4),
+        rows("b", 2),
+        rows("c", 3),
+        rows("zz", new JSONArray(List.of(3, 4))));
+  }
+
   // 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.

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java

@@ -124,8 +124,9 @@ private boolean sortByFieldsOnly(LogicalSort sort) {
     return sort.getSortList().stream()
         .map(sortItem -> sortItem.getRight() instanceof ReferenceExpression
         || (sortItem.getRight() instanceof FunctionExpression
-        && ((FunctionExpression)sortItem.getRight()).getFunctionName().getFunctionName().equalsIgnoreCase(
-            BuiltinFunctionName.NESTED.name())))
+        && ((FunctionExpression)sortItem.getRight())
+            .getFunctionName().getFunctionName()
+            .equalsIgnoreCase(BuiltinFunctionName.NESTED.name())))
         .reduce(true, Boolean::logicalAnd);
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/sort/SortQueryBuilder.java

@@ -6,6 +6,8 @@
 
 package org.opensearch.sql.opensearch.storage.script.sort;
 
+import static org.opensearch.sql.analysis.NestedAnalyzer.generatePath;
+
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 import org.opensearch.search.sort.FieldSortBuilder;
@@ -16,13 +18,10 @@
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.FunctionExpression;
-import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.opensearch.data.type.OpenSearchTextType;
 
-import static org.opensearch.sql.analysis.NestedAnalyzer.generatePath;
-
 /**
  * Builder of {@link SortBuilder}.
  */
@@ -62,11 +61,13 @@ public SortBuilder<?> build(Expression expression, Sort.SortOption option) {
     } else if (expression instanceof FunctionExpression
         && ((FunctionExpression)expression).getFunctionName().getFunctionName().equalsIgnoreCase(
         BuiltinFunctionName.NESTED.name())) {
+
       validateNestedArgs((FunctionExpression) expression);
       String orderByName = ((FunctionExpression)expression).getArguments().get(0).toString();
-      ReferenceExpression path = ((FunctionExpression)expression).getArguments().size() == 2 ?
-          (ReferenceExpression) ((FunctionExpression)expression).getArguments().get(1) :
-          generatePath(orderByName);
+      // Generate path if argument not supplied in function.
+      ReferenceExpression path = ((FunctionExpression)expression).getArguments().size() == 2
+          ? (ReferenceExpression) ((FunctionExpression)expression).getArguments().get(1)
+          : generatePath(orderByName);
       return SortBuilders.fieldSort(orderByName)
               .order(sortOrderMap.get(option.getSortOrder()))
               .setNestedSort(new NestedSortBuilder(path.toString()));

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java

@@ -19,6 +19,7 @@
 import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
 import static org.opensearch.sql.data.type.ExprCoreType.LONG;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.expression.DSL.literal;
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation;
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter;
 import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight;
@@ -58,6 +59,7 @@
 import org.opensearch.search.aggregations.AggregationBuilders;
 import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
 import org.opensearch.search.aggregations.bucket.composite.TermsValuesSourceBuilder;
+import org.opensearch.search.sort.NestedSortBuilder;
 import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.search.sort.SortBuilders;
 import org.opensearch.search.sort.SortOrder;
@@ -574,6 +576,76 @@ void only_one_project_should_be_push() {
     );
   }
 
+  @Test
+  void test_nested_sort_filter_push_down() {
+    assertEqualsAfterOptimization(
+        project(
+            indexScanBuilder(
+                withFilterPushedDown(QueryBuilders.termQuery("intV", 1)),
+                withSortPushedDown(
+                    SortBuilders.fieldSort("message.info")
+                        .order(SortOrder.ASC)
+                        .setNestedSort(new NestedSortBuilder("message")))),
+            DSL.named("intV", DSL.ref("intV", INTEGER))
+        ),
+        project(
+                sort(
+                    filter(
+                        relation("schema", table),
+                        DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1)))
+                    ),
+                    Pair.of(
+                        SortOption.DEFAULT_ASC, DSL.nested(DSL.ref("message.info", STRING))
+                    )
+                ),
+            DSL.named("intV", DSL.ref("intV", INTEGER))
+        )
+    );
+  }
+
+  @Test
+  void test_function_expression_sort_returns_optimized_logical_sort() {
+    // Invalid use case coverage OpenSearchIndexScanBuilder::sortByFieldsOnly returns false
+    assertEqualsAfterOptimization(
+        sort(
+            indexScanBuilder(),
+            Pair.of(
+                SortOption.DEFAULT_ASC,
+                DSL.match(DSL.namedArgument("field", literal("message")))
+            )
+        ),
+        sort(
+            relation("schema", table),
+            Pair.of(
+                SortOption.DEFAULT_ASC,
+                DSL.match(DSL.namedArgument("field", literal("message"))
+                )
+            )
+        )
+    );
+  }
+
+  @Test
+  void test_non_field_sort_returns_optimized_logical_sort() {
+    // Invalid use case coverage OpenSearchIndexScanBuilder::sortByFieldsOnly returns false
+    assertEqualsAfterOptimization(
+        sort(
+            indexScanBuilder(),
+            Pair.of(
+                SortOption.DEFAULT_ASC,
+                DSL.literal("field")
+            )
+        ),
+        sort(
+            relation("schema", table),
+            Pair.of(
+                SortOption.DEFAULT_ASC,
+                DSL.literal("field")
+            )
+        )
+    );
+  }
+
   @Test
   void sort_with_expression_cannot_merge_with_relation() {
     assertEqualsAfterOptimization(

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/sort/SortQueryBuilderTest.java

@@ -10,11 +10,14 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 
 import org.hamcrest.Matchers;
 import org.junit.jupiter.api.Test;
 import org.opensearch.sql.ast.tree.Sort;
+import org.opensearch.sql.data.model.ExprShortValue;
 import org.opensearch.sql.expression.DSL;
+import org.opensearch.sql.expression.LiteralExpression;
 
 class SortQueryBuilderTest {
 
@@ -25,6 +28,62 @@ void build_sortbuilder_from_reference() {
     assertNotNull(sortQueryBuilder.build(DSL.ref("intV", INTEGER), Sort.SortOption.DEFAULT_ASC));
   }
 
+  @Test
+  void build_sortbuilder_from_nested_function() {
+    assertNotNull(
+        sortQueryBuilder.build(
+            DSL.nested(DSL.ref("message.info", STRING)),
+            Sort.SortOption.DEFAULT_ASC
+        )
+    );
+  }
+
+  @Test
+  void build_sortbuilder_from_nested_function_with_path_param() {
+    assertNotNull(
+        sortQueryBuilder.build(
+            DSL.nested(DSL.ref("message.info", STRING), DSL.ref("message", STRING)),
+            Sort.SortOption.DEFAULT_ASC
+        )
+    );
+  }
+
+  @Test
+  void nested_with_too_many_args_throws_exception() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> sortQueryBuilder.build(
+            DSL.nested(
+                DSL.ref("message.info", STRING),
+                DSL.ref("message", STRING),
+                DSL.ref("message", STRING)
+            ),
+            Sort.SortOption.DEFAULT_ASC
+        )
+    );
+  }
+
+  @Test
+  void nested_with_invalid_arg_type_throws_exception() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> sortQueryBuilder.build(
+            DSL.nested(
+                DSL.literal(1)
+            ),
+            Sort.SortOption.DEFAULT_ASC
+        )
+    );
+  }
+
+  @Test
+  void build_sortbuilder_from_expression_should_throw_exception() {
+    final IllegalStateException exception =
+        assertThrows(IllegalStateException.class, () -> sortQueryBuilder.build(
+            new LiteralExpression(new ExprShortValue(1)), Sort.SortOption.DEFAULT_ASC));
+    assertThat(exception.getMessage(), Matchers.containsString("unsupported expression"));
+  }
+
   @Test
   void build_sortbuilder_from_function_should_throw_exception() {
     final IllegalStateException exception =

sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java

@@ -646,7 +646,20 @@ public void can_parse_nested_function() {
     assertNotNull(
         parser.parse("SELECT NESTED(PATH.INNER_FIELD, PATH) FROM TEST"));
     assertNotNull(
-        parser.parse("SELECT NESTED(PATH.INNER_FIELD, PATH) FROM TEST ORDER BY nested(PATH.INNER_FIELD, PATH)"));
+        parser.parse(
+            "SELECT * FROM TEST WHERE NESTED(PATH.INNER_FIELDS) = 'A'"
+        )
+    );
+    assertNotNull(
+        parser.parse(
+            "SELECT * FROM TEST WHERE NESTED(PATH.INNER_FIELDS, PATH) = 'A'"
+        )
+    );
+    assertNotNull(
+        parser.parse(
+        "SELECT FIELD FROM TEST ORDER BY nested(PATH.INNER_FIELD, PATH)"
+        )
+    );
   }
 
   @Test