Skip to content

Commit ce3163b

Browse files
LantaoJinLantaoJin
committed
Support composite aggregation paginating in HAVING clause
Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 885230f commit ce3163b

22 files changed

Lines changed: 469 additions & 72 deletions

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1271,6 +1271,30 @@ public void testExplainCountsByAgg() throws IOException {
12711271
TEST_INDEX_ACCOUNT)));
12721272
}
12731273

1274+
@Test
1275+
public void testHaving() throws IOException {
1276+
enabledOnlyWhenPushdownIsEnabled();
1277+
String expected = loadExpectedPlan("explain_agg_having1.yaml");
1278+
assertYamlEqualsIgnoreId(
1279+
expected,
1280+
explainQueryYaml(
1281+
"source=opensearch-sql_test_index_account | stats count() as c by"
1282+
+ " state | where c > 10"));
1283+
expected = loadExpectedPlan("explain_agg_having2.yaml");
1284+
assertYamlEqualsIgnoreId(
1285+
expected,
1286+
explainQueryYaml(
1287+
"source=opensearch-sql_test_index_account | stats bucket_nullable = false count() by"
1288+
+ " state | where `count()` > 10"));
1289+
expected = loadExpectedPlan("explain_agg_having3.yaml");
1290+
assertYamlEqualsIgnoreId(
1291+
expected,
1292+
explainQueryYaml(
1293+
"source=opensearch-sql_test_index_account | stats avg(balance) as avg, count() as cnt"
1294+
+ " by state | eval new_avg = avg + 1000, new_cnt = cnt + 1 | where new_avg > 1000"
1295+
+ " or new_cnt > 1"));
1296+
}
1297+
12741298
@Test
12751299
public void testExplainSortOnMeasure() throws IOException {
12761300
enabledOnlyWhenPushdownIsEnabled();

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

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

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

8+
import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
9+
import static org.opensearch.sql.util.MatcherUtils.rows;
10+
import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
11+
12+
import java.io.IOException;
13+
import org.json.JSONObject;
14+
import org.junit.jupiter.api.Test;
815
import org.opensearch.sql.ppl.StatsCommandIT;
916

1017
public class CalciteStatsCommandIT extends StatsCommandIT {
@@ -14,4 +21,20 @@ public void init() throws Exception {
1421
enableCalcite();
1522
setQuerySizeLimit(2000);
1623
}
24+
25+
@Test
26+
public void testStatsHaving() throws IOException {
27+
try {
28+
setQueryBucketSize(2);
29+
JSONObject response =
30+
executeQuery(
31+
String.format(
32+
"source=%s | stats sum(balance) as a by state | where a > 780000",
33+
TEST_INDEX_ACCOUNT));
34+
System.out.println(response);
35+
verifyDataRows(response, rows(782199, "TX"));
36+
} finally {
37+
resetQueryBucketSize();
38+
}
39+
}
1740
}

‎integ-test/src/test/resources/expectedOutput/calcite/explain_agg_having1.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalFilter(condition=[>($0, 10)])
5+
LogicalProject(c=[$1], state=[$0])
6+
LogicalAggregate(group=[{0}], c=[COUNT()])
7+
LogicalProject(state=[$7])
8+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
9+
physical: |
10+
EnumerableLimit(fetch=[10000])
11+
EnumerableCalc(expr#0..1=[{inputs}], expr#2=[10], expr#3=[>($t0, $t2)], proj#0..1=[{exprs}], $condition=[$t3])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},c=COUNT()), HAVING->>($0, 10), PROJECT->[c, state]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0, paginatingAgg=true)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_agg_having2.yaml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalFilter(condition=[>($0, 10)])
5+
LogicalProject(count()=[$1], state=[$0])
6+
LogicalAggregate(group=[{0}], count()=[COUNT()])
7+
LogicalProject(state=[$7])
8+
LogicalFilter(condition=[IS NOT NULL($7)])
9+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
10+
physical: |
11+
EnumerableLimit(fetch=[10000])
12+
EnumerableCalc(expr#0..1=[{inputs}], expr#2=[10], expr#3=[>($t0, $t2)], proj#0..1=[{exprs}], $condition=[$t3])
13+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), HAVING->>($0, 10), PROJECT->[count(), state]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0, paginatingAgg=true)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_agg_having3.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalFilter(condition=[OR(>($3, 1000), >($4, 1))])
5+
LogicalProject(avg=[$1], cnt=[$2], state=[$0], new_avg=[+($1, 1000)], new_cnt=[+($2, 1)])
6+
LogicalAggregate(group=[{0}], avg=[AVG($1)], cnt=[COUNT()])
7+
LogicalProject(state=[$7], balance=[$3])
8+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
9+
physical: |
10+
EnumerableLimit(fetch=[10000])
11+
EnumerableCalc(expr#0..3=[{inputs}], expr#4=[CAST($t1):DOUBLE], expr#5=[/($t4, $t2)], expr#6=[1000], expr#7=[+($t5, $t6)], expr#8=[1], expr#9=[+($t3, $t8)], expr#10=[>($t7, $t6)], expr#11=[>($t9, $t8)], expr#12=[OR($t10, $t11)], avg=[$t5], cnt=[$t3], state=[$t0], new_avg=[$t7], new_cnt=[$t9], $condition=[$t12])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},agg#0=SUM($1),agg#1=COUNT($1),cnt=COUNT()), HAVING->OR(>(+($1, 1000), 1000), >(+($2, 1), 1))], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"state":{"terms":{"field":"state.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"$f1":{"sum":{"field":"balance"}},"$f2":{"value_count":{"field":"balance"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0, paginatingAgg=true)])

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/HavingPushdownRule.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.planner.rules;
7+
8+
import java.util.function.Predicate;
9+
import org.apache.calcite.plan.RelOptRuleCall;
10+
import org.apache.calcite.plan.RelRule;
11+
import org.apache.calcite.rel.logical.LogicalFilter;
12+
import org.apache.calcite.rel.logical.LogicalProject;
13+
import org.immutables.value.Value;
14+
import org.opensearch.sql.calcite.plan.OpenSearchRuleConfig;
15+
import org.opensearch.sql.opensearch.storage.scan.AbstractCalciteIndexScan;
16+
import org.opensearch.sql.opensearch.storage.scan.CalciteLogicalIndexScan;
17+
18+
@Value.Enclosing
19+
public class HavingPushdownRule extends RelRule<HavingPushdownRule.Config> {
20+
protected HavingPushdownRule(HavingPushdownRule.Config config) {
21+
super(config);
22+
}
23+
24+
@Override
25+
public void onMatch(RelOptRuleCall call) {
26+
final LogicalFilter filter = call.rel(0);
27+
final LogicalProject project = call.rel(1);
28+
final CalciteLogicalIndexScan scan = call.rel(2);
29+
CalciteLogicalIndexScan newScan = scan.pushDownHavingFlag(filter);
30+
if (newScan != null) {
31+
call.transformTo(
32+
call.builder()
33+
.push(newScan)
34+
.project(project.getProjects())
35+
.filter(filter.getCondition())
36+
.build());
37+
}
38+
}
39+
40+
@Value.Immutable
41+
public interface Config extends OpenSearchRuleConfig {
42+
HavingPushdownRule.Config DEFAULT =
43+
ImmutableHavingPushdownRule.Config.builder()
44+
.build()
45+
.withDescription("Filter-Project-TableScan(AggPushed)")
46+
.withOperandSupplier(
47+
b0 ->
48+
b0.operand(LogicalFilter.class)
49+
.oneInput(
50+
b1 ->
51+
b1.operand(LogicalProject.class)
52+
.oneInput(
53+
b2 ->
54+
b2.operand(CalciteLogicalIndexScan.class)
55+
.predicate(
56+
Predicate.not(
57+
AbstractCalciteIndexScan
58+
::noAggregatePushed)
59+
.and(
60+
Predicate.not(
61+
AbstractCalciteIndexScan
62+
::isHavingFlagPushed)))
63+
.noInputs())));
64+
65+
@Override
66+
default HavingPushdownRule toRule() {
67+
return new HavingPushdownRule(this);
68+
}
69+
}
70+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/OpenSearchIndexRules.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ public class OpenSearchIndexRules {
4141
RareTopPushdownRule.Config.DEFAULT.toRule();
4242
private static final SortExprIndexScanRule SORT_EXPR_INDEX_SCAN =
4343
SortExprIndexScanRule.Config.DEFAULT.toRule();
44+
private static final HavingPushdownRule HAVING_PUSHDOWN =
45+
HavingPushdownRule.Config.DEFAULT.toRule();
4446

4547
// Rule that always pushes down relevance functions regardless of pushdown settings
4648
public static final RelevanceFunctionPushdownRule RELEVANCE_FUNCTION_PUSHDOWN =
@@ -63,7 +65,8 @@ public class OpenSearchIndexRules {
6365
SORT_AGGREGATION_METRICS_RULE,
6466
RARE_TOP_PUSH_DOWN,
6567
EXPAND_COLLATION_ON_PROJECT_EXPR,
66-
SORT_EXPR_INDEX_SCAN);
68+
SORT_EXPR_INDEX_SCAN,
69+
HAVING_PUSHDOWN);
6770

6871
// prevent instantiation
6972
private OpenSearchIndexRules() {}

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

Lines changed: 86 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import java.io.IOException;
1313
import java.util.Collections;
1414
import java.util.List;
15+
import java.util.Map;
1516
import java.util.function.Consumer;
1617
import java.util.function.Function;
1718
import lombok.EqualsAndHashCode;
@@ -28,6 +29,9 @@
2829
import org.opensearch.search.SearchHit;
2930
import org.opensearch.search.SearchHits;
3031
import org.opensearch.search.SearchModule;
32+
import org.opensearch.search.aggregations.Aggregation;
33+
import org.opensearch.search.aggregations.bucket.composite.CompositeAggregation;
34+
import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
3135
import org.opensearch.search.builder.PointInTimeBuilder;
3236
import org.opensearch.search.builder.SearchSourceBuilder;
3337
import org.opensearch.search.sort.FieldSortBuilder;
@@ -75,34 +79,42 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
7579

7680
private SearchResponse searchResponse = null;
7781

78-
/** Constructor of OpenSearchQueryRequest. */
79-
public OpenSearchQueryRequest(
80-
String indexName, int size, OpenSearchExprValueFactory factory, List<String> includes) {
81-
this(new IndexName(indexName), size, factory, includes);
82-
}
82+
private boolean paginatingAgg;
8383

84-
/** Constructor of OpenSearchQueryRequest. */
85-
public OpenSearchQueryRequest(
86-
IndexName indexName, int size, OpenSearchExprValueFactory factory, List<String> includes) {
87-
this.indexName = indexName;
88-
this.sourceBuilder = new SearchSourceBuilder();
84+
private Map<String, Object> afterKey;
85+
86+
/** For testing only. */
87+
static OpenSearchQueryRequest of(
88+
String indexName, int size, OpenSearchExprValueFactory factory, List<String> includes) {
89+
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
8990
sourceBuilder.from(0);
9091
sourceBuilder.size(size);
9192
sourceBuilder.timeout(DEFAULT_QUERY_TIMEOUT);
92-
this.exprValueFactory = factory;
93-
this.includes = includes;
93+
return of(indexName, sourceBuilder, factory, includes);
9494
}
9595

96-
/** Constructor of OpenSearchQueryRequest. */
96+
/** For testing only. */
97+
static OpenSearchQueryRequest of(
98+
String indexName,
99+
SearchSourceBuilder sourceBuilder,
100+
OpenSearchExprValueFactory factory,
101+
List<String> includes) {
102+
return new OpenSearchQueryRequest(
103+
new IndexName(indexName), sourceBuilder, factory, includes, false);
104+
}
105+
106+
/** Constructor of OpenSearchQueryRequest without PIT support. */
97107
public OpenSearchQueryRequest(
98108
IndexName indexName,
99109
SearchSourceBuilder sourceBuilder,
100110
OpenSearchExprValueFactory factory,
101-
List<String> includes) {
111+
List<String> includes,
112+
boolean paginatingAgg) {
102113
this.indexName = indexName;
103114
this.sourceBuilder = sourceBuilder;
104115
this.exprValueFactory = factory;
105116
this.includes = includes;
117+
this.paginatingAgg = paginatingAgg;
106118
}
107119

108120
/** Constructor of OpenSearchQueryRequest with PIT support. */
@@ -119,6 +131,7 @@ public OpenSearchQueryRequest(
119131
this.includes = includes;
120132
this.cursorKeepAlive = cursorKeepAlive;
121133
this.pitId = pitId;
134+
this.paginatingAgg = false; // mutex with pitId
122135
}
123136

124137
/** true if the request is a count aggregation request. */
@@ -170,28 +183,67 @@ public OpenSearchQueryRequest(StreamInput in, OpenSearchStorageEngine engine) th
170183
public OpenSearchResponse search(
171184
Function<SearchRequest, SearchResponse> searchAction,
172185
Function<SearchScrollRequest, SearchResponse> scrollAction) {
173-
if (this.pitId == null) {
174-
// When SearchRequest doesn't contain PitId, fetch single page request
175-
if (searchDone) {
176-
return new OpenSearchResponse(
177-
SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
178-
} else {
179-
// get the value before set searchDone = true
180-
boolean isCountAggRequest = isCountAggRequest();
181-
searchDone = true;
182-
return new OpenSearchResponse(
183-
searchAction.apply(
184-
new SearchRequest().indices(indexName.getIndexNames()).source(sourceBuilder)),
185-
exprValueFactory,
186-
includes,
187-
isCountAggRequest);
188-
}
186+
if (paginatingAgg) {
187+
return searchPaginatingAgg(searchAction);
188+
} else if (this.pitId == null) {
189+
return searchSinglePage(searchAction);
189190
} else {
190191
// Search with PIT instead of scroll API
191192
return searchWithPIT(searchAction);
192193
}
193194
}
194195

196+
private OpenSearchResponse searchSinglePage(
197+
Function<SearchRequest, SearchResponse> searchAction) {
198+
// When SearchRequest doesn't contain PitId, fetch single page request
199+
if (searchDone) {
200+
return new OpenSearchResponse(
201+
SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
202+
} else {
203+
// get the value before set searchDone = true
204+
boolean isCountAggRequest = isCountAggRequest();
205+
searchDone = true;
206+
return new OpenSearchResponse(
207+
searchAction.apply(
208+
new SearchRequest().indices(indexName.getIndexNames()).source(sourceBuilder)),
209+
exprValueFactory,
210+
includes,
211+
isCountAggRequest);
212+
}
213+
}
214+
215+
private OpenSearchResponse searchPaginatingAgg(
216+
Function<SearchRequest, SearchResponse> searchAction) {
217+
OpenSearchResponse openSearchResponse;
218+
if (searchDone) {
219+
openSearchResponse =
220+
new OpenSearchResponse(
221+
SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
222+
} else {
223+
SearchRequest searchRequest =
224+
new SearchRequest().indices(indexName.getIndexNames()).source(this.sourceBuilder);
225+
this.searchResponse = searchAction.apply(searchRequest);
226+
227+
openSearchResponse =
228+
new OpenSearchResponse(
229+
this.searchResponse, exprValueFactory, includes, isCountAggRequest());
230+
// get the value before set searchDone = true
231+
boolean isCountAggRequest = isCountAggRequest();
232+
needClean = openSearchResponse.isEmpty();
233+
searchDone = openSearchResponse.isEmpty();
234+
Aggregation agg = this.searchResponse.getAggregations().asList().get(0);
235+
if (agg instanceof CompositeAggregation compositeAgg) {
236+
afterKey = compositeAgg.afterKey();
237+
if (afterKey != null && !afterKey.isEmpty()) {
238+
this.sourceBuilder.aggregations().getAggregatorFactories().stream()
239+
.filter(b -> b instanceof CompositeAggregationBuilder)
240+
.forEach(c -> ((CompositeAggregationBuilder) c).aggregateAfter(afterKey));
241+
}
242+
}
243+
}
244+
return openSearchResponse;
245+
}
246+
195247
public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse> searchAction) {
196248
OpenSearchResponse openSearchResponse;
197249
if (searchDone) {
@@ -252,6 +304,8 @@ public void clean(Consumer<String> cleanAction) {
252304
}
253305
} finally {
254306
this.pitId = null;
307+
this.searchAfter = null;
308+
this.afterKey = null;
255309
}
256310
}
257311

@@ -264,6 +318,8 @@ public void forceClean(Consumer<String> cleanAction) {
264318
}
265319
} finally {
266320
this.pitId = null;
321+
this.searchAfter = null;
322+
this.afterKey = null;
267323
}
268324
}
269325

0 commit comments

Comments
 (0)