Skip to content

Eliminate duplicate expression-tree traversal #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
siddharthteotia wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Eliminate duplicate expression-tree traversal #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
85a677d
Eliminate duplicated expression-tree traversal across broker function…
siddharthteotia Mar 27, 2026
98afe2f
Add regression tests for broker function override traversal
siddharthteotia Mar 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
271 changes: 96 additions & 175 deletions ...main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1310,51 +1311,64 @@ 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<Expression> selectList = pinotQuery.getSelectList();
for (Expression expression : selectList) {
handleHLLLog2mOverride(expression, hllLog2mOverride);
private static void traverseAggregateFunctions(PinotQuery pinotQuery, Predicate<Function> visitor) {
for (Expression expression : pinotQuery.getSelectList()) {
traverseFunctions(expression, visitor);
}
List<Expression> 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<Function> 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.
*/
@VisibleForTesting
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.
*/
Expand All @@ -1374,126 +1388,52 @@ static void handleSegmentPartitionedDistinctCountOverride(PinotQuery pinotQuery,
if (segmentPartitionedColumns.isEmpty()) {
return;
}
for (Expression expression : pinotQuery.getSelectList()) {
handleSegmentPartitionedDistinctCountOverride(expression, segmentPartitionedColumns);
}
List<Expression> 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<String> segmentPartitionedColumns) {
Function function = expression.getFunctionCall();
if (function == null) {
return;
}
if (function.getOperator().equals("distinctcount")) {
List<Expression> 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<Expression> 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<Expression> 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));
@VisibleForTesting
static void handleDistinctCountBitmapOverride(PinotQuery pinotQuery) {
traverseAggregateFunctions(pinotQuery, function -> {
if (function.getOperator().equals("distinctcount")) {
function.setOperator("distinctcountbitmap");
return false;
}
}
Expression havingExpression = pinotQuery.getHavingExpression();
if (havingExpression != null) {
handleDistinctCountBitmapOverride(havingExpression);
}
return true;
});
}

/**
* Rewrites selected 'Distinct' prefixed function to 'Distinct----MV' function for the field of multivalued type.
*/
@VisibleForTesting
static void handleDistinctMultiValuedOverride(PinotQuery pinotQuery, Schema tableSchema) {
for (Expression expression : pinotQuery.getSelectList()) {
handleDistinctMultiValuedOverride(expression, tableSchema);
}
List<Expression> 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<Expression> 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<Expression> 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,
Expand Down Expand Up @@ -1613,56 +1553,37 @@ private static void groovySecureAnalysis(Function function) {
*/
@VisibleForTesting
static void handleApproximateFunctionOverride(PinotQuery pinotQuery) {
for (Expression expression : pinotQuery.getSelectList()) {
handleApproximateFunctionOverride(expression);
}
List<Expression> 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,
Expand Down
Loading
Loading