From 85a677dd8b8c990217d99a493dafe22cd549ae06 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 27 Mar 2026 08:57:00 -0700 Subject: [PATCH 1/2] Eliminate duplicated expression-tree traversal across broker function override methods All five aggregate-function override methods (handleHLLLog2mOverride, handleSegmentPartitionedDistinctCountOverride, handleDistinctCountBitmapOverride, handleDistinctMultiValuedOverride, handleApproximateFunctionOverride) shared an identical two-layer pattern: a PinotQuery-level entry that fans out to select/orderBy/having, and a recursive expression-level method that either transforms a matched function or recurses into operands. Introduce two private static helpers: - traverseAggregateFunctions(PinotQuery, Predicate): fans out over select, orderBy (unwrapping the sort wrapper), and having - traverseFunctions(Expression, Predicate): recurses depth-first; the predicate returns true to continue into operands, false to prune the subtree Rewrite each override method as a single call to traverseAggregateFunctions with a lambda that encodes only its transform logic. This removes five expression-level private overloads (~100 lines) and makes the traversal contract explicit and testable in isolation. Co-Authored-By: Claude Sonnet 4.6 --- .../BaseSingleStageBrokerRequestHandler.java | 267 ++++++------------ 1 file changed, 93 insertions(+), 174 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 31105bbb02..32c20d1e5b 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -31,6 +31,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Predicate; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CompletionService; import java.util.concurrent.ConcurrentHashMap; @@ -1310,51 +1311,63 @@ private void setTableName(BrokerRequest brokerRequest, String tableName) { } /** - * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given query. + * Visits every {@link Function} node reachable from the query's aggregate-context expressions (SELECT list, + * ORDER BY, HAVING). The {@code visitor} returns {@code true} to continue recursing into the function's operands, + * or {@code false} to prune the subtree (e.g. after transforming a matched function). */ - private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { - List selectList = pinotQuery.getSelectList(); - for (Expression expression : selectList) { - handleHLLLog2mOverride(expression, hllLog2mOverride); + private static void traverseAggregateFunctions(PinotQuery pinotQuery, Predicate visitor) { + for (Expression expression : pinotQuery.getSelectList()) { + traverseFunctions(expression, visitor); } List orderByList = pinotQuery.getOrderByList(); if (orderByList != null) { for (Expression expression : orderByList) { // NOTE: Order-by is always a Function with the ordering of the Expression - handleHLLLog2mOverride(expression.getFunctionCall().getOperands().get(0), hllLog2mOverride); + traverseFunctions(expression.getFunctionCall().getOperands().get(0), visitor); } } Expression havingExpression = pinotQuery.getHavingExpression(); if (havingExpression != null) { - handleHLLLog2mOverride(havingExpression, hllLog2mOverride); + traverseFunctions(havingExpression, visitor); } } /** - * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given expression. + * Recursively visits each {@link Function} in an expression tree, invoking {@code visitor} on every function node. + * The visitor returns {@code true} to recurse into the function's operands, {@code false} to stop. */ - private static void handleHLLLog2mOverride(Expression expression, int hllLog2mOverride) { + private static void traverseFunctions(Expression expression, Predicate visitor) { Function function = expression.getFunctionCall(); if (function == null) { return; } - switch (function.getOperator()) { - case "distinctcounthll": - case "distinctcounthllmv": - case "distinctcountrawhll": - case "distinctcountrawhllmv": - if (function.getOperandsSize() == 1) { - function.addToOperands(RequestUtils.getLiteralExpression(hllLog2mOverride)); - } - return; - default: - break; - } - for (Expression operand : function.getOperands()) { - handleHLLLog2mOverride(operand, hllLog2mOverride); + if (visitor.test(function)) { + for (Expression operand : function.getOperands()) { + traverseFunctions(operand, visitor); + } } } + /** + * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given query. + */ + private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { + traverseAggregateFunctions(pinotQuery, function -> { + switch (function.getOperator()) { + case "distinctcounthll": + case "distinctcounthllmv": + case "distinctcountrawhll": + case "distinctcountrawhllmv": + if (function.getOperandsSize() == 1) { + function.addToOperands(RequestUtils.getLiteralExpression(hllLog2mOverride)); + } + return false; + default: + return true; + } + }); + } + /** * Overrides the LIMIT of the given query if it exceeds the query limit. */ @@ -1374,63 +1387,30 @@ static void handleSegmentPartitionedDistinctCountOverride(PinotQuery pinotQuery, if (segmentPartitionedColumns.isEmpty()) { return; } - for (Expression expression : pinotQuery.getSelectList()) { - handleSegmentPartitionedDistinctCountOverride(expression, segmentPartitionedColumns); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleSegmentPartitionedDistinctCountOverride(expression.getFunctionCall().getOperands().get(0), - segmentPartitionedColumns); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleSegmentPartitionedDistinctCountOverride(havingExpression, segmentPartitionedColumns); - } - } - - /** - * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given expression. - */ - private static void handleSegmentPartitionedDistinctCountOverride(Expression expression, - Set segmentPartitionedColumns) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - if (function.getOperator().equals("distinctcount")) { - List operands = function.getOperands(); - if (operands.size() == 1 && operands.get(0).isSetIdentifier() && segmentPartitionedColumns.contains( - operands.get(0).getIdentifier().getName())) { - function.setOperator("segmentpartitioneddistinctcount"); - } - } else { - for (Expression operand : function.getOperands()) { - handleSegmentPartitionedDistinctCountOverride(operand, segmentPartitionedColumns); + traverseAggregateFunctions(pinotQuery, function -> { + if (function.getOperator().equals("distinctcount")) { + List operands = function.getOperands(); + if (operands.size() == 1 && operands.get(0).isSetIdentifier() + && segmentPartitionedColumns.contains(operands.get(0).getIdentifier().getName())) { + function.setOperator("segmentpartitioneddistinctcount"); + } + return false; } - } + return true; + }); } /** * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given query. */ private static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { - for (Expression expression : pinotQuery.getSelectList()) { - handleDistinctCountBitmapOverride(expression); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleDistinctCountBitmapOverride(expression.getFunctionCall().getOperands().get(0)); + traverseAggregateFunctions(pinotQuery, function -> { + if (function.getOperator().equals("distinctcount")) { + function.setOperator("distinctcountbitmap"); + return false; } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleDistinctCountBitmapOverride(havingExpression); - } + return true; + }); } /** @@ -1438,62 +1418,20 @@ private static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { */ @VisibleForTesting static void handleDistinctMultiValuedOverride(PinotQuery pinotQuery, Schema tableSchema) { - for (Expression expression : pinotQuery.getSelectList()) { - handleDistinctMultiValuedOverride(expression, tableSchema); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleDistinctMultiValuedOverride(expression.getFunctionCall().getOperands().get(0), tableSchema); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleDistinctMultiValuedOverride(havingExpression, tableSchema); - } - } - - /** - * Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type. - */ - private static void handleDistinctMultiValuedOverride(Expression expression, Schema tableSchema) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - - String overrideOperator = DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP.get(function.getOperator()); - if (overrideOperator != null) { - List operands = function.getOperands(); - if (operands.size() >= 1 && operands.get(0).isSetIdentifier() && isMultiValueColumn(tableSchema, - operands.get(0).getIdentifier().getName())) { - // we are only checking the first operand that if its a MV column as all the overriding agg. fn.'s have - // first operator is column name - function.setOperator(overrideOperator); - } - } else { - for (Expression operand : function.getOperands()) { - handleDistinctMultiValuedOverride(operand, tableSchema); - } - } - } - - /** - * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given expression. - */ - private static void handleDistinctCountBitmapOverride(Expression expression) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - if (function.getOperator().equals("distinctcount")) { - function.setOperator("distinctcountbitmap"); - } else { - for (Expression operand : function.getOperands()) { - handleDistinctCountBitmapOverride(operand); + traverseAggregateFunctions(pinotQuery, function -> { + String overrideOperator = DISTINCT_MV_COL_FUNCTION_OVERRIDE_MAP.get(function.getOperator()); + if (overrideOperator != null) { + List operands = function.getOperands(); + if (operands.size() >= 1 && operands.get(0).isSetIdentifier() + && isMultiValueColumn(tableSchema, operands.get(0).getIdentifier().getName())) { + // we are only checking the first operand that if its a MV column as all the overriding agg. fn.'s have + // first operator is column name + function.setOperator(overrideOperator); + } + return false; } - } + return true; + }); } private HandlerContext getHandlerContext(@Nullable QueryConfig offlineTableQueryConfig, @@ -1613,56 +1551,37 @@ private static void groovySecureAnalysis(Function function) { */ @VisibleForTesting static void handleApproximateFunctionOverride(PinotQuery pinotQuery) { - for (Expression expression : pinotQuery.getSelectList()) { - handleApproximateFunctionOverride(expression); - } - List orderByExpressions = pinotQuery.getOrderByList(); - if (orderByExpressions != null) { - for (Expression expression : orderByExpressions) { - // NOTE: Order-by is always a Function with the ordering of the Expression - handleApproximateFunctionOverride(expression.getFunctionCall().getOperands().get(0)); - } - } - Expression havingExpression = pinotQuery.getHavingExpression(); - if (havingExpression != null) { - handleApproximateFunctionOverride(havingExpression); - } - } - - private static void handleApproximateFunctionOverride(Expression expression) { - Function function = expression.getFunctionCall(); - if (function == null) { - return; - } - String functionName = function.getOperator(); - if (functionName.equals("distinctcount") || functionName.equals("distinctcountmv")) { - function.setOperator("distinctcountsmarthll"); - } else if (functionName.startsWith("percentile")) { - String remainingFunctionName = functionName.substring(10); - if (remainingFunctionName.isEmpty() || remainingFunctionName.equals("mv")) { - function.setOperator("percentilesmarttdigest"); - } else if (remainingFunctionName.matches("\\d+")) { - try { - int percentile = Integer.parseInt(remainingFunctionName); - function.setOperator("percentilesmarttdigest"); - function.addToOperands(RequestUtils.getLiteralExpression(percentile)); - } catch (Exception e) { - throw new BadQueryRequestException("Illegal function name: " + functionName); - } - } else if (remainingFunctionName.matches("\\d+mv")) { - try { - int percentile = Integer.parseInt(remainingFunctionName.substring(0, remainingFunctionName.length() - 2)); + traverseAggregateFunctions(pinotQuery, function -> { + String functionName = function.getOperator(); + if (functionName.equals("distinctcount") || functionName.equals("distinctcountmv")) { + function.setOperator("distinctcountsmarthll"); + return false; + } else if (functionName.startsWith("percentile")) { + String remainingFunctionName = functionName.substring(10); + if (remainingFunctionName.isEmpty() || remainingFunctionName.equals("mv")) { function.setOperator("percentilesmarttdigest"); - function.addToOperands(RequestUtils.getLiteralExpression(percentile)); - } catch (Exception e) { - throw new BadQueryRequestException("Illegal function name: " + functionName); + } else if (remainingFunctionName.matches("\\d+")) { + try { + int percentile = Integer.parseInt(remainingFunctionName); + function.setOperator("percentilesmarttdigest"); + function.addToOperands(RequestUtils.getLiteralExpression(percentile)); + } catch (Exception e) { + throw new BadQueryRequestException("Illegal function name: " + functionName); + } + } else if (remainingFunctionName.matches("\\d+mv")) { + try { + int percentile = + Integer.parseInt(remainingFunctionName.substring(0, remainingFunctionName.length() - 2)); + function.setOperator("percentilesmarttdigest"); + function.addToOperands(RequestUtils.getLiteralExpression(percentile)); + } catch (Exception e) { + throw new BadQueryRequestException("Illegal function name: " + functionName); + } } + return false; } - } else { - for (Expression operand : function.getOperands()) { - handleApproximateFunctionOverride(operand); - } - } + return true; + }); } private static void handleExpressionOverride(PinotQuery pinotQuery, From 98afe2f2fea3644487facd4632f84b67da551aaf Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 27 Mar 2026 09:32:13 -0700 Subject: [PATCH 2/2] Add regression tests for broker function override traversal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promote handleHLLLog2mOverride and handleDistinctCountBitmapOverride to @VisibleForTesting package-private so they can be exercised directly. Add QueryOverrideTest cases covering: - All four HLL variants; no-op when log2m already set; non-HLL untouched - distinctcount → distinctcountbitmap rewrite; non-matching untouched - Traversal descends into non-matching wrapper functions (ADD) - Subtree pruning after a match (inner HLL is not visited) - orderBy and having coverage for segmentPartitioned and multiValued overrides Co-Authored-By: Claude Sonnet 4.6 --- .../BaseSingleStageBrokerRequestHandler.java | 6 +- .../requesthandler/QueryOverrideTest.java | 126 ++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 32c20d1e5b..6b990ffaad 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -1351,7 +1351,8 @@ private static void traverseFunctions(Expression expression, Predicate /** * Sets HyperLogLog log2m for DistinctCountHLL functions if not explicitly set for the given query. */ - private static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { + @VisibleForTesting + static void handleHLLLog2mOverride(PinotQuery pinotQuery, int hllLog2mOverride) { traverseAggregateFunctions(pinotQuery, function -> { switch (function.getOperator()) { case "distinctcounthll": @@ -1403,7 +1404,8 @@ static void handleSegmentPartitionedDistinctCountOverride(PinotQuery pinotQuery, /** * Rewrites 'DistinctCount' to 'DistinctCountBitmap' for the given query. */ - private static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { + @VisibleForTesting + static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) { traverseAggregateFunctions(pinotQuery, function -> { if (function.getOperator().equals("distinctcount")) { function.setOperator("distinctcountbitmap"); diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java index 22137c137b..0e62c62a5b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryOverrideTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableSet; import java.io.IOException; import java.util.Arrays; +import org.apache.pinot.common.request.Function; import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.spi.data.Schema; @@ -54,6 +55,131 @@ private void testLimitOverride(String query, int expectedLimit) { assertEquals(pinotQuery.getLimit(), expectedLimit); } + @Test + public void testHLLLog2mOverride() { + int log2m = 8; + + // log2m is appended when the function has exactly one operand (the column) + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT_HLL(col1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, log2m); + Function f = pinotQuery.getSelectList().get(0).getFunctionCall(); + assertEquals(f.getOperandsSize(), 2); + assertEquals(f.getOperands().get(1), RequestUtils.getLiteralExpression(log2m)); + + // log2m is NOT overridden when already present (two operands) + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT_HLL(col1, 5) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, log2m); + f = pinotQuery.getSelectList().get(0).getFunctionCall(); + assertEquals(f.getOperandsSize(), 2); + assertEquals(f.getOperands().get(1), RequestUtils.getLiteralExpression(5)); + + // All four HLL variants are matched + for (String func : Arrays.asList( + "DISTINCT_COUNT_HLL", "DISTINCT_COUNT_HLL_MV", + "DISTINCT_COUNT_RAW_HLL", "DISTINCT_COUNT_RAW_HLL_MV")) { + pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT " + func + "(col1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, log2m); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperandsSize(), 2, func); + } + + // Non-HLL function is not touched + pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT DISTINCT_COUNT(col1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, log2m); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperandsSize(), 1); + + // Applied across select, orderBy, and having + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT_HLL(col1) FROM myTable GROUP BY col2 " + + "HAVING DISTINCT_COUNT_HLL(col1) > 10 ORDER BY DISTINCT_COUNT_HLL(col1) DESC"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, log2m); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperandsSize(), 2); + assertEquals(pinotQuery.getOrderByList().get(0).getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperandsSize(), 2); + assertEquals(pinotQuery.getHavingExpression().getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperandsSize(), 2); + } + + @Test + public void testDistinctCountBitmapOverride() { + // Rewrites distinctcount to distinctcountbitmap + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT(col1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleDistinctCountBitmapOverride(pinotQuery); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountbitmap"); + + // Non-matching function is not touched + pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT SUM(col1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleDistinctCountBitmapOverride(pinotQuery); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "sum"); + + // Applied across select, orderBy, and having + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT(col1) FROM myTable GROUP BY col2 " + + "HAVING DISTINCT_COUNT(col1) > 10 ORDER BY DISTINCT_COUNT(col1) DESC"); + BaseSingleStageBrokerRequestHandler.handleDistinctCountBitmapOverride(pinotQuery); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountbitmap"); + assertEquals(pinotQuery.getOrderByList().get(0).getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "distinctcountbitmap"); + assertEquals(pinotQuery.getHavingExpression().getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "distinctcountbitmap"); + } + + @Test + public void testFunctionOverrideTraversalBehavior() { + // Non-matching wrapper: traversal must descend into ADD to find and rewrite the nested DISTINCT_COUNT + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT ADD(DISTINCT_COUNT(col1), 1) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleDistinctCountBitmapOverride(pinotQuery); + Function add = pinotQuery.getSelectList().get(0).getFunctionCall(); + assertEquals(add.getOperator(), "add"); + assertEquals(add.getOperands().get(0).getFunctionCall().getOperator(), "distinctcountbitmap"); + + // Matched function: subtree is pruned — the inner DISTINCT_COUNT_HLL must not receive log2m + // because traversal stops after rewriting the outer one + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT_HLL(DISTINCT_COUNT_HLL(col1)) FROM myTable"); + BaseSingleStageBrokerRequestHandler.handleHLLLog2mOverride(pinotQuery, 8); + Function outer = pinotQuery.getSelectList().get(0).getFunctionCall(); + assertEquals(outer.getOperandsSize(), 2); // outer got log2m + Function inner = outer.getOperands().get(0).getFunctionCall(); + assertEquals(inner.getOperator(), "distinctcounthll"); + assertEquals(inner.getOperandsSize(), 1); // inner was NOT visited + + // orderBy and having coverage for segment-partitioned override + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT(col1) FROM myTable GROUP BY col2 " + + "HAVING DISTINCT_COUNT(col1) > 10 ORDER BY DISTINCT_COUNT(col1) DESC"); + BaseSingleStageBrokerRequestHandler.handleSegmentPartitionedDistinctCountOverride( + pinotQuery, ImmutableSet.of("col1")); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), + "segmentpartitioneddistinctcount"); + assertEquals(pinotQuery.getOrderByList().get(0).getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "segmentpartitioneddistinctcount"); + assertEquals(pinotQuery.getHavingExpression().getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "segmentpartitioneddistinctcount"); + + // orderBy and having coverage for multi-valued override + Schema tableSchema; + try { + tableSchema = Schema.fromString("{\"schemaName\":\"testSchema\"," + + "\"dimensionFieldSpecs\":[{\"name\":\"col1\",\"dataType\":\"LONG\"," + + "\"singleValueField\":\"false\"}]}"); + } catch (IOException e) { + throw new RuntimeException(e); + } + pinotQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT DISTINCT_COUNT(col1) FROM myTable GROUP BY col2 " + + "HAVING DISTINCT_COUNT(col1) > 10 ORDER BY DISTINCT_COUNT(col1) DESC"); + BaseSingleStageBrokerRequestHandler.handleDistinctMultiValuedOverride(pinotQuery, tableSchema); + assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "distinctcountmv"); + assertEquals(pinotQuery.getOrderByList().get(0).getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "distinctcountmv"); + assertEquals(pinotQuery.getHavingExpression().getFunctionCall().getOperands().get(0) + .getFunctionCall().getOperator(), "distinctcountmv"); + } + @Test public void testDistinctCountOverride() { String query = "SELECT DISTINCT_COUNT(col1) FROM myTable";