Skip to content

Commit 7e35a10

Browse files
LantaoJinLantaoJin
committed
fix IT
Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 5487bc4 commit 7e35a10

10 files changed

Lines changed: 84 additions & 23 deletions

File tree

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_dedup_keepempty_false_push.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(account_number=[$0], gender=[$1], age=[$2])
5+
LogicalFilter(condition=[<=($3, 1)])
6+
LogicalProject(account_number=[$0], gender=[$1], age=[$2], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $1 ORDER BY $1)])
7+
LogicalFilter(condition=[IS NOT NULL($1)])
8+
LogicalProject(account_number=[$0], gender=[$4], age=[$8])
9+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
10+
physical: |
11+
EnumerableLimit(fetch=[10000])
12+
EnumerableCalc(expr#0..17=[{inputs}], expr#18=[1], expr#19=[<=($t17, $t18)], account_number=[$t0], gender=[$t4], age=[$t8], $condition=[$t19])
13+
EnumerableWindow(window#0=[window(partition {4} order by [4] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
14+
EnumerableCalc(expr#0..16=[{inputs}], expr#17=[IS NOT NULL($t4)], proj#0..16=[{exprs}], $condition=[$t17])
15+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_dedup_keepempty_true_not_pushed.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+
LogicalProject(account_number=[$0], gender=[$1], age=[$2])
5+
LogicalFilter(condition=[OR(IS NULL($1), <=($3, 1))])
6+
LogicalProject(account_number=[$0], gender=[$4], age=[$8], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $4 ORDER BY $4)])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
8+
physical: |
9+
EnumerableLimit(fetch=[10000])
10+
EnumerableCalc(expr#0..17=[{inputs}], expr#18=[IS NULL($t4)], expr#19=[1], expr#20=[<=($t17, $t19)], expr#21=[OR($t18, $t20)], account_number=[$t0], gender=[$t4], age=[$t8], $condition=[$t21])
11+
EnumerableWindow(window#0=[window(partition {4} order by [4] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_dedup_push.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(account_number=[$0], gender=[$1], age=[$2])
5+
LogicalFilter(condition=[<=($3, 1)])
6+
LogicalProject(account_number=[$0], gender=[$1], age=[$2], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $1 ORDER BY $1)])
7+
LogicalFilter(condition=[IS NOT NULL($1)])
8+
LogicalProject(account_number=[$0], gender=[$4], age=[$8])
9+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
10+
physical: |
11+
EnumerableLimit(fetch=[10000])
12+
EnumerableCalc(expr#0..17=[{inputs}], expr#18=[1], expr#19=[<=($t17, $t18)], account_number=[$t0], gender=[$t4], age=[$t8], $condition=[$t19])
13+
EnumerableWindow(window#0=[window(partition {4} order by [4] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
14+
EnumerableCalc(expr#0..16=[{inputs}], expr#17=[IS NOT NULL($t4)], proj#0..16=[{exprs}], $condition=[$t17])
15+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])

‎opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,7 @@ private ExprValue parseContent(Content content) {
299299
* value. For example, {"empty_field": []}.
300300
*/
301301
private Optional<ExprType> type(String field) {
302-
return Optional.ofNullable(typeMapping.get(field));
302+
return Optional.ofNullable(typeMapping.get(field)).map(ExprType::getOriginalExprType);
303303
}
304304

305305
/**

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -426,7 +426,7 @@ private static Pair<AggregationBuilder, MetricParser> createRegularAggregation(
426426
.sort(
427427
helper.inferNamedField(args.getFirst()).getReferenceForTermQuery(),
428428
SortOrder.ASC),
429-
new TopHitsParser(aggFieldName, true));
429+
new TopHitsParser(aggFieldName, true, false));
430430
}
431431
}
432432
case MAX -> {
@@ -445,7 +445,7 @@ private static Pair<AggregationBuilder, MetricParser> createRegularAggregation(
445445
.sort(
446446
helper.inferNamedField(args.getFirst()).getReferenceForTermQuery(),
447447
SortOrder.DESC),
448-
new TopHitsParser(aggFieldName, true));
448+
new TopHitsParser(aggFieldName, true, false));
449449
}
450450
}
451451
case VAR_SAMP ->
@@ -495,15 +495,15 @@ yield switch (functionName) {
495495
helper.inferNamedField(args.getFirst()).getReferenceForTermQuery())
496496
.size(helper.inferValue(args.getLast(), Integer.class))
497497
.from(0),
498-
new TopHitsParser(aggFieldName, false));
498+
new TopHitsParser(aggFieldName, false, true));
499499
case FIRST -> {
500500
TopHitsAggregationBuilder firstBuilder =
501501
AggregationBuilders.topHits(aggFieldName).size(1).from(0);
502502
if (!args.isEmpty()) {
503503
firstBuilder.fetchField(
504504
helper.inferNamedField(args.getFirst()).getReferenceForTermQuery());
505505
}
506-
yield Pair.of(firstBuilder, new TopHitsParser(aggFieldName, true));
506+
yield Pair.of(firstBuilder, new TopHitsParser(aggFieldName, true, false));
507507
}
508508
case LAST -> {
509509
TopHitsAggregationBuilder lastBuilder =
@@ -515,7 +515,7 @@ yield switch (functionName) {
515515
lastBuilder.fetchField(
516516
helper.inferNamedField(args.getFirst()).getReferenceForTermQuery());
517517
}
518-
yield Pair.of(lastBuilder, new TopHitsParser(aggFieldName, true));
518+
yield Pair.of(lastBuilder, new TopHitsParser(aggFieldName, true, false));
519519
}
520520
case PERCENTILE_APPROX -> {
521521
PercentilesAggregationBuilder aggBuilder =
@@ -564,7 +564,7 @@ yield switch (functionName) {
564564
aggCall.getAggregation(), rex.getKind()));
565565
}
566566
});
567-
yield Pair.of(topHitsAggregationBuilder, new TopHitsParser(aggFieldName, false));
567+
yield Pair.of(topHitsAggregationBuilder, new TopHitsParser(aggFieldName, false, false));
568568
}
569569
default ->
570570
throw new AggregateAnalyzer.AggregateAnalyzerException(

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

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
package org.opensearch.sql.opensearch.response.agg;
77

8-
import java.util.Arrays;
98
import java.util.Collections;
109
import java.util.HashMap;
1110
import java.util.List;
@@ -32,15 +31,9 @@ public List<Map<String, Object>> parse(Aggregation agg) {
3231
}
3332

3433
// Get value from fields (fetchField)
35-
List<Map<String, Object>> res =
36-
Arrays.stream(hits)
37-
.filter(hit -> hit.getFields() != null && hit.getFields().isEmpty())
38-
.map(hit -> hit.getFields().values().iterator().next().getValue())
39-
.map(v -> new HashMap<>(Collections.singletonMap(agg.getName(), v)))
40-
.map(v -> (Map<String, Object>) v)
41-
.toList();
42-
if (!res.isEmpty()) {
43-
return res;
34+
if (hits[0].getFields() != null && !hits[0].getFields().isEmpty()) {
35+
Object value = hits[0].getFields().values().iterator().next().getValue();
36+
return Collections.singletonList(Collections.singletonMap(agg.getName(), value));
4437
} else {
4538
return Collections.singletonList(
4639
new HashMap<>(Collections.singletonMap(agg.getName(), null)));

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

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@
1111
import java.util.LinkedHashMap;
1212
import java.util.List;
1313
import java.util.Map;
14+
import java.util.Objects;
1415
import java.util.stream.Collectors;
1516
import lombok.EqualsAndHashCode;
1617
import lombok.Getter;
18+
import org.opensearch.common.document.DocumentField;
1719
import org.opensearch.search.SearchHit;
1820
import org.opensearch.search.aggregations.Aggregation;
1921
import org.opensearch.search.aggregations.metrics.TopHits;
@@ -24,10 +26,12 @@ public class TopHitsParser implements MetricParser {
2426

2527
@Getter private final String name;
2628
private final boolean returnSingleValue;
29+
private final boolean returnMergeValue;
2730

28-
public TopHitsParser(String name, boolean returnSingleValue) {
31+
public TopHitsParser(String name, boolean returnSingleValue, boolean returnMergeValue) {
2932
this.name = name;
3033
this.returnSingleValue = returnSingleValue;
34+
this.returnMergeValue = returnMergeValue;
3135
}
3236

3337
@Override
@@ -42,13 +46,33 @@ public List<Map<String, Object>> parse(Aggregation agg) {
4246

4347
if (hits[0].getFields() == null || hits[0].getFields().isEmpty()) {
4448
return Collections.singletonList(
45-
new HashMap<>(Collections.singletonMap(agg.getName(), Collections.emptyList())));
49+
new HashMap<>(
50+
Collections.singletonMap(
51+
agg.getName(), returnSingleValue ? null : Collections.emptyList())));
4652
}
4753
if (returnSingleValue) {
54+
if (hits[0].getFields() == null || hits[0].getFields().isEmpty()) {
55+
return Collections.singletonList(
56+
new HashMap<>(Collections.singletonMap(agg.getName(), null)));
57+
}
4858
// Extract the single value from the first (and only) hit from fields (fetchField)
4959
Object value = hits[0].getFields().values().iterator().next().getValue();
5060
return Collections.singletonList(
5161
new HashMap<>(Collections.singletonMap(agg.getName(), value)));
62+
} else if (returnMergeValue) {
63+
if (hits[0].getFields() == null || hits[0].getFields().isEmpty()) {
64+
return Collections.singletonList(
65+
new HashMap<>(Collections.singletonMap(agg.getName(), Collections.emptyList())));
66+
}
67+
// Return all values as a list from fields (fetchField)
68+
return Collections.singletonList(
69+
Collections.singletonMap(
70+
agg.getName(),
71+
Arrays.stream(hits)
72+
.flatMap(h -> h.getFields().values().stream())
73+
.map(DocumentField::getValue)
74+
.filter(Objects::nonNull) // Filter out null values
75+
.collect(Collectors.toList())));
5276
} else {
5377
// "hits": {
5478
// "hits": [

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ public Pair<AggregationBuilder, MetricParser> visitNamedAggregator(
156156
node.getArguments().get(1),
157157
condition,
158158
name,
159-
new TopHitsParser(name, false));
159+
new TopHitsParser(name, false, true));
160160
case "percentile":
161161
case "percentile_approx":
162162
return make(

‎opensearch/src/test/java/org/opensearch/sql/opensearch/request/AggregateAnalyzerTest.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,8 @@ void analyze_firstAggregation() throws ExpressionNotAnalyzableException {
359359
.withAggCall(b -> b.aggregateCall(PPLBuiltinOperators.FIRST, b.field("a")).as("first_a"))
360360
.expectDslQuery(
361361
"[{\"first_a\":{\"top_hits\":{\"from\":0,\"size\":1,\"version\":false,\"seq_no_primary_term\":false,\"explain\":false}}}]")
362-
.expectResponseParser(new MetricParserHelper(List.of(new TopHitsParser("first_a", true))))
362+
.expectResponseParser(
363+
new MetricParserHelper(List.of(new TopHitsParser("first_a", true, false))))
363364
.verify();
364365
}
365366

@@ -369,7 +370,8 @@ void analyze_lastAggregation() throws ExpressionNotAnalyzableException {
369370
.withAggCall(b -> b.aggregateCall(PPLBuiltinOperators.LAST, b.field("b")).as("last_b"))
370371
.expectDslQuery(
371372
"[{\"last_b\":{\"top_hits\":{\"from\":0,\"size\":1,\"version\":false,\"seq_no_primary_term\":false,\"explain\":false,\"sort\":[{\"_doc\":{\"order\":\"desc\"}}]}}}]")
372-
.expectResponseParser(new MetricParserHelper(List.of(new TopHitsParser("last_b", true))))
373+
.expectResponseParser(
374+
new MetricParserHelper(List.of(new TopHitsParser("last_b", true, false))))
373375
.verify();
374376
}
375377

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ void top_hits_aggregation_should_pass() {
295295
+ " }\n"
296296
+ "}";
297297
OpenSearchAggregationResponseParser parser =
298-
new CompositeAggregationParser(new TopHitsParser("take", false));
298+
new CompositeAggregationParser(new TopHitsParser("take", false, false));
299299
assertThat(
300300
parse(parser, response),
301301
contains(

0 commit comments

Comments
 (0)