Skip to content

Commit 055b2ed

Browse files
xinyualxinyual
committed
Add DISTINCT_COUNT_APPROX function (opensearch-project#3654)
* add approx distinct count Signed-off-by: xinyual <xinyual@amazon.com> * add IT Signed-off-by: xinyual <xinyual@amazon.com> * fix IT name Signed-off-by: xinyual <xinyual@amazon.com> * change parser logic Signed-off-by: xinyual <xinyual@amazon.com> * apply spotless Signed-off-by: xinyual <xinyual@amazon.com> * add doc Signed-off-by: xinyual <xinyual@amazon.com> * add java doc Signed-off-by: xinyual <xinyual@amazon.com> * fix doc test Signed-off-by: xinyual <xinyual@amazon.com> * revert useless change Signed-off-by: xinyual <xinyual@amazon.com> * add version to doc Signed-off-by: xinyual <xinyual@amazon.com> * change to use opensearch core code Signed-off-by: xinyual <xinyual@amazon.com> * remove dependency Signed-off-by: xinyual <xinyual@amazon.com> * refactor code to move function to opensearch package Signed-off-by: xinyual <xinyual@amazon.com> * refactor code Signed-off-by: xinyual <xinyual@amazon.com> * add Doc and change to concurrency map Signed-off-by: xinyual <xinyual@amazon.com> --------- Signed-off-by: xinyual <xinyual@amazon.com>
1 parent a696618 commit 055b2ed

13 files changed

Lines changed: 337 additions & 86 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1029,7 +1029,7 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
10291029

10301030
RexNode thenExpr;
10311031
switch (trendlineComputation.getComputationType()) {
1032-
case TrendlineType.SMA:
1032+
case SMA:
10331033
// THEN avg(field) over (ROWS (windowSize-1) PRECEDING)
10341034
thenExpr =
10351035
PlanUtils.makeOver(
@@ -1041,7 +1041,7 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
10411041
List.of(),
10421042
windowFrame);
10431043
break;
1044-
case TrendlineType.WMA:
1044+
case WMA:
10451045
// THEN wma expression
10461046
thenExpr =
10471047
buildWmaRexNode(

‎core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java‎

Lines changed: 2 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,6 @@
1010
import static org.apache.calcite.rex.RexWindowBounds.UNBOUNDED_PRECEDING;
1111
import static org.apache.calcite.rex.RexWindowBounds.following;
1212
import static org.apache.calcite.rex.RexWindowBounds.preceding;
13-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_POP_NULLABLE;
14-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_SAMP_NULLABLE;
15-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_POP_NULLABLE;
16-
import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_SAMP_NULLABLE;
17-
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.TransferUserDefinedAggFunction;
1813

1914
import com.google.common.collect.ImmutableList;
2015
import java.util.ArrayList;
@@ -26,7 +21,6 @@
2621
import org.apache.calcite.rex.RexVisitorImpl;
2722
import org.apache.calcite.rex.RexWindowBound;
2823
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
29-
import org.apache.calcite.sql.type.ReturnTypes;
3024
import org.apache.calcite.sql.type.SqlTypeName;
3125
import org.apache.calcite.tools.RelBuilder;
3226
import org.opensearch.sql.ast.AbstractNodeVisitor;
@@ -38,9 +32,8 @@
3832
import org.opensearch.sql.ast.tree.Relation;
3933
import org.opensearch.sql.ast.tree.UnresolvedPlan;
4034
import org.opensearch.sql.calcite.CalcitePlanContext;
41-
import org.opensearch.sql.calcite.udf.udaf.PercentileApproxFunction;
42-
import org.opensearch.sql.calcite.udf.udaf.TakeAggFunction;
4335
import org.opensearch.sql.expression.function.BuiltinFunctionName;
36+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
4437

4538
public interface PlanUtils {
4639

@@ -268,56 +261,7 @@ static RelBuilder.AggCall makeAggCall(
268261
boolean distinct,
269262
RexNode field,
270263
List<RexNode> argList) {
271-
switch (functionName) {
272-
case MAX:
273-
return context.relBuilder.max(field);
274-
case MIN:
275-
return context.relBuilder.min(field);
276-
case AVG:
277-
return context.relBuilder.avg(distinct, null, field);
278-
case COUNT:
279-
return context.relBuilder.count(
280-
distinct, null, field == null ? ImmutableList.of() : ImmutableList.of(field));
281-
case SUM:
282-
return context.relBuilder.sum(distinct, null, field);
283-
// case MEAN:
284-
// throw new UnsupportedOperationException("MEAN is not supported in PPL");
285-
// case STDDEV:
286-
// return context.relBuilder.aggregateCall(SqlStdOperatorTable.STDDEV,
287-
// field);
288-
case VARSAMP:
289-
return context.relBuilder.aggregateCall(VAR_SAMP_NULLABLE, field);
290-
case VARPOP:
291-
return context.relBuilder.aggregateCall(VAR_POP_NULLABLE, field);
292-
case STDDEV_POP:
293-
return context.relBuilder.aggregateCall(STDDEV_POP_NULLABLE, field);
294-
case STDDEV_SAMP:
295-
return context.relBuilder.aggregateCall(STDDEV_SAMP_NULLABLE, field);
296-
// case PERCENTILE_APPROX:
297-
// return
298-
// context.relBuilder.aggregateCall(SqlStdOperatorTable.PERCENTILE_CONT, field);
299-
case TAKE:
300-
return TransferUserDefinedAggFunction(
301-
TakeAggFunction.class,
302-
"TAKE",
303-
UserDefinedFunctionUtils.getReturnTypeInferenceForArray(),
304-
List.of(field),
305-
argList,
306-
context.relBuilder);
307-
case PERCENTILE_APPROX:
308-
List<RexNode> newArgList = new ArrayList<>(argList);
309-
newArgList.add(context.rexBuilder.makeFlag(field.getType().getSqlTypeName()));
310-
return TransferUserDefinedAggFunction(
311-
PercentileApproxFunction.class,
312-
"percentile_approx",
313-
ReturnTypes.ARG0_FORCE_NULLABLE,
314-
List.of(field),
315-
newArgList,
316-
context.relBuilder);
317-
default:
318-
throw new UnsupportedOperationException(
319-
"Unexpected aggregation: " + functionName.getName().getFunctionName());
320-
}
264+
return PPLFuncImpTable.INSTANCE.resolveAgg(functionName, distinct, field, argList, context);
321265
}
322266

323267
/** Get all uniq input references from a RexNode. */

‎core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public static RelBuilder.AggCall TransferUserDefinedAggFunction(
7878
return relBuilder.aggregateCall(sqlUDAF, addArgList);
7979
}
8080

81-
static SqlReturnTypeInference getReturnTypeInferenceForArray() {
81+
public static SqlReturnTypeInference getReturnTypeInferenceForArray() {
8282
return opBinding -> {
8383
RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
8484

‎core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ public enum BuiltinFunctionName {
195195
TAKE(FunctionName.of("take")),
196196
// t-digest percentile which is used in OpenSearch core by default.
197197
PERCENTILE_APPROX(FunctionName.of("percentile_approx")),
198+
DISTINCT_COUNT_APPROX(FunctionName.of("distinct_count_approx")),
198199
// Not always an aggregation query
199200
NESTED(FunctionName.of("nested")),
200201

@@ -334,6 +335,7 @@ public enum BuiltinFunctionName {
334335
.put("take", BuiltinFunctionName.TAKE)
335336
.put("percentile", BuiltinFunctionName.PERCENTILE_APPROX)
336337
.put("percentile_approx", BuiltinFunctionName.PERCENTILE_APPROX)
338+
.put("distinct_count_approx", BuiltinFunctionName.DISTINCT_COUNT_APPROX)
337339
.build();
338340

339341
private static final Map<String, BuiltinFunctionName> WINDOW_FUNC_MAPPING =

0 commit comments

Comments
 (0)