Skip to content

Commit 554b55f

Browse files
forestmveyforestmvey
authored
Add Support for Nested Function in Order By Clause (#280)
* Adding order by clause support for nested function. Signed-off-by: forestmvey <[email protected]> * Adding test coverage for nested in ORDER BY clause. Signed-off-by: forestmvey <[email protected]> * Added nested function validation to NestedAnalyzer. Signed-off-by: forestmvey <[email protected]> --------- Signed-off-by: forestmvey <[email protected]>
1 parent 6c3744e commit 554b55f

File tree

12 files changed

+267
-39
lines changed

12 files changed

+267
-39
lines changed

‎core/src/main/java/org/opensearch/sql/analysis/NestedAnalyzer.java

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
import org.opensearch.sql.ast.expression.Function;
1818
import org.opensearch.sql.ast.expression.QualifiedName;
1919
import org.opensearch.sql.ast.expression.UnresolvedExpression;
20+
import org.opensearch.sql.expression.Expression;
21+
import org.opensearch.sql.expression.FunctionExpression;
2022
import org.opensearch.sql.expression.NamedExpression;
2123
import org.opensearch.sql.expression.ReferenceExpression;
2224
import org.opensearch.sql.expression.function.BuiltinFunctionName;
@@ -105,7 +107,18 @@ private void validateArgs(List<UnresolvedExpression> args) {
105107
* @param field : Nested field to generate path of.
106108
* @return : Path of field derived from last level of nesting.
107109
*/
108-
private ReferenceExpression generatePath(String field) {
110+
public static ReferenceExpression generatePath(String field) {
109111
return new ReferenceExpression(field.substring(0, field.lastIndexOf(".")), STRING);
110112
}
113+
114+
/**
115+
* Check if supplied expression is a nested function.
116+
* @param expr Expression checking if is nested function.
117+
* @return True if expression is a nested function.
118+
*/
119+
public static Boolean isNestedFunction(Expression expr) {
120+
return (expr instanceof FunctionExpression
121+
&& ((FunctionExpression) expr).getFunctionName().getFunctionName()
122+
.equalsIgnoreCase(BuiltinFunctionName.NESTED.name()));
123+
}
111124
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,11 @@
99
import static java.util.Collections.emptyList;
1010
import static org.junit.jupiter.api.Assertions.assertAll;
1111
import static org.junit.jupiter.api.Assertions.assertEquals;
12+
import static org.junit.jupiter.api.Assertions.assertFalse;
1213
import static org.junit.jupiter.api.Assertions.assertThrows;
1314
import static org.junit.jupiter.api.Assertions.assertTrue;
1415
import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME;
16+
import static org.opensearch.sql.analysis.NestedAnalyzer.isNestedFunction;
1517
import static org.opensearch.sql.ast.dsl.AstDSL.aggregate;
1618
import static org.opensearch.sql.ast.dsl.AstDSL.alias;
1719
import static org.opensearch.sql.ast.dsl.AstDSL.argument;
@@ -39,6 +41,7 @@
3941
import static org.opensearch.sql.data.type.ExprCoreType.LONG;
4042
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
4143
import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;
44+
import static org.opensearch.sql.expression.DSL.literal;
4245
import static org.opensearch.sql.utils.MLCommonsConstants.ACTION;
4346
import static org.opensearch.sql.utils.MLCommonsConstants.ALGO;
4447
import static org.opensearch.sql.utils.MLCommonsConstants.ASYNC;
@@ -574,6 +577,10 @@ public void project_nested_field_arg() {
574577
function("nested", qualifiedName("message", "info")), null)
575578
)
576579
);
580+
581+
assertTrue(isNestedFunction(DSL.nested(DSL.ref("message.info", STRING))));
582+
assertFalse(isNestedFunction(DSL.literal("fieldA")));
583+
assertFalse(isNestedFunction(DSL.match(DSL.namedArgument("field", literal("message")))));
577584
}
578585

579586
@Test

‎docs/user/dql/functions.rst

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4469,6 +4469,17 @@ Example with ``field`` and ``path`` parameters in the SELECT and WHERE clause::
44694469
| b |
44704470
+---------------------------------+
44714471

4472+
Example with ``field`` and ``path`` parameters in the SELECT and ORDER BY clause::
4473+
4474+
os> SELECT nested(message.info, message) FROM nested ORDER BY nested(message.info, message) DESC;
4475+
fetched rows / total rows = 2/2
4476+
+---------------------------------+
4477+
| nested(message.info, message) |
4478+
|---------------------------------|
4479+
| b |
4480+
| a |
4481+
+---------------------------------+
4482+
44724483

44734484
System Functions
44744485
================

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,40 @@ public void nested_function_with_order_by_clause() {
187187
rows("zz"));
188188
}
189189

190+
@Test
191+
public void nested_function_with_order_by_clause_desc() {
192+
String query =
193+
"SELECT nested(message.info) FROM " + TEST_INDEX_NESTED_TYPE
194+
+ " ORDER BY nested(message.info, message) DESC";
195+
JSONObject result = executeJdbcRequest(query);
196+
197+
assertEquals(6, result.getInt("total"));
198+
verifyDataRows(result,
199+
rows("zz"),
200+
rows("c"),
201+
rows("c"),
202+
rows("a"),
203+
rows("b"),
204+
rows("a"));
205+
}
206+
207+
@Test
208+
public void nested_function_and_field_with_order_by_clause() {
209+
String query =
210+
"SELECT nested(message.info), myNum FROM " + TEST_INDEX_NESTED_TYPE
211+
+ " ORDER BY nested(message.info, message), myNum";
212+
JSONObject result = executeJdbcRequest(query);
213+
214+
assertEquals(6, result.getInt("total"));
215+
verifyDataRows(result,
216+
rows("a", 1),
217+
rows("c", 4),
218+
rows("a", 4),
219+
rows("b", 2),
220+
rows("c", 3),
221+
rows("zz", new JSONArray(List.of(3, 4))));
222+
}
223+
190224
// Nested function in GROUP BY clause is not yet implemented for JDBC format. This test ensures
191225
// that the V2 engine falls back to legacy implementation.
192226
// 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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package org.opensearch.sql.opensearch.storage.scan;
77

8+
import static org.opensearch.sql.analysis.NestedAnalyzer.isNestedFunction;
9+
810
import java.util.function.Function;
911
import lombok.EqualsAndHashCode;
1012
import org.opensearch.sql.expression.ReferenceExpression;
@@ -113,9 +115,15 @@ public boolean pushDownNested(LogicalNested nested) {
113115
return delegate.pushDownNested(nested);
114116
}
115117

118+
/**
119+
* Valid if sorting is only by fields.
120+
* @param sort Logical sort
121+
* @return True if sorting by fields only
122+
*/
116123
private boolean sortByFieldsOnly(LogicalSort sort) {
117124
return sort.getSortList().stream()
118-
.map(sortItem -> sortItem.getRight() instanceof ReferenceExpression)
125+
.map(sortItem -> sortItem.getRight() instanceof ReferenceExpression
126+
|| isNestedFunction(sortItem.getRight()))
119127
.reduce(true, Boolean::logicalAnd);
120128
}
121129
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/LuceneQuery.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
package org.opensearch.sql.opensearch.storage.script.filter.lucene;
88

9+
import static org.opensearch.sql.analysis.NestedAnalyzer.isNestedFunction;
10+
911
import com.google.common.collect.ImmutableMap;
1012
import java.util.Map;
1113
import java.util.function.Function;
@@ -62,10 +64,7 @@ public boolean canSupport(FunctionExpression func) {
6264
* @return return true if function has supported nested function expression.
6365
*/
6466
public boolean isNestedPredicate(FunctionExpression func) {
65-
return ((func.getArguments().get(0) instanceof FunctionExpression
66-
&& ((FunctionExpression)func.getArguments().get(0))
67-
.getFunctionName().getFunctionName().equalsIgnoreCase(BuiltinFunctionName.NESTED.name()))
68-
);
67+
return isNestedFunction(func.getArguments().get(0));
6968
}
7069

7170
/**

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,21 @@
66

77
package org.opensearch.sql.opensearch.storage.script.sort;
88

9+
import static org.opensearch.sql.analysis.NestedAnalyzer.generatePath;
10+
import static org.opensearch.sql.analysis.NestedAnalyzer.isNestedFunction;
11+
912
import com.google.common.collect.ImmutableMap;
1013
import java.util.Map;
1114
import org.opensearch.search.sort.FieldSortBuilder;
15+
import org.opensearch.search.sort.NestedSortBuilder;
1216
import org.opensearch.search.sort.SortBuilder;
1317
import org.opensearch.search.sort.SortBuilders;
1418
import org.opensearch.search.sort.SortOrder;
1519
import org.opensearch.sql.ast.tree.Sort;
1620
import org.opensearch.sql.expression.Expression;
21+
import org.opensearch.sql.expression.FunctionExpression;
1722
import org.opensearch.sql.expression.ReferenceExpression;
23+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
1824
import org.opensearch.sql.opensearch.data.type.OpenSearchTextType;
1925

2026
/**
@@ -53,11 +59,44 @@ public SortBuilder<?> build(Expression expression, Sort.SortOption option) {
5359
return SortBuilders.scoreSort().order(sortOrderMap.get(option.getSortOrder()));
5460
}
5561
return fieldBuild((ReferenceExpression) expression, option);
62+
} else if (isNestedFunction(expression)) {
63+
64+
validateNestedArgs((FunctionExpression) expression);
65+
String orderByName = ((FunctionExpression)expression).getArguments().get(0).toString();
66+
// Generate path if argument not supplied in function.
67+
ReferenceExpression path = ((FunctionExpression)expression).getArguments().size() == 2
68+
? (ReferenceExpression) ((FunctionExpression)expression).getArguments().get(1)
69+
: generatePath(orderByName);
70+
return SortBuilders.fieldSort(orderByName)
71+
.order(sortOrderMap.get(option.getSortOrder()))
72+
.setNestedSort(new NestedSortBuilder(path.toString()));
5673
} else {
5774
throw new IllegalStateException("unsupported expression " + expression.getClass());
5875
}
5976
}
6077

78+
/**
79+
* Validate semantics for arguments in nested function.
80+
* @param nestedFunc Nested function expression.
81+
*/
82+
private void validateNestedArgs(FunctionExpression nestedFunc) {
83+
if (nestedFunc.getArguments().size() > 2) {
84+
throw new IllegalArgumentException(
85+
"nested function supports 2 parameters (field, path) or 1 parameter (field)"
86+
);
87+
}
88+
89+
for (Expression arg : nestedFunc.getArguments()) {
90+
if (!(arg instanceof ReferenceExpression)) {
91+
throw new IllegalArgumentException(
92+
String.format("Illegal nested field name: %s",
93+
arg.toString()
94+
)
95+
);
96+
}
97+
}
98+
}
99+
61100
private FieldSortBuilder fieldBuild(ReferenceExpression ref, Sort.SortOption option) {
62101
return SortBuilders.fieldSort(
63102
OpenSearchTextType.convertTextToKeyword(ref.getAttr(), ref.type()))

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
2020
import static org.opensearch.sql.data.type.ExprCoreType.LONG;
2121
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
22+
import static org.opensearch.sql.expression.DSL.literal;
2223
import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation;
2324
import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter;
2425
import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight;
@@ -58,6 +59,7 @@
5859
import org.opensearch.search.aggregations.AggregationBuilders;
5960
import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
6061
import org.opensearch.search.aggregations.bucket.composite.TermsValuesSourceBuilder;
62+
import org.opensearch.search.sort.NestedSortBuilder;
6163
import org.opensearch.search.sort.SortBuilder;
6264
import org.opensearch.search.sort.SortBuilders;
6365
import org.opensearch.search.sort.SortOrder;
@@ -574,6 +576,76 @@ void only_one_project_should_be_push() {
574576
);
575577
}
576578

579+
@Test
580+
void test_nested_sort_filter_push_down() {
581+
assertEqualsAfterOptimization(
582+
project(
583+
indexScanBuilder(
584+
withFilterPushedDown(QueryBuilders.termQuery("intV", 1)),
585+
withSortPushedDown(
586+
SortBuilders.fieldSort("message.info")
587+
.order(SortOrder.ASC)
588+
.setNestedSort(new NestedSortBuilder("message")))),
589+
DSL.named("intV", DSL.ref("intV", INTEGER))
590+
),
591+
project(
592+
sort(
593+
filter(
594+
relation("schema", table),
595+
DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1)))
596+
),
597+
Pair.of(
598+
SortOption.DEFAULT_ASC, DSL.nested(DSL.ref("message.info", STRING))
599+
)
600+
),
601+
DSL.named("intV", DSL.ref("intV", INTEGER))
602+
)
603+
);
604+
}
605+
606+
@Test
607+
void test_function_expression_sort_returns_optimized_logical_sort() {
608+
// Invalid use case coverage OpenSearchIndexScanBuilder::sortByFieldsOnly returns false
609+
assertEqualsAfterOptimization(
610+
sort(
611+
indexScanBuilder(),
612+
Pair.of(
613+
SortOption.DEFAULT_ASC,
614+
DSL.match(DSL.namedArgument("field", literal("message")))
615+
)
616+
),
617+
sort(
618+
relation("schema", table),
619+
Pair.of(
620+
SortOption.DEFAULT_ASC,
621+
DSL.match(DSL.namedArgument("field", literal("message"))
622+
)
623+
)
624+
)
625+
);
626+
}
627+
628+
@Test
629+
void test_non_field_sort_returns_optimized_logical_sort() {
630+
// Invalid use case coverage OpenSearchIndexScanBuilder::sortByFieldsOnly returns false
631+
assertEqualsAfterOptimization(
632+
sort(
633+
indexScanBuilder(),
634+
Pair.of(
635+
SortOption.DEFAULT_ASC,
636+
DSL.literal("field")
637+
)
638+
),
639+
sort(
640+
relation("schema", table),
641+
Pair.of(
642+
SortOption.DEFAULT_ASC,
643+
DSL.literal("field")
644+
)
645+
)
646+
);
647+
}
648+
577649
@Test
578650
void sort_with_expression_cannot_merge_with_relation() {
579651
assertEqualsAfterOptimization(

0 commit comments

Comments
 (0)