From 40c86362b18f28c460573323d500a405bdebfb9c Mon Sep 17 00:00:00 2001 From: Prerana Singhal Date: Thu, 15 Dec 2022 15:51:13 +0530 Subject: [PATCH 1/2] Added support for accurate distinct count on attributes through config Existing config supports configuring the distinct count aggregation function for a given scope for all attributes. But we might need accurate distinct count on some attributes if the default is set to the approximate function. Defined a way to configure such attributes. --- .../converters/PinotFunctionConverter.java | 14 ++++---- .../PinotFunctionConverterConfig.java | 27 ++++++++++++-- .../PinotFunctionConverterTest.java | 36 ++++++++++++++----- 3 files changed, 59 insertions(+), 18 deletions(-) diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverter.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverter.java index 76eef51f..9e5e41ed 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverter.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverter.java @@ -44,7 +44,7 @@ public String convert( // and reasonably accurate, so support selecting the implementation to use return this.functionToString(this.toPinotPercentile(function), argumentConverter); case QUERY_FUNCTION_DISTINCTCOUNT: - return this.functionToString(this.toPinotDistinctCount(function), argumentConverter); + return this.functionToStringForDistinctCount(function, argumentConverter); case QUERY_FUNCTION_CONCAT: return this.functionToString(this.toPinotConcat(function), argumentConverter); case QUERY_FUNCTION_AVGRATE: @@ -72,6 +72,12 @@ private String functionToString( return function.getFunctionName() + "(" + argumentString + ")"; } + private String functionToStringForDistinctCount( + Function function, java.util.function.Function argumentConverter) { + String columnName = argumentConverter.apply(function.getArgumentsList().get(0)); + return this.config.getDistinctCountFunction(columnName) + "(" + columnName + ")"; + } + private String functionToStringForAvgRate( Function function, java.util.function.Function argumentConverter, @@ -119,12 +125,6 @@ private Function toPinotConcat(Function function) { return Function.newBuilder(function).setFunctionName(PINOT_CONCAT_FUNCTION).build(); } - private Function toPinotDistinctCount(Function function) { - return Function.newBuilder(function) - .setFunctionName(this.config.getDistinctCountAggregationFunction()) - .build(); - } - private boolean isHardcodedPercentile(Function function) { String functionName = function.getFunctionName().toUpperCase(); return functionName.startsWith(QUERY_FUNCTION_PERCENTILE) diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java index 54f55467..31c9436e 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java @@ -1,22 +1,30 @@ package org.hypertrace.core.query.service.pinot.converters; +import com.google.common.collect.ImmutableSet; import com.typesafe.config.Config; import com.typesafe.config.ConfigFactory; +import java.util.Collections; +import java.util.Set; +import javax.annotation.Nonnull; +import lombok.AccessLevel; import lombok.AllArgsConstructor; import lombok.Value; @Value -@AllArgsConstructor +@AllArgsConstructor(access = AccessLevel.PACKAGE) public class PinotFunctionConverterConfig { private static final String PERCENTILE_AGGREGATION_FUNCTION_CONFIG = "percentileAggFunction"; private static final String DISTINCT_COUNT_AGGREGATION_FUNCTION_CONFIG = "distinctCountAggFunction"; + private static final String ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG = + "argsForAccurateDistinctCountAgg"; private static final String DEFAULT_PERCENTILE_AGGREGATION_FUNCTION = "PERCENTILETDIGEST"; - private static final String DEFAULT_DISTINCT_COUNT_AGGREGATION_FUNCTION = "DISTINCTCOUNT"; + private static final String ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION = "DISTINCTCOUNT"; String percentileAggregationFunction; String distinctCountAggregationFunction; + @Nonnull Set argsForAccurateDistinctCountAgg; public PinotFunctionConverterConfig(Config config) { if (config.hasPath(PERCENTILE_AGGREGATION_FUNCTION_CONFIG)) { @@ -28,11 +36,24 @@ public PinotFunctionConverterConfig(Config config) { this.distinctCountAggregationFunction = config.getString(DISTINCT_COUNT_AGGREGATION_FUNCTION_CONFIG); } else { - this.distinctCountAggregationFunction = DEFAULT_DISTINCT_COUNT_AGGREGATION_FUNCTION; + this.distinctCountAggregationFunction = ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION; + } + if (config.hasPath(ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG)) { + this.argsForAccurateDistinctCountAgg = + ImmutableSet.copyOf(config.getStringList(ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG)); + } else { + this.argsForAccurateDistinctCountAgg = Collections.emptySet(); } } public PinotFunctionConverterConfig() { this(ConfigFactory.empty()); } + + public String getDistinctCountFunction(String arg) { + if (argsForAccurateDistinctCountAgg.contains(arg)) { + return ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION; + } + return distinctCountAggregationFunction; + } } diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java index dbf5ed5e..cd3279e0 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java @@ -20,7 +20,9 @@ import java.time.Duration; import java.util.Arrays; +import java.util.Collections; import java.util.Optional; +import java.util.Set; import java.util.stream.Collectors; import org.hypertrace.core.query.service.ExecutionContext; import org.hypertrace.core.query.service.api.Expression; @@ -104,7 +106,8 @@ void acceptsCustomPercentileFunctions() { assertEquals( expected, - new PinotFunctionConverter(new PinotFunctionConverterConfig("CUSTOMPERCENTILE", null)) + new PinotFunctionConverter( + new PinotFunctionConverterConfig("CUSTOMPERCENTILE", null, Collections.emptySet())) .convert(mockingExecutionContext, percentileFunction, this.mockArgumentConverter)); } @@ -233,25 +236,42 @@ void convertAvgRateFunction() { @Test void convertsDistinctCountFunction() { - Expression column = createColumnExpression("foo").build(); + Expression column1 = createColumnExpression("foo").build(); + Expression column2 = createColumnExpression("bar").build(); - when(this.mockArgumentConverter.apply(column)).thenReturn("foo"); + when(this.mockArgumentConverter.apply(column1)).thenReturn("foo"); + when(this.mockArgumentConverter.apply(column2)).thenReturn("bar"); assertEquals( "DISTINCTCOUNT(foo)", new PinotFunctionConverter() .convert( mockingExecutionContext, - buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column.toBuilder()), + buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column1.toBuilder()), this.mockArgumentConverter)); - assertEquals( - "CUSTOM_DC(foo)", - new PinotFunctionConverter(new PinotFunctionConverterConfig(null, "CUSTOM_DC")) + "DISTINCTCOUNT(bar)", + new PinotFunctionConverter() .convert( mockingExecutionContext, - buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column.toBuilder()), + buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column2.toBuilder()), this.mockArgumentConverter)); + + PinotFunctionConverter converter = + new PinotFunctionConverter( + new PinotFunctionConverterConfig(null, "CUSTOM_DC", Set.of("bar"))); + assertEquals( + "CUSTOM_DC(foo)", + converter.convert( + mockingExecutionContext, + buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column1.toBuilder()), + this.mockArgumentConverter)); + assertEquals( + "DISTINCTCOUNT(bar)", + converter.convert( + mockingExecutionContext, + buildFunction(QUERY_FUNCTION_DISTINCTCOUNT, column2.toBuilder()), + this.mockArgumentConverter)); } @Test From f51ccb788a0f8a35f0d92cfe42d6ca3bf8a7f5b9 Mon Sep 17 00:00:00 2001 From: Prerana Singhal Date: Tue, 10 Jan 2023 10:06:16 +0530 Subject: [PATCH 2/2] Added generic overrides map --- .../PinotFunctionConverterConfig.java | 32 ++++++++++--------- .../PinotFunctionConverterTest.java | 7 ++-- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java index 31c9436e..72417c17 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterConfig.java @@ -1,10 +1,11 @@ package org.hypertrace.core.query.service.pinot.converters; -import com.google.common.collect.ImmutableSet; import com.typesafe.config.Config; import com.typesafe.config.ConfigFactory; import java.util.Collections; -import java.util.Set; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Collectors; import javax.annotation.Nonnull; import lombok.AccessLevel; import lombok.AllArgsConstructor; @@ -17,14 +18,13 @@ public class PinotFunctionConverterConfig { private static final String PERCENTILE_AGGREGATION_FUNCTION_CONFIG = "percentileAggFunction"; private static final String DISTINCT_COUNT_AGGREGATION_FUNCTION_CONFIG = "distinctCountAggFunction"; - private static final String ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG = - "argsForAccurateDistinctCountAgg"; + private static final String DISTINCT_COUNT_AGGREGATION_OVERRIDES = "distinctCountAggOverrides"; private static final String DEFAULT_PERCENTILE_AGGREGATION_FUNCTION = "PERCENTILETDIGEST"; - private static final String ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION = "DISTINCTCOUNT"; + private static final String DEFAULT_DISTINCT_COUNT_AGGREGATION_FUNCTION = "DISTINCTCOUNT"; String percentileAggregationFunction; String distinctCountAggregationFunction; - @Nonnull Set argsForAccurateDistinctCountAgg; + @Nonnull Map distinctCountAggOverrides; public PinotFunctionConverterConfig(Config config) { if (config.hasPath(PERCENTILE_AGGREGATION_FUNCTION_CONFIG)) { @@ -36,13 +36,17 @@ public PinotFunctionConverterConfig(Config config) { this.distinctCountAggregationFunction = config.getString(DISTINCT_COUNT_AGGREGATION_FUNCTION_CONFIG); } else { - this.distinctCountAggregationFunction = ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION; + this.distinctCountAggregationFunction = DEFAULT_DISTINCT_COUNT_AGGREGATION_FUNCTION; } - if (config.hasPath(ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG)) { - this.argsForAccurateDistinctCountAgg = - ImmutableSet.copyOf(config.getStringList(ARGS_FOR_ACCURATE_DISTINCT_COUNT_AGG_CONFIG)); + if (config.hasPath(DISTINCT_COUNT_AGGREGATION_OVERRIDES)) { + Config overridesConfig = config.getConfig(DISTINCT_COUNT_AGGREGATION_OVERRIDES); + this.distinctCountAggOverrides = + overridesConfig.entrySet().stream() + .collect( + Collectors.toMap( + entry -> entry.getKey(), entry -> overridesConfig.getString(entry.getKey()))); } else { - this.argsForAccurateDistinctCountAgg = Collections.emptySet(); + this.distinctCountAggOverrides = Collections.emptyMap(); } } @@ -51,9 +55,7 @@ public PinotFunctionConverterConfig() { } public String getDistinctCountFunction(String arg) { - if (argsForAccurateDistinctCountAgg.contains(arg)) { - return ACCURATE_DISTINCT_COUNT_AGGREGATION_FUNCTION; - } - return distinctCountAggregationFunction; + return Optional.ofNullable(distinctCountAggOverrides.get(arg)) + .orElse(distinctCountAggregationFunction); } } diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java index cd3279e0..899cfa9e 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/converters/PinotFunctionConverterTest.java @@ -18,11 +18,11 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.when; +import com.typesafe.config.ConfigFactory; import java.time.Duration; import java.util.Arrays; import java.util.Collections; import java.util.Optional; -import java.util.Set; import java.util.stream.Collectors; import org.hypertrace.core.query.service.ExecutionContext; import org.hypertrace.core.query.service.api.Expression; @@ -107,7 +107,7 @@ void acceptsCustomPercentileFunctions() { assertEquals( expected, new PinotFunctionConverter( - new PinotFunctionConverterConfig("CUSTOMPERCENTILE", null, Collections.emptySet())) + new PinotFunctionConverterConfig("CUSTOMPERCENTILE", null, Collections.emptyMap())) .convert(mockingExecutionContext, percentileFunction, this.mockArgumentConverter)); } @@ -259,7 +259,8 @@ void convertsDistinctCountFunction() { PinotFunctionConverter converter = new PinotFunctionConverter( - new PinotFunctionConverterConfig(null, "CUSTOM_DC", Set.of("bar"))); + new PinotFunctionConverterConfig( + ConfigFactory.parseString("distinctCountAggOverrides = {foo=CUSTOM_DC, xyz=XYZ}"))); assertEquals( "CUSTOM_DC(foo)", converter.convert(