Support distinct count aggregation (#167) · WonRay/sql@b3dfc49 · GitHub
Skip to content

Commit b3dfc49

Browse files
chloe-zhpenghuo
chloe-zh
and
penghuo
authored
Support distinct count aggregation (opensearch-project#167)
* Support construct AggregationResponseParser during Aggregator build stage (opensearch-project#108) * Support construct AggregationResponseParser during Aggregator build stage * modify the doc Signed-off-by: penghuo <penghuo@gmail.com> * support distinct count aggregation Signed-off-by: chloe-zh <chloezh1102@gmail.com> * fixed tests Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Merge remote-tracking branch 'upstream/develop' into issue/opensearch-project#100 Signed-off-by: chloe-zh <chloezh1102@gmail.com> # Conflicts: # opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/dsl/MetricAggregationBuilder.java * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * updated user doc Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Update: support only count for distinct aggregations Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Update doc; removed distinct start Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Removed unnecessary methods Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Impl stddev and variance function in SQL and PPL (opensearch-project#115) * impl variance frontend and backend * Support construct AggregationResponseParser during Aggregator build stage * add var and varp for PPL Signed-off-by: penghuo <penghuo@gmail.com> * add UT Signed-off-by: penghuo <penghuo@gmail.com> * fix UT Signed-off-by: penghuo <penghuo@gmail.com> * fix doc format Signed-off-by: penghuo <penghuo@gmail.com> * fix doc format Signed-off-by: penghuo <penghuo@gmail.com> * fix the doc Signed-off-by: penghuo <penghuo@gmail.com> * add stddev_samp and stddev_pop Signed-off-by: penghuo <penghuo@gmail.com> * fix UT coverage * address comments Signed-off-by: penghuo <penghuo@gmail.com> * Fix the aggregation filter missing in named aggregators (opensearch-project#123) * Take the condition expression as property to the named aggregator when wrapping the delegated aggregator Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * Added test case where filtered agg is not pushed down Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * modified comparison test Signed-off-by: chloe-zh <chloezh1102@gmail.com> * removed a comparison test and added it to aggregationIT Signed-off-by: chloe-zh <chloezh1102@gmail.com> * added ppl IT test cases; added window function test cases Signed-off-by: chloe-zh <chloezh1102@gmail.com> * moved distinct window function test cases to WindowsIT Signed-off-by: chloe-zh <chloezh1102@gmail.com> * added ut Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * update Signed-off-by: chloe-zh <chloezh1102@gmail.com> * addressed comments Signed-off-by: chloe-zh <chloezh1102@gmail.com> * added test cases to meet the coverage requirement Signed-off-by: chloe-zh <chloezh1102@gmail.com> * added test cases for distinct count map and array types Signed-off-by: chloe-zh <chloezh1102@gmail.com> Co-authored-by: Peng Huo <penghuo@gmail.com>
1 parent bcdd3f5 commit b3dfc49

29 files changed

Lines changed: 428 additions & 32 deletions

File tree

‎core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java‎

Lines changed: 3 additions & 2 deletions

‎core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,16 @@ public static UnresolvedExpression aggregate(
211211

212212
public static UnresolvedExpression filteredAggregate(
213213
String func, UnresolvedExpression field, UnresolvedExpression condition) {
214-
return new AggregateFunction(func, field, condition);
214+
return new AggregateFunction(func, field).condition(condition);
215+
}
216+
217+
public static UnresolvedExpression distinctAggregate(String func, UnresolvedExpression field) {
218+
return new AggregateFunction(func, field, true);
219+
}
220+
221+
public static UnresolvedExpression filteredDistinctCount(
222+
String func, UnresolvedExpression field, UnresolvedExpression condition) {
223+
return new AggregateFunction(func, field, true).condition(condition);
215224
}
216225

217226
public static Function function(String funcName, UnresolvedExpression... funcArgs) {

‎core/src/main/java/org/opensearch/sql/ast/expression/AggregateFunction.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,12 @@
2828

2929
import java.util.Collections;
3030
import java.util.List;
31+
import javax.annotation.Nullable;
3132
import lombok.EqualsAndHashCode;
3233
import lombok.Getter;
3334
import lombok.RequiredArgsConstructor;
35+
import lombok.Setter;
36+
import lombok.experimental.Accessors;
3437
import org.opensearch.sql.ast.AbstractNodeVisitor;
3538
import org.opensearch.sql.common.utils.StringUtils;
3639

@@ -45,7 +48,10 @@ public class AggregateFunction extends UnresolvedExpression {
4548
private final String funcName;
4649
private final UnresolvedExpression field;
4750
private final List<UnresolvedExpression> argList;
51+
@Setter
52+
@Accessors(fluent = true)
4853
private UnresolvedExpression condition;
54+
private Boolean distinct = false;
4955

5056
/**
5157
* Constructor.
@@ -62,14 +68,13 @@ public AggregateFunction(String funcName, UnresolvedExpression field) {
6268
* Constructor.
6369
* @param funcName function name.
6470
* @param field {@link UnresolvedExpression}.
65-
* @param condition condition in aggregation filter.
71+
* @param distinct whether distinct field is specified or not.
6672
*/
67-
public AggregateFunction(String funcName, UnresolvedExpression field,
68-
UnresolvedExpression condition) {
73+
public AggregateFunction(String funcName, UnresolvedExpression field, Boolean distinct) {
6974
this.funcName = funcName;
7075
this.field = field;
7176
this.argList = Collections.emptyList();
72-
this.condition = condition;
77+
this.distinct = distinct;
7378
}
7479

7580
@Override

‎core/src/main/java/org/opensearch/sql/data/model/ExprValueUtils.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,14 @@ public static ExprValue fromObjectValue(Object o, ExprCoreType type) {
157157
}
158158
}
159159

160+
public static Byte getByteValue(ExprValue exprValue) {
161+
return exprValue.byteValue();
162+
}
163+
164+
public static Short getShortValue(ExprValue exprValue) {
165+
return exprValue.shortValue();
166+
}
167+
160168
public static Integer getIntegerValue(ExprValue exprValue) {
161169
return exprValue.integerValue();
162170
}

‎core/src/main/java/org/opensearch/sql/expression/DSL.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,10 @@ public Aggregator count(Expression... expressions) {
500500
return aggregate(BuiltinFunctionName.COUNT, expressions);
501501
}
502502

503+
public Aggregator distinctCount(Expression... expressions) {
504+
return count(expressions).distinct(true);
505+
}
506+
503507
public Aggregator varSamp(Expression... expressions) {
504508
return aggregate(BuiltinFunctionName.VARSAMP, expressions);
505509
}

‎core/src/main/java/org/opensearch/sql/expression/aggregation/Aggregator.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,10 @@ public abstract class Aggregator<S extends AggregationState>
6464
@Getter
6565
@Accessors(fluent = true)
6666
protected Expression condition;
67+
@Setter
68+
@Getter
69+
@Accessors(fluent = true)
70+
protected Boolean distinct = false;
6771

6872
/**
6973
* Create an {@link AggregationState} which will be used for aggregation.

‎core/src/main/java/org/opensearch/sql/expression/aggregation/CountAggregator.java‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@
2828

2929
import static org.opensearch.sql.utils.ExpressionUtils.format;
3030

31+
import java.util.HashSet;
3132
import java.util.List;
3233
import java.util.Locale;
34+
import java.util.Set;
3335
import org.opensearch.sql.data.model.ExprValue;
3436
import org.opensearch.sql.data.model.ExprValueUtils;
3537
import org.opensearch.sql.data.type.ExprCoreType;
@@ -45,33 +47,51 @@ public CountAggregator(List<Expression> arguments, ExprCoreType returnType) {
4547

4648
@Override
4749
public CountAggregator.CountState create() {
48-
return new CountState();
50+
return distinct ? new DistinctCountState() : new CountState();
4951
}
5052

5153
@Override
5254
protected CountState iterate(ExprValue value, CountState state) {
53-
state.count++;
55+
state.count(value);
5456
return state;
5557
}
5658

5759
@Override
5860
public String toString() {
59-
return String.format(Locale.ROOT, "count(%s)", format(getArguments()));
61+
return distinct
62+
? String.format(Locale.ROOT, "count(distinct %s)", format(getArguments()))
63+
: String.format(Locale.ROOT, "count(%s)", format(getArguments()));
6064
}
6165

6266
/**
6367
* Count State.
6468
*/
6569
protected static class CountState implements AggregationState {
66-
private int count;
70+
protected int count;
6771

6872
CountState() {
6973
this.count = 0;
7074
}
7175

76+
public void count(ExprValue value) {
77+
count++;
78+
}
79+
7280
@Override
7381
public ExprValue result() {
7482
return ExprValueUtils.integerValue(count);
7583
}
7684
}
85+
86+
protected static class DistinctCountState extends CountState {
87+
private final Set<ExprValue> distinctValues = new HashSet<>();
88+
89+
@Override
90+
public void count(ExprValue value) {
91+
if (!distinctValues.contains(value)) {
92+
distinctValues.add(value);
93+
count++;
94+
}
95+
}
96+
}
7797
}

‎core/src/main/java/org/opensearch/sql/expression/aggregation/NamedAggregator.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ public class NamedAggregator extends Aggregator<AggregationState> {
5454

5555
/**
5656
* NamedAggregator.
57-
* The aggregator properties {@link #condition} is inherited by named aggregator
58-
* to avoid errors introduced by the property inconsistency.
57+
* The aggregator properties {@link #condition} and {@link #distinct}
58+
* are inherited by named aggregator to avoid errors introduced by the property inconsistency.
5959
*
6060
* @param name name
6161
* @param delegated delegated
@@ -67,6 +67,7 @@ public NamedAggregator(
6767
this.name = name;
6868
this.delegated = delegated;
6969
this.condition = delegated.condition;
70+
this.distinct = delegated.distinct;
7071
}
7172

7273
@Override

‎core/src/test/java/org/opensearch/sql/analysis/ExpressionAnalyzerTest.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,24 @@ public void variance_mapto_varPop() {
300300
);
301301
}
302302

303+
@Test
304+
public void distinct_count() {
305+
assertAnalyzeEqual(
306+
dsl.distinctCount(DSL.ref("integer_value", INTEGER)),
307+
AstDSL.distinctAggregate("count", qualifiedName("integer_value"))
308+
);
309+
}
310+
311+
@Test
312+
public void filtered_distinct_count() {
313+
assertAnalyzeEqual(
314+
dsl.distinctCount(DSL.ref("integer_value", INTEGER))
315+
.condition(dsl.greater(DSL.ref("integer_value", INTEGER), DSL.literal(1))),
316+
AstDSL.filteredDistinctCount("count", qualifiedName("integer_value"), function(
317+
">", qualifiedName("integer_value"), intLiteral(1)))
318+
);
319+
}
320+
303321
protected Expression analyze(UnresolvedExpression unresolvedExpression) {
304322
return expressionAnalyzer.analyze(unresolvedExpression, analysisContext);
305323
}

‎core/src/test/java/org/opensearch/sql/data/model/ExprValueUtilsTest.java‎

Lines changed: 2 additions & 2 deletions

0 commit comments

Comments
 (0)