Jackie-Jiang commented on a change in pull request #7664:
URL: https://github.com/apache/pinot/pull/7664#discussion_r740512233
##########
File path:
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
##########
@@ -1155,6 +1218,59 @@ private static void
handleDistinctCountBitmapOverride(BrokerRequest brokerReques
}
}
+ /**
+ * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given
SQL query.
+ */
+ @VisibleForTesting
+ static void handleSegmentPartitionedDistinctCountOverride(PinotQuery
pinotQuery,
+ Set<String> segmentPartitionedColumns) {
+ 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
SQL expression.
+ */
+ private static void handleSegmentPartitionedDistinctCountOverride(Expression
expression,
+ Set<String> segmentPartitionedColumns) {
+ Function function = expression.getFunctionCall();
+ if (function == null) {
+ return;
+ }
+ if (StringUtils.remove(function.getOperator(), '_')
+ .equalsIgnoreCase(AggregationFunctionType.DISTINCTCOUNT.name())) {
+ final Set<String> identifiers =
+ expression.getFunctionCall().getOperands().stream().filter(expr ->
expr.isSetIdentifier())
+ .map(expr ->
expr.getIdentifier().getName()).collect(Collectors.toUnmodifiableSet());
Review comment:
Suggest checking if there is a single identifier to avoid unexpected
rewrite
```suggestion
List<Expression> operands = function.getOperands();
if (operands.size() == 1 && operands.get(0).getType() == IDENTIFIER &&
segmentPartitionedColumns.contains(operands.get(0).getIdentifier.getName())) {
function.setOperator(AggregationFunctionType.SEGMENTPARTITIONEDDISTINCTCOUNT.name());
}
```
##########
File path:
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
##########
@@ -1139,6 +1176,32 @@ static void handleQueryLimitOverride(PinotQuery
pinotQuery, int queryLimit) {
}
}
+ /**
+ * Rewrites 'DistinctCount' to 'SegmentPartitionDistinctCount' for the given
PQL broker request.
+ */
+ @Deprecated
+ @VisibleForTesting
+ static void handleSegmentPartitionedDistinctCountOverride(BrokerRequest
brokerRequest,
+ Set<String> segmentPartitionedColumns) {
+ if (segmentPartitionedColumns.isEmpty()) {
+ return;
+ }
+ List<AggregationInfo> aggregationsInfo =
brokerRequest.getAggregationsInfo();
+ if (aggregationsInfo != null) {
+ for (AggregationInfo aggregationInfo : aggregationsInfo) {
+ if (StringUtils.remove(aggregationInfo.getAggregationType(), '_')
+ .equalsIgnoreCase(AggregationFunctionType.DISTINCTCOUNT.name())) {
+ for (String expr : aggregationInfo.getExpressions()) {
Review comment:
Suggest checking if there is a single identifier to avoid unexpected
rewrite
```suggestion
List<String> expressions = aggregationInfo.getExpressions();
if (expressions.size() == 1 &&
segmentPartitionedColumns.contains(expressions.get(0))) {
aggregationInfo.setAggregationType(AggregationFunctionType.SEGMENTPARTITIONEDDISTINCTCOUNT.name());
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
