Skip to content

Commit

Permalink
increase test coverage
Browse files Browse the repository at this point in the history
Signed-off-by: Lantao Jin <[email protected]>
  • Loading branch information
LantaoJin committed May 28, 2024
1 parent b0fbd8d commit 9d57c20
Show file tree
Hide file tree
Showing 5 changed files with 265 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,8 @@ public Pair<AggregationBuilder, MetricParser> 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));
Expand Down Expand Up @@ -229,11 +230,15 @@ private Pair<AggregationBuilder, MetricParser> 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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Map<String, Object>> parse(OpenSearchAggregationResponseParser parser, String json) {
return parser.parse(fromJson(json));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 9d57c20

Please sign in to comment.