From 9d57c20bf7d82b892ed767d4804a4b6902ee3343 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Tue, 28 May 2024 18:26:12 +0800 Subject: [PATCH] increase test coverage Signed-off-by: Lantao Jin --- .../dsl/MetricAggregationBuilder.java | 7 +- .../request/OpenSearchRequestBuilderTest.java | 20 +++ .../response/AggregationResponseUtils.java | 5 + ...enSearchAggregationResponseParserTest.java | 144 ++++++++++++++++++ .../dsl/MetricAggregationBuilderTest.java | 90 +++++++++++ 5 files changed, 265 insertions(+), 1 deletion(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java index ee0f458f65..65dca20772 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java @@ -162,7 +162,8 @@ public Pair visitNamedAggregator( return make( AggregationBuilders.percentiles(name), expression, - node.getArguments().get(1), + node.getArguments().get(1), // quantile + node.getArguments().size() >= 3 ? node.getArguments().get(2) : null, // compression condition, name, new SinglePercentileParser(name)); @@ -229,11 +230,15 @@ private Pair make( PercentilesAggregationBuilder builder, Expression expression, Expression quantile, + Expression compression, Expression condition, String name, MetricParser parser) { PercentilesAggregationBuilder aggregationBuilder = helper.build(expression, builder::field, builder::script); + if (compression != null) { + aggregationBuilder.compression(compression.valueOf().doubleValue()); + } aggregationBuilder.percentiles(quantile.valueOf().doubleValue()); if (condition != null) { return Pair.of( diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 5bb0a2207b..742e76cbd0 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -58,6 +58,7 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.CompositeAggregationParser; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; +import org.opensearch.sql.opensearch.response.agg.SinglePercentileParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; import org.opensearch.sql.planner.logical.LogicalNested; @@ -165,6 +166,25 @@ void test_push_down_aggregation() { verify(exprValueFactory).setParser(responseParser); } + @Test + void test_push_down_percentile_aggregation() { + AggregationBuilder aggBuilder = + AggregationBuilders.composite( + "composite_buckets", Collections.singletonList(new TermsValuesSourceBuilder("longA"))); + OpenSearchAggregationResponseParser responseParser = + new CompositeAggregationParser(new SinglePercentileParser("PERCENTILE(intA, 50)")); + requestBuilder.pushDownAggregation(Pair.of(List.of(aggBuilder), responseParser)); + + assertEquals( + new SearchSourceBuilder() + .from(DEFAULT_OFFSET) + .size(0) + .timeout(DEFAULT_QUERY_TIMEOUT) + .aggregation(aggBuilder), + requestBuilder.getSourceBuilder()); + verify(exprValueFactory).setParser(responseParser); + } + @Test void test_push_down_query_and_sort() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/AggregationResponseUtils.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/AggregationResponseUtils.java index 76148b9395..ccdfdce7a4 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/AggregationResponseUtils.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/AggregationResponseUtils.java @@ -41,8 +41,10 @@ import org.opensearch.search.aggregations.metrics.ParsedMax; import org.opensearch.search.aggregations.metrics.ParsedMin; import org.opensearch.search.aggregations.metrics.ParsedSum; +import org.opensearch.search.aggregations.metrics.ParsedTDigestPercentiles; import org.opensearch.search.aggregations.metrics.ParsedTopHits; import org.opensearch.search.aggregations.metrics.ParsedValueCount; +import org.opensearch.search.aggregations.metrics.PercentilesAggregationBuilder; import org.opensearch.search.aggregations.metrics.SumAggregationBuilder; import org.opensearch.search.aggregations.metrics.TopHitsAggregationBuilder; import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder; @@ -56,6 +58,9 @@ public class AggregationResponseUtils { .put(MaxAggregationBuilder.NAME, (p, c) -> ParsedMax.fromXContent(p, (String) c)) .put(SumAggregationBuilder.NAME, (p, c) -> ParsedSum.fromXContent(p, (String) c)) .put(AvgAggregationBuilder.NAME, (p, c) -> ParsedAvg.fromXContent(p, (String) c)) + .put( + PercentilesAggregationBuilder.NAME, + (p, c) -> ParsedTDigestPercentiles.fromXContent(p, (String) c)) .put( ExtendedStatsAggregationBuilder.NAME, (p, c) -> ParsedExtendedStats.fromXContent(p, (String) c)) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchAggregationResponseParserTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchAggregationResponseParserTest.java index 1a15e57c55..938c19fa59 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchAggregationResponseParserTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchAggregationResponseParserTest.java @@ -26,6 +26,8 @@ import org.opensearch.sql.opensearch.response.agg.FilterParser; import org.opensearch.sql.opensearch.response.agg.NoBucketAggregationParser; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; +import org.opensearch.sql.opensearch.response.agg.PercentilesParser; +import org.opensearch.sql.opensearch.response.agg.SinglePercentileParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; import org.opensearch.sql.opensearch.response.agg.StatsParser; import org.opensearch.sql.opensearch.response.agg.TopHitsParser; @@ -309,6 +311,148 @@ void top_hits_aggregation_should_pass() { contains(ImmutableMap.of("type", "take", "take", ImmutableList.of("m", "f")))); } + /** SELECT PERCENTILES(age, 50) FROM accounts. */ + @Test + void no_bucket_one_metric_percentile_should_pass() { + String response = + "{\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"50.0\": 35.0\n" + + " }\n" + + " }\n" + + " }"; + NoBucketAggregationParser parser = + new NoBucketAggregationParser(new SinglePercentileParser("percentiles")); + assertThat(parse(parser, response), contains(entry("percentiles", 35.0))); + } + + /** SELECT PERCENTILES(age) FROM accounts. */ + @Test + void no_bucket_one_metric_percentiles_should_pass() { + String response = + "{\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"1.0\": 21.0,\n" + + " \"5.0\": 27.0,\n" + + " \"25.0\": 30.0,\n" + + " \"50.0\": 35.0,\n" + + " \"75.0\": 55.0,\n" + + " \"95.0\": 58.0,\n" + + " \"99.0\": 60.0\n" + + " }\n" + + " }\n" + + " }"; + NoBucketAggregationParser parser = + new NoBucketAggregationParser(new PercentilesParser("percentiles")); + assertThat( + parse(parser, response), + contains(entry("percentiles", List.of(21.0, 27.0, 30.0, 35.0, 55.0, 58.0, 60.0)))); + } + + @Test + void one_bucket_one_metric_percentile_should_pass() { + String response = + "{\n" + + " \"composite#composite_buckets\": {\n" + + " \"after_key\": {\n" + + " \"type\": \"sale\"\n" + + " },\n" + + " \"buckets\": [\n" + + " {\n" + + " \"key\": {\n" + + " \"type\": \"cost\"\n" + + " },\n" + + " \"doc_count\": 2,\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"50.0\": 40.0\n" + + " }\n" + + " }\n" + + " },\n" + + " {\n" + + " \"key\": {\n" + + " \"type\": \"sale\"\n" + + " },\n" + + " \"doc_count\": 2,\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"50.0\": 100.0\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + " }\n" + + "}"; + + OpenSearchAggregationResponseParser parser = + new CompositeAggregationParser(new SinglePercentileParser("percentiles")); + assertThat( + parse(parser, response), + containsInAnyOrder( + ImmutableMap.of("type", "cost", "percentiles", 40d), + ImmutableMap.of("type", "sale", "percentiles", 100d))); + } + + @Test + void one_bucket_one_metric_percentiles_should_pass() { + String response = + "{\n" + + " \"composite#composite_buckets\": {\n" + + " \"after_key\": {\n" + + " \"type\": \"sale\"\n" + + " },\n" + + " \"buckets\": [\n" + + " {\n" + + " \"key\": {\n" + + " \"type\": \"cost\"\n" + + " },\n" + + " \"doc_count\": 2,\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"1.0\": 21.0,\n" + + " \"5.0\": 27.0,\n" + + " \"25.0\": 30.0,\n" + + " \"50.0\": 35.0,\n" + + " \"75.0\": 55.0,\n" + + " \"95.0\": 58.0,\n" + + " \"99.0\": 60.0\n" + + " }\n" + + " }\n" + + " },\n" + + " {\n" + + " \"key\": {\n" + + " \"type\": \"sale\"\n" + + " },\n" + + " \"doc_count\": 2,\n" + + " \"percentiles#percentiles\": {\n" + + " \"values\": {\n" + + " \"1.0\": 21.0,\n" + + " \"5.0\": 27.0,\n" + + " \"25.0\": 30.0,\n" + + " \"50.0\": 35.0,\n" + + " \"75.0\": 55.0,\n" + + " \"95.0\": 58.0,\n" + + " \"99.0\": 60.0\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + " }\n" + + "}"; + + OpenSearchAggregationResponseParser parser = + new CompositeAggregationParser(new PercentilesParser("percentiles")); + assertThat( + parse(parser, response), + containsInAnyOrder( + ImmutableMap.of( + "type", "cost", "percentiles", List.of(21.0, 27.0, 30.0, 35.0, 55.0, 58.0, 60.0)), + ImmutableMap.of( + "type", "sale", "percentiles", List.of(21.0, 27.0, 30.0, 35.0, 55.0, 58.0, 60.0)))); + } + public List> parse(OpenSearchAggregationResponseParser parser, String json) { return parser.parse(fromJson(json)); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java index 7f302c9c53..6d792dec25 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilderTest.java @@ -10,6 +10,7 @@ import static org.mockito.Mockito.when; import static org.opensearch.sql.common.utils.StringUtils.format; import static org.opensearch.sql.data.type.ExprCoreType.ARRAY; +import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.STRING; import static org.opensearch.sql.expression.DSL.literal; @@ -39,6 +40,7 @@ import org.opensearch.sql.expression.aggregation.MaxAggregator; import org.opensearch.sql.expression.aggregation.MinAggregator; import org.opensearch.sql.expression.aggregation.NamedAggregator; +import org.opensearch.sql.expression.aggregation.PercentileApproximateAggregator; import org.opensearch.sql.expression.aggregation.SumAggregator; import org.opensearch.sql.expression.aggregation.TakeAggregator; import org.opensearch.sql.expression.function.FunctionName; @@ -215,6 +217,94 @@ void should_build_varSamp_aggregation() { varianceSample(Arrays.asList(ref("age", INTEGER)), INTEGER))))); } + @Test + void should_build_percentile_aggregation() { + assertEquals( + format( + "{%n" + + " \"percentile(age, 50)\" : {%n" + + " \"percentiles\" : {%n" + + " \"field\" : \"age\",%n" + + " \"percents\" : [ 50.0 ],%n" + + " \"keyed\" : true,%n" + + " \"tdigest\" : {%n" + + " \"compression\" : 100.0%n" + + " }%n" + + " }%n" + + " }%n" + + "}"), + buildQuery( + Arrays.asList( + named( + "percentile(age, 50)", + new PercentileApproximateAggregator( + Arrays.asList(ref("age", INTEGER), literal(50)), DOUBLE))))); + } + + @Test + void should_build_percentile_with_compression_aggregation() { + assertEquals( + format( + "{%n" + + " \"percentile(age, 50)\" : {%n" + + " \"percentiles\" : {%n" + + " \"field\" : \"age\",%n" + + " \"percents\" : [ 50.0 ],%n" + + " \"keyed\" : true,%n" + + " \"tdigest\" : {%n" + + " \"compression\" : 0.1%n" + + " }%n" + + " }%n" + + " }%n" + + "}"), + buildQuery( + Arrays.asList( + named( + "percentile(age, 50)", + new PercentileApproximateAggregator( + Arrays.asList(ref("age", INTEGER), literal(50), literal(0.1)), DOUBLE))))); + } + + @Test + void should_build_filtered_percentile_aggregation() { + assertEquals( + format( + "{%n" + + " \"percentile(age, 50)\" : {%n" + + " \"filter\" : {%n" + + " \"range\" : {%n" + + " \"age\" : {%n" + + " \"from\" : 30,%n" + + " \"to\" : null,%n" + + " \"include_lower\" : false,%n" + + " \"include_upper\" : true,%n" + + " \"boost\" : 1.0%n" + + " }%n" + + " }%n" + + " },%n" + + " \"aggregations\" : {%n" + + " \"percentile(age, 50)\" : {%n" + + " \"percentiles\" : {%n" + + " \"field\" : \"age\",%n" + + " \"percents\" : [ 50.0 ],%n" + + " \"keyed\" : true,%n" + + " \"tdigest\" : {%n" + + " \"compression\" : 100.0%n" + + " }%n" + + " }%n" + + " }%n" + + " }%n" + + " }%n" + + "}"), + buildQuery( + Arrays.asList( + named( + "percentile(age, 50)", + new PercentileApproximateAggregator( + Arrays.asList(ref("age", INTEGER), literal(50)), DOUBLE) + .condition(DSL.greater(ref("age", INTEGER), literal(30))))))); + } + @Test void should_build_stddevPop_aggregation() { assertEquals(