Skip to content

Commit b56d813

Browse files
LantaoJinLantaoJin
committed
Support nested aggregation when calcite enabled
Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 92cb089 commit b56d813

10 files changed

Lines changed: 130 additions & 8 deletions

File tree

‎docs/user/ppl/cmd/stats.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,3 +503,7 @@ PPL query::
503503
| 36 | 30 | M |
504504
+-----+----------+--------+
505505

506+
507+
Limitation
508+
==========
509+
From 3.1.0, the ``stats`` command can preform a `nested aggregation <https://docs.opensearch.org/docs/latest/aggregations/bucket/nested/>`_ only when ``plugins.calcite.enabled`` is true.

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,11 @@
55

66
package org.opensearch.sql.calcite.remote;
77

8+
import static org.opensearch.sql.util.MatcherUtils.assertJsonEqualsIgnoreId;
9+
810
import java.io.IOException;
911
import org.junit.Ignore;
12+
import org.junit.jupiter.api.Test;
1013
import org.opensearch.sql.ppl.ExplainIT;
1114

1215
public class CalciteExplainIT extends ExplainIT {
@@ -43,6 +46,18 @@ public void testTrendlineWithSortPushDownExplain() throws Exception {
4346
"https://github.com/opensearch-project/sql/issues/3466");
4447
}
4548

49+
@Test
50+
public void testNestedAggPushDownExplain() throws Exception {
51+
String expected = loadFromFile("expectedOutput/calcite/explain_nested_agg_push.json");
52+
53+
assertJsonEqualsIgnoreId(
54+
expected,
55+
explainQueryToString(
56+
"source=opensearch-sql_test_index_nested_simple| stats count(address.area) as"
57+
+ " count_area, min(address.area) as min_area, max(address.area) as max_area,"
58+
+ " avg(address.area) as avg_area, avg(age) as avg_age by name"));
59+
}
60+
4661
@Override
4762
@Ignore("test only in v2")
4863
public void testExplainModeUnsupportedInV2() throws IOException {}

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteStatsCommandIT.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,16 @@
55

66
package org.opensearch.sql.calcite.remote;
77

8+
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
9+
import static org.opensearch.sql.util.MatcherUtils.rows;
10+
import static org.opensearch.sql.util.MatcherUtils.schema;
11+
import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
12+
import static org.opensearch.sql.util.MatcherUtils.verifySchemaInOrder;
13+
14+
import java.io.IOException;
15+
import org.json.JSONObject;
16+
import org.junit.Ignore;
17+
import org.junit.jupiter.api.Test;
818
import org.opensearch.sql.ppl.StatsCommandIT;
919

1020
public class CalciteStatsCommandIT extends StatsCommandIT {
@@ -13,5 +23,63 @@ public void init() throws Exception {
1323
super.init();
1424
enableCalcite();
1525
disallowCalciteFallback();
26+
27+
loadIndex(Index.NESTED_SIMPLE);
28+
}
29+
30+
@Test
31+
public void testNestedAggregation() throws IOException {
32+
JSONObject actual =
33+
executeQuery(
34+
String.format(
35+
"source=%s | stats count(address.area) as count_area, min(address.area) as"
36+
+ " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
37+
+ " avg(age) as avg_age",
38+
TEST_INDEX_NESTED_SIMPLE));
39+
verifySchemaInOrder(
40+
actual,
41+
isCalciteEnabled() ? schema("count_area", "bigint") : schema("count_area", "int"),
42+
schema("min_area", "double"),
43+
schema("max_area", "double"),
44+
schema("avg_area", "double"),
45+
schema("avg_age", "double"));
46+
verifyDataRows(actual, rows(9, 9.99, 1000.99, 300.11555555555555, 25.2));
47+
}
48+
49+
@Test
50+
public void testNestedAggregationBy() throws IOException {
51+
JSONObject actual =
52+
executeQuery(
53+
String.format(
54+
"source=%s | stats count(address.area) as count_area, min(address.area) as"
55+
+ " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
56+
+ " avg(age) as avg_age by name",
57+
TEST_INDEX_NESTED_SIMPLE));
58+
verifySchemaInOrder(
59+
actual,
60+
isCalciteEnabled() ? schema("count_area", "bigint") : schema("count_area", "int"),
61+
schema("min_area", "double"),
62+
schema("max_area", "double"),
63+
schema("avg_area", "double"),
64+
schema("avg_age", "double"),
65+
schema("name", "string"));
66+
verifyDataRows(
67+
actual,
68+
rows(4, 10.24, 400.99, 209.69, 24, "abbas"),
69+
rows(0, null, null, null, 19, "andy"),
70+
rows(2, 9.99, 1000.99, 505.49, 32, "chen"),
71+
rows(1, 190.5, 190.5, 190.5, 25, "david"),
72+
rows(2, 231.01, 429.79, 330.4, 26, "peng"));
73+
}
74+
75+
@Ignore("https://github.com/opensearch-project/sql/issues/3384")
76+
public void testNestedAggregationBySpan() throws IOException {
77+
JSONObject actual =
78+
executeQuery(
79+
String.format(
80+
"source=%s | stats count(address.area) as count_area, min(address.area) as"
81+
+ " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
82+
+ " avg(age) as avg_age by span(age, 10)",
83+
TEST_INDEX_NESTED_SIMPLE));
1684
}
1785
}

‎integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public class ExplainIT extends PPLIntegTestCase {
2323
public void init() throws Exception {
2424
super.init();
2525
loadIndex(Index.ACCOUNT);
26+
loadIndex(Index.NESTED_SIMPLE);
2627
}
2728

2829
@Test
@@ -150,7 +151,7 @@ public void testTrendlineWithSortPushDownExplain() throws Exception {
150151
+ "| fields ageTrend"));
151152
}
152153

153-
String loadFromFile(String filename) throws Exception {
154+
protected String loadFromFile(String filename) throws Exception {
154155
URI uri = Resources.getResource(filename).toURI();
155156
return new String(Files.readAllBytes(Paths.get(uri)));
156157
}

‎integ-test/src/test/resources/expectedOutput/calcite/explain_nested_agg_push.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalProject(count_area=[$1], min_area=[$2], max_area=[$3], avg_area=[$4], avg_age=[$5], name=[$0])\n LogicalAggregate(group=[{0}], count_area=[COUNT($1)], min_area=[MIN($1)], max_area=[MAX($1)], avg_area=[AVG($1)], avg_age=[AVG($2)])\n LogicalProject(name=[$0], address.area=[$2], age=[$8])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_nested_simple]])\n",
4+
"physical": "EnumerableCalc(expr#0..5=[{inputs}], count_area=[$t1], min_area=[$t2], max_area=[$t3], avg_area=[$t4], avg_age=[$t5], name=[$t0])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_nested_simple]], PushDownContext=[[PROJECT->[name, address.area, age], AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count_area=COUNT($1),min_area=MIN($1),max_area=MAX($1),avg_area=AVG($1),avg_age=AVG($2))], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"name\",\"address.area\",\"age\"],\"excludes\":[]},\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"name\":{\"terms\":{\"field\":\"name.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"nested_count_area\":{\"nested\":{\"path\":\"address\"},\"aggregations\":{\"count_area\":{\"value_count\":{\"field\":\"address.area\"}}}},\"nested_min_area\":{\"nested\":{\"path\":\"address\"},\"aggregations\":{\"min_area\":{\"min\":{\"field\":\"address.area\"}}}},\"nested_max_area\":{\"nested\":{\"path\":\"address\"},\"aggregations\":{\"max_area\":{\"max\":{\"field\":\"address.area\"}}}},\"nested_avg_area\":{\"nested\":{\"path\":\"address\"},\"aggregations\":{\"avg_area\":{\"avg\":{\"field\":\"address.area\"}}}},\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
5+
}
6+
}

‎integ-test/src/test/resources/indexDefinitions/nested_simple_index_mapping.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
"format": "basic_date_time"
3030
}
3131
}
32+
},
33+
"area" : {
34+
"type": "double"
3235
}
3336
}
3437
},

‎integ-test/src/test/resources/nested_simple.json‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
{"index":{"_id":"1"}}
2-
{"name":"abbas","age":24,"address":[{"city":"New york city","state":"NY","moveInDate":{"dateAndTime":"19840412T090742.000Z"}},{"city":"bellevue","state":"WA","moveInDate":[{"dateAndTime":"20230503T080742.000Z"},{"dateAndTime":"20011111T040744.000Z"}]},{"city":"seattle","state":"WA","moveInDate":{"dateAndTime":"19660319T030455.000Z"}},{"city":"chicago","state":"IL","moveInDate":{"dateAndTime":"20110601T010142.000Z"}}]}
2+
{"name":"abbas","age":24,"address":[{"city":"New york city","state":"NY","moveInDate":{"dateAndTime":"19840412T090742.000Z"},"area":300.13},{"city":"bellevue","state":"WA","moveInDate":[{"dateAndTime":"20230503T080742.000Z"},{"dateAndTime":"20011111T040744.000Z"}],"area":400.99},{"city":"seattle","state":"WA","moveInDate":{"dateAndTime":"19660319T030455.000Z"},"area":127.4},{"city":"chicago","state":"IL","moveInDate":{"dateAndTime":"20110601T010142.000Z"},"area":10.24}]}
33
{"index":{"_id":"2"}}
4-
{"name":"chen","age":32,"address":[{"city":"Miami","state":"Florida","moveInDate":{"dateAndTime":"19010811T040333.000Z"}},{"city":"los angeles","state":"CA","moveInDate":{"dateAndTime":"20230503T080742.000Z"}}]}
4+
{"name":"chen","age":32,"address":[{"city":"Miami","state":"Florida","moveInDate":{"dateAndTime":"19010811T040333.000Z"},"area":1000.99},{"city":"los angeles","state":"CA","moveInDate":{"dateAndTime":"20230503T080742.000Z"},"area":9.99}]}
55
{"index":{"_id":"3"}}
6-
{"name":"peng","age":26,"address":[{"city":"san diego","state":"CA","moveInDate":{"dateAndTime":"20011111T040744.000Z"}},{"city":"austin","state":"TX","moveInDate":{"dateAndTime":"19770713T090441.000Z"}}]}
6+
{"name":"peng","age":26,"address":[{"city":"san diego","state":"CA","moveInDate":{"dateAndTime":"20011111T040744.000Z"},"area":231.01},{"city":"austin","state":"TX","moveInDate":{"dateAndTime":"19770713T090441.000Z"},"area":429.79}]}
77
{"index":{"_id":"4"}}
88
{"name":"andy","age":19,"id":4,"address":[{"city":"houston","state":"TX","moveInDate":{"dateAndTime":"19331212T050545.000Z"}}]}
99
{"index":{"_id":"5"}}
10-
{"name":"david","age":25,"address":[{"city":"raleigh","state":"NC","moveInDate":{"dateAndTime":"19090617T010421.000Z"}},{"city":"charlotte","state":"SC","moveInDate":[{"dateAndTime":"20011111T040744.000Z"}]}]}
10+
{"name":"david","age":25,"address":[{"city":"raleigh","state":"NC","moveInDate":{"dateAndTime":"19090617T010421.000Z"},"area":190.5},{"city":"charlotte","state":"SC","moveInDate":[{"dateAndTime":"20011111T040744.000Z"}]}]}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/AggregateAnalyzer.java‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
package org.opensearch.sql.opensearch.request;
2828

2929
import static java.util.Objects.requireNonNull;
30+
import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
3031
import static org.opensearch.sql.data.type.ExprCoreType.DATE;
3132
import static org.opensearch.sql.data.type.ExprCoreType.TIME;
3233
import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;
@@ -41,6 +42,7 @@
4142
import org.apache.calcite.rel.core.Aggregate;
4243
import org.apache.calcite.rel.core.AggregateCall;
4344
import org.apache.calcite.sql.SqlKind;
45+
import org.apache.commons.lang3.StringUtils;
4446
import org.apache.commons.lang3.tuple.Pair;
4547
import org.opensearch.search.aggregations.AggregationBuilder;
4648
import org.opensearch.search.aggregations.AggregationBuilders;
@@ -118,7 +120,11 @@ public static Pair<List<AggregationBuilder>, OpenSearchAggregationResponseParser
118120
// Process all aggregate calls
119121
Pair<Builder, List<MetricParser>> builderAndParser =
120122
processAggregateCalls(
121-
groupList.size(), aggregate.getAggCallList(), fieldExpressionCreator, outputFields);
123+
groupList.size(),
124+
aggregate.getAggCallList(),
125+
fieldExpressionCreator,
126+
outputFields,
127+
fieldTypes);
122128
Builder metricBuilder = builderAndParser.getLeft();
123129
List<MetricParser> metricParserList = builderAndParser.getRight();
124130

@@ -146,7 +152,8 @@ private static Pair<Builder, List<MetricParser>> processAggregateCalls(
146152
int groupOffset,
147153
List<AggregateCall> aggCalls,
148154
FieldExpressionCreator fieldExpressionCreator,
149-
List<String> outputFields) {
155+
List<String> outputFields,
156+
Map<String, ExprType> fieldTypes) {
150157
assert aggCalls.size() + groupOffset == outputFields.size()
151158
: "groups size and agg calls size should match with output fields";
152159
Builder metricBuilder = new AggregatorFactories.Builder();
@@ -164,7 +171,17 @@ private static Pair<Builder, List<MetricParser>> processAggregateCalls(
164171

165172
Pair<ValuesSourceAggregationBuilder<?>, MetricParser> builderAndParser =
166173
createAggregationBuilderAndParser(aggCall, argStr, aggField);
167-
metricBuilder.addAggregator(builderAndParser.getLeft());
174+
// Nested aggregation (https://docs.opensearch.org/docs/latest/aggregations/bucket/nested/)
175+
// works as expected only when pushdown is triggerred. If aggregates a nested field without
176+
// pushdown, the result could be incorrect. TODO fix it later.
177+
String rootStr = StringUtils.substringBefore(argStr, ".");
178+
if (fieldTypes.get(rootStr) != null && fieldTypes.get(rootStr) == ARRAY) {
179+
metricBuilder.addAggregator(
180+
AggregationBuilders.nested(String.format("nested_%s", aggCall.getName()), rootStr)
181+
.subAggregation(builderAndParser.getLeft()));
182+
} else {
183+
metricBuilder.addAggregator(builderAndParser.getLeft());
184+
}
168185
metricParserList.add(builderAndParser.getRight());
169186
}
170187
return Pair.of(metricBuilder, metricParserList);

‎opensearch/src/main/java/org/opensearch/sql/opensearch/response/agg/MetricParserHelper.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import lombok.RequiredArgsConstructor;
2323
import org.opensearch.search.aggregations.Aggregation;
2424
import org.opensearch.search.aggregations.Aggregations;
25+
import org.opensearch.search.aggregations.bucket.nested.Nested;
2526
import org.opensearch.sql.common.utils.StringUtils;
2627

2728
/** Parse multiple metrics in one bucket. */
@@ -46,6 +47,9 @@ public MetricParserHelper(List<MetricParser> metricParserList) {
4647
public Map<String, Object> parse(Aggregations aggregations) {
4748
Map<String, Object> resultMap = new HashMap<>();
4849
for (Aggregation aggregation : aggregations) {
50+
if (aggregation instanceof Nested) {
51+
aggregation = ((Nested) aggregation).getAggregations().asList().getFirst();
52+
}
4953
if (metricParserMap.containsKey(aggregation.getName())) {
5054
resultMap.putAll(metricParserMap.get(aggregation.getName()).parse(aggregation));
5155
} else {

‎opensearch/src/test/java/org/opensearch/sql/opensearch/response/AggregationResponseUtils.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.opensearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder;
2727
import org.opensearch.search.aggregations.bucket.histogram.ParsedDateHistogram;
2828
import org.opensearch.search.aggregations.bucket.histogram.ParsedHistogram;
29+
import org.opensearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
30+
import org.opensearch.search.aggregations.bucket.nested.ParsedNested;
2931
import org.opensearch.search.aggregations.bucket.terms.DoubleTerms;
3032
import org.opensearch.search.aggregations.bucket.terms.LongTerms;
3133
import org.opensearch.search.aggregations.bucket.terms.ParsedDoubleTerms;
@@ -87,6 +89,8 @@ public class AggregationResponseUtils {
8789
.put(
8890
TopHitsAggregationBuilder.NAME,
8991
(p, c) -> ParsedTopHits.fromXContent(p, (String) c))
92+
.put(
93+
NestedAggregationBuilder.NAME, (p, c) -> ParsedNested.fromXContent(p, (String) c))
9094
.build()
9195
.entrySet()
9296
.stream()

0 commit comments

Comments
 (0)