From 081b818fb60b13e14a1b532685475aaf820dbd41 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 27 Mar 2026 04:41:21 -0700 Subject: [PATCH] Refactor BaseSingleStageBrokerRequestHandler.doHandleRequest() into focused phases Extract five cohesive private methods from the 495-line doHandleRequest(): - prepareOfflineAndRealtimeRequests(): hybrid/offline/realtime branching, time-boundary attachment, expression/timestamp overrides, optimizer, always-false/true pruning - collectRoutingWarnings(): disabled table names and unavailable segment warnings - configureQueryTimeouts(): per-table timeout configuration (throws TimeoutException) - configureServerResponseOptions(): max response size and single-server final-result flag - finalizeResponse(): replica groups, exception propagation, metrics emission, query logging doHandleRequest() is now ~175 lines (down from ~495) with each phase clearly labeled. No behavioral changes; timing semantics preserved by capturing executionEndTimeNs before finalizeResponse(). Co-Authored-By: Claude Sonnet 4.6 --- .../BaseSingleStageBrokerRequestHandler.java | 499 ++++++++++-------- 1 file changed, 287 insertions(+), 212 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..6370df1361 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 @@ -511,13 +511,8 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn return BrokerResponseNative.NO_TABLE_RESULT; } - String offlineTableName = routeInfo.getOfflineTableName(); - String realtimeTableName = routeInfo.getRealtimeTableName(); - TableConfig offlineTableConfig = routeInfo.getOfflineTableConfig(); - TableConfig realtimeTableConfig = routeInfo.getRealtimeTableConfig(); QueryConfig offlineTableQueryConfig = routeInfo.getOfflineTableQueryConfig(); QueryConfig realtimeTableQueryConfig = routeInfo.getRealtimeTableQueryConfig(); - TimeBoundaryInfo timeBoundaryInfo = routeInfo.getTimeBoundaryInfo(); HandlerContext handlerContext = getHandlerContext(offlineTableQueryConfig, realtimeTableQueryConfig); validateGroovyScript(serverPinotQuery, handlerContext._disableGroovy); @@ -550,77 +545,16 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn } // Prepare OFFLINE and REALTIME requests - BrokerRequest offlineBrokerRequest = null; - BrokerRequest realtimeBrokerRequest = null; - - if (routeInfo.isHybrid()) { - // Hybrid - PinotQuery offlinePinotQuery = serverPinotQuery.deepCopy(); - offlinePinotQuery.getDataSource().setTableName(offlineTableName); - assert timeBoundaryInfo != null; - attachTimeBoundary(offlinePinotQuery, timeBoundaryInfo, true); - handleExpressionOverride(offlinePinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); - handleTimestampIndexOverride(offlinePinotQuery, offlineTableConfig); - _queryOptimizer.optimize(offlinePinotQuery, offlineTableConfig, schema); - offlineBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(offlinePinotQuery); - - PinotQuery realtimePinotQuery = serverPinotQuery.deepCopy(); - realtimePinotQuery.getDataSource().setTableName(realtimeTableName); - attachTimeBoundary(realtimePinotQuery, timeBoundaryInfo, false); - handleExpressionOverride(realtimePinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); - handleTimestampIndexOverride(realtimePinotQuery, realtimeTableConfig); - _queryOptimizer.optimize(realtimePinotQuery, realtimeTableConfig, schema); - realtimeBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(realtimePinotQuery); - - requestContext.setFanoutType(RequestContext.FanoutType.HYBRID); - requestContext.setOfflineServerTenant(getServerTenant(offlineTableConfig)); - requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableConfig)); - } else if (routeInfo.isOffline()) { - // OFFLINE only - setTableName(serverBrokerRequest, offlineTableName); - handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); - handleTimestampIndexOverride(serverPinotQuery, offlineTableConfig); - _queryOptimizer.optimize(serverPinotQuery, offlineTableConfig, schema); - offlineBrokerRequest = serverBrokerRequest; - - requestContext.setFanoutType(RequestContext.FanoutType.OFFLINE); - requestContext.setOfflineServerTenant(getServerTenant(offlineTableConfig)); - } else { - // REALTIME only - setTableName(serverBrokerRequest, realtimeTableName); - handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); - handleTimestampIndexOverride(serverPinotQuery, realtimeTableConfig); - _queryOptimizer.optimize(serverPinotQuery, realtimeTableConfig, schema); - realtimeBrokerRequest = serverBrokerRequest; - - requestContext.setFanoutType(RequestContext.FanoutType.REALTIME); - requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableConfig)); - } - - // Check if response can be sent without server query evaluation. - if (offlineBrokerRequest != null && isFilterAlwaysFalse(offlineBrokerRequest.getPinotQuery())) { - // We don't need to evaluate offline request - offlineBrokerRequest = null; - } - if (realtimeBrokerRequest != null && isFilterAlwaysFalse(realtimeBrokerRequest.getPinotQuery())) { - // We don't need to evaluate realtime request - realtimeBrokerRequest = null; - } + Pair tableRequests = + prepareOfflineAndRealtimeRequests(routeInfo, serverPinotQuery, serverBrokerRequest, schema, requestContext); + BrokerRequest offlineBrokerRequest = tableRequests.getLeft(); + BrokerRequest realtimeBrokerRequest = tableRequests.getRight(); if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { return getEmptyBrokerOnlyResponse(pinotQuery, serverPinotQuery, requestContext, tableName, requesterIdentity, schema, query, database, workloadName); } - if (offlineBrokerRequest != null && isFilterAlwaysTrue(offlineBrokerRequest.getPinotQuery())) { - // Drop offline request filter since it is always true - offlineBrokerRequest.getPinotQuery().setFilterExpression(null); - } - if (realtimeBrokerRequest != null && isFilterAlwaysTrue(realtimeBrokerRequest.getPinotQuery())) { - // Drop realtime request filter since it is always true - realtimeBrokerRequest.getPinotQuery().setFilterExpression(null); - } - // Calculate routing table for the query // TODO: Modify RoutingManager interface to directly take PinotQuery long routingStartTimeNs = System.nanoTime(); @@ -629,7 +563,6 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn Set offlineExecutionServers = routeInfo.getOfflineExecutionServers(); Set realtimeExecutionServers = routeInfo.getRealtimeExecutionServers(); - List unavailableSegments = routeInfo.getUnavailableSegments(); int numPrunedSegmentsTotal = routeInfo.getNumPrunedSegmentsTotal(); // Rewrite the broker requests as the rest of the code expects them to be null or not based on whether the routing @@ -637,8 +570,6 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn offlineBrokerRequest = routeInfo.getOfflineBrokerRequest(); realtimeBrokerRequest = routeInfo.getRealtimeBrokerRequest(); - List errorMsgs = new ArrayList<>(); - // If all tables in a hybrid are disabled then return an error. // Note that if a query is for one of OFFLINE or REALTIME table and the other physical table is disabled, // we still want to run the query. So this condition is false. @@ -648,33 +579,9 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn return BrokerResponseNative.TABLE_IS_DISABLED; } - List disabledTableNames = routeInfo.getDisabledTableNames(); - if (disabledTableNames != null) { - for (String name : disabledTableNames) { - String errorMessage = String.format("%s Table is disabled", name); - LOGGER.info("{}: {}", errorMessage, query); - errorMsgs.add(new QueryProcessingException(QueryErrorCode.TABLE_IS_DISABLED, errorMessage)); - } - } - - int numUnavailableSegments = unavailableSegments.size(); - requestContext.setNumUnavailableSegments(numUnavailableSegments); - - if (numUnavailableSegments > 0) { - String errorMessage; - if (numUnavailableSegments > MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION) { - errorMessage = String.format("%d segments unavailable, sampling %d: %s", numUnavailableSegments, - MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION, - unavailableSegments.subList(0, MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION)); - } else { - errorMessage = String.format("%d segments unavailable: %s", numUnavailableSegments, unavailableSegments); - } - String realtimeRoutingPolicy = realtimeBrokerRequest != null ? getRoutingPolicy(realtimeTableConfig) : null; - String offlineRoutingPolicy = offlineBrokerRequest != null ? getRoutingPolicy(offlineTableConfig) : null; - errorMessage = addRoutingPolicyInErrMsg(errorMessage, realtimeRoutingPolicy, offlineRoutingPolicy); - errorMsgs.add(new QueryProcessingException(QueryErrorCode.BROKER_SEGMENT_UNAVAILABLE, errorMessage)); - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_UNAVAILABLE_SEGMENTS, 1); - } + List errorMsgs = + collectRoutingWarnings(routeInfo, offlineBrokerRequest, realtimeBrokerRequest, rawTableName, query, + requestContext); if (offlineBrokerRequest == null && realtimeBrokerRequest == null) { if (!errorMsgs.isEmpty()) { @@ -698,26 +605,16 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn QueryThreadContext.checkTerminationAndSampleUsage("Broker request routing"); // Set timeout in the requests - long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); - // Remaining time in milliseconds for the server query execution // NOTE: For hybrid use case, in most cases offline table and real-time table should have the same query timeout // configured, but if necessary, we also allow different timeout for them. // If the timeout is not the same for offline table and real-time table, use the max of offline table // remaining time and realtime table remaining time. Server side will have different remaining time set for // each table type, and broker should wait for both types to return. - long remainingTimeMs = 0; + long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); + long remainingTimeMs; try { - Long logicalTableQueryTimeout = logicalTableConfig != null && logicalTableConfig.getQueryConfig() != null - && logicalTableConfig.getQueryConfig().getTimeoutMs() != null ? logicalTableConfig.getQueryConfig() - .getTimeoutMs() : null; - if (offlineBrokerRequest != null) { - remainingTimeMs = setQueryTimeout(offlineTableName, logicalTableQueryTimeout, - offlineBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs); - } - if (realtimeBrokerRequest != null) { - remainingTimeMs = Math.max(remainingTimeMs, setQueryTimeout(realtimeTableName, logicalTableQueryTimeout, - realtimeBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs)); - } + remainingTimeMs = configureQueryTimeouts(offlineBrokerRequest, realtimeBrokerRequest, + routeInfo.getOfflineTableName(), routeInfo.getRealtimeTableName(), logicalTableConfig, timeSpentMs); } catch (TimeoutException e) { String errorMessage = e.getMessage(); LOGGER.info("{} {}: {}", errorMessage, requestId, query); @@ -728,39 +625,8 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn // Set the maximum serialized response size per server, and ask server to directly return final response when only // one server is queried - int numServers = 0; - if (offlineExecutionServers != null) { - numServers += offlineExecutionServers.size(); - } - if (realtimeExecutionServers != null) { - numServers += realtimeExecutionServers.size(); - } - if (offlineBrokerRequest != null) { - Map queryOptions = offlineBrokerRequest.getPinotQuery().getQueryOptions(); - setMaxServerResponseSizeBytes(numServers, queryOptions, offlineTableQueryConfig); - // Set the query option to directly return final result for single server query unless it is explicitly disabled - if (numServers == 1) { - // Set the same flag in the original server request to be used in the reduce phase for hybrid table - if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null - && offlineBrokerRequest != serverBrokerRequest) { - serverBrokerRequest.getPinotQuery().getQueryOptions() - .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); - } - } - } - if (realtimeBrokerRequest != null) { - Map queryOptions = realtimeBrokerRequest.getPinotQuery().getQueryOptions(); - setMaxServerResponseSizeBytes(numServers, queryOptions, realtimeTableQueryConfig); - // Set the query option to directly return final result for single server query unless it is explicitly disabled - if (numServers == 1) { - // Set the same flag in the original server request to be used in the reduce phase for hybrid table - if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null - && realtimeBrokerRequest != serverBrokerRequest) { - serverBrokerRequest.getPinotQuery().getQueryOptions() - .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); - } - } - } + configureServerResponseOptions(offlineBrokerRequest, realtimeBrokerRequest, offlineExecutionServers, + realtimeExecutionServers, offlineTableQueryConfig, realtimeTableQueryConfig, serverBrokerRequest); // Execute the query // TODO: Replace ServerStats with ServerRoutingStatsEntry. @@ -804,73 +670,12 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, routeInfo, remainingTimeMs, serverStats, requestContext); } - brokerResponse.setTablesQueried(Set.of(rawTableName)); - brokerResponse.setReplicaGroups(Stream.concat( - offlineExecutionServers != null ? offlineExecutionServers.stream() : Stream.empty(), - realtimeExecutionServers != null ? realtimeExecutionServers.stream() : Stream.empty()) - .map(ServerInstance::getPool) - .collect(Collectors.toSet())); - - for (QueryProcessingException errorMsg : errorMsgs) { - brokerResponse.addException(errorMsg); - } - brokerResponse.setNumSegmentsPrunedByBroker(numPrunedSegmentsTotal); long executionEndTimeNs = System.nanoTime(); - _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_EXECUTION, - executionEndTimeNs - routingEndTimeNs); - - // Track number of queries with number of groups limit reached - if (brokerResponse.isNumGroupsLimitReached()) { - _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_NUM_GROUPS_LIMIT_REACHED, - 1); - } - // server returns STRING as default dataType for all columns in (some) scenarios where no rows are returned - // this is an attempt to return more faithful information based on other sources - fillEmptyResponseSchema(pinotQuery, brokerResponse, schema, database, query); - - // Set total query processing time - long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); - brokerResponse.setTimeUsedMs(totalTimeMs); - augmentStatistics(requestContext, brokerResponse); - // include both broker side errorMsgs and server side errorMsgs - List brokerExceptions = brokerResponse.getExceptions(); - brokerExceptions.stream() - .filter(exception -> exception.getErrorCode() == QueryErrorCode.QUERY_VALIDATION.getId()) - .findFirst() - .ifPresent(exception -> _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1)); - if (!pinotQuery.isExplain() && QueryOptionsUtils.shouldDropResults(pinotQuery.getQueryOptions())) { - brokerResponse.setResultTable(null); - } - if (QueryOptionsUtils.isSecondaryWorkload(pinotQuery.getQueryOptions())) { - _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.SECONDARY_WORKLOAD_QUERY_TOTAL_TIME_MS, totalTimeMs, - TimeUnit.MILLISECONDS); - _brokerMetrics.addTimedValue(BrokerTimer.SECONDARY_WORKLOAD_QUERY_TOTAL_TIME_MS, totalTimeMs, - TimeUnit.MILLISECONDS); - } else { - _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, - TimeUnit.MILLISECONDS); - _brokerMetrics.addTimedValue(BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, TimeUnit.MILLISECONDS); - } - _brokerMetrics.addTimedValue(workloadName, BrokerTimer.WORKLOAD_TOTAL_QUERY_TIME_MS, totalTimeMs, - TimeUnit.MILLISECONDS, ImmutableMap.of(MetricAttributeConstants.WORKLOAD_NAME, workloadName)); - - for (int group : brokerResponse.getReplicaGroups()) { - String replicaGroupTag = BrokerMetrics.getTagForPreferredGroup(sqlNodeAndOptions.getOptions()); - String replicaGroupId = String.valueOf(group); - - List tags = ImmutableList.of(replicaGroupTag, replicaGroupId); - Map attributes = ImmutableMap.of( - MetricAttributeConstants.REPLICA_GROUP_TAG, replicaGroupTag, - MetricAttributeConstants.REPLICA_GROUP_ID, replicaGroupId - ); - _brokerMetrics.addMeteredValue(BrokerMeter.REPLICA_QUERIES, 1, tags, attributes); - } - - // Log query and stats - _queryLogger.log( - new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, - QueryLogger.QueryLogParams.QueryEngine.SINGLE_STAGE, requesterIdentity, serverStats, workloadName)); + // Populate response metadata, emit metrics, and log query + finalizeResponse(brokerResponse, errorMsgs, offlineExecutionServers, realtimeExecutionServers, + numPrunedSegmentsTotal, pinotQuery, schema, database, query, rawTableName, workloadName, tableName, + requestContext, requesterIdentity, serverStats, sqlNodeAndOptions, routingEndTimeNs, executionEndTimeNs); return brokerResponse; } @@ -1011,6 +816,276 @@ private CompileResult compileRequest(long requestId, String query, SqlNodeAndOpt return new CompileResult(pinotQuery, serverPinotQuery, schema, tableName, rawTableName); } + /** + * Applies per-table transformations (time boundary, expression override, timestamp index override, query optimizer) + * and prunes requests whose filter is always false. Returns a pair of (offlineBrokerRequest, realtimeBrokerRequest); + * either entry may be null when pruned. + */ + private Pair prepareOfflineAndRealtimeRequests(TableRouteInfo routeInfo, + PinotQuery serverPinotQuery, BrokerRequest serverBrokerRequest, @Nullable Schema schema, + RequestContext requestContext) { + String offlineTableName = routeInfo.getOfflineTableName(); + String realtimeTableName = routeInfo.getRealtimeTableName(); + TableConfig offlineTableConfig = routeInfo.getOfflineTableConfig(); + TableConfig realtimeTableConfig = routeInfo.getRealtimeTableConfig(); + TimeBoundaryInfo timeBoundaryInfo = routeInfo.getTimeBoundaryInfo(); + + BrokerRequest offlineBrokerRequest = null; + BrokerRequest realtimeBrokerRequest = null; + + if (routeInfo.isHybrid()) { + // Hybrid + PinotQuery offlinePinotQuery = serverPinotQuery.deepCopy(); + offlinePinotQuery.getDataSource().setTableName(offlineTableName); + assert timeBoundaryInfo != null; + attachTimeBoundary(offlinePinotQuery, timeBoundaryInfo, true); + handleExpressionOverride(offlinePinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); + handleTimestampIndexOverride(offlinePinotQuery, offlineTableConfig); + _queryOptimizer.optimize(offlinePinotQuery, offlineTableConfig, schema); + offlineBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(offlinePinotQuery); + + PinotQuery realtimePinotQuery = serverPinotQuery.deepCopy(); + realtimePinotQuery.getDataSource().setTableName(realtimeTableName); + attachTimeBoundary(realtimePinotQuery, timeBoundaryInfo, false); + handleExpressionOverride(realtimePinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); + handleTimestampIndexOverride(realtimePinotQuery, realtimeTableConfig); + _queryOptimizer.optimize(realtimePinotQuery, realtimeTableConfig, schema); + realtimeBrokerRequest = CalciteSqlCompiler.convertToBrokerRequest(realtimePinotQuery); + + requestContext.setFanoutType(RequestContext.FanoutType.HYBRID); + requestContext.setOfflineServerTenant(getServerTenant(offlineTableConfig)); + requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableConfig)); + } else if (routeInfo.isOffline()) { + // OFFLINE only + setTableName(serverBrokerRequest, offlineTableName); + handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(offlineTableName)); + handleTimestampIndexOverride(serverPinotQuery, offlineTableConfig); + _queryOptimizer.optimize(serverPinotQuery, offlineTableConfig, schema); + offlineBrokerRequest = serverBrokerRequest; + + requestContext.setFanoutType(RequestContext.FanoutType.OFFLINE); + requestContext.setOfflineServerTenant(getServerTenant(offlineTableConfig)); + } else { + // REALTIME only + setTableName(serverBrokerRequest, realtimeTableName); + handleExpressionOverride(serverPinotQuery, _tableCache.getExpressionOverrideMap(realtimeTableName)); + handleTimestampIndexOverride(serverPinotQuery, realtimeTableConfig); + _queryOptimizer.optimize(serverPinotQuery, realtimeTableConfig, schema); + realtimeBrokerRequest = serverBrokerRequest; + + requestContext.setFanoutType(RequestContext.FanoutType.REALTIME); + requestContext.setRealtimeServerTenant(getServerTenant(realtimeTableConfig)); + } + + // Prune requests whose filter is always false + if (offlineBrokerRequest != null && isFilterAlwaysFalse(offlineBrokerRequest.getPinotQuery())) { + offlineBrokerRequest = null; + } + if (realtimeBrokerRequest != null && isFilterAlwaysFalse(realtimeBrokerRequest.getPinotQuery())) { + realtimeBrokerRequest = null; + } + + // Drop always-true filters (no-op but avoids sending them to servers) + if (offlineBrokerRequest != null && isFilterAlwaysTrue(offlineBrokerRequest.getPinotQuery())) { + offlineBrokerRequest.getPinotQuery().setFilterExpression(null); + } + if (realtimeBrokerRequest != null && isFilterAlwaysTrue(realtimeBrokerRequest.getPinotQuery())) { + realtimeBrokerRequest.getPinotQuery().setFilterExpression(null); + } + + return Pair.of(offlineBrokerRequest, realtimeBrokerRequest); + } + + /** + * Builds the list of non-fatal routing warnings (partially-disabled tables, unavailable segments) after routing + * has been calculated. Also sets the unavailable segment count on {@code requestContext}. + */ + private List collectRoutingWarnings(TableRouteInfo routeInfo, + @Nullable BrokerRequest offlineBrokerRequest, @Nullable BrokerRequest realtimeBrokerRequest, + String rawTableName, String query, RequestContext requestContext) { + List errorMsgs = new ArrayList<>(); + + List disabledTableNames = routeInfo.getDisabledTableNames(); + if (disabledTableNames != null) { + for (String name : disabledTableNames) { + String errorMessage = String.format("%s Table is disabled", name); + LOGGER.info("{}: {}", errorMessage, query); + errorMsgs.add(new QueryProcessingException(QueryErrorCode.TABLE_IS_DISABLED, errorMessage)); + } + } + + List unavailableSegments = routeInfo.getUnavailableSegments(); + int numUnavailableSegments = unavailableSegments.size(); + requestContext.setNumUnavailableSegments(numUnavailableSegments); + + if (numUnavailableSegments > 0) { + String errorMessage; + if (numUnavailableSegments > MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION) { + errorMessage = String.format("%d segments unavailable, sampling %d: %s", numUnavailableSegments, + MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION, + unavailableSegments.subList(0, MAX_UNAVAILABLE_SEGMENTS_TO_PRINT_IN_QUERY_EXCEPTION)); + } else { + errorMessage = String.format("%d segments unavailable: %s", numUnavailableSegments, unavailableSegments); + } + String realtimeRoutingPolicy = + realtimeBrokerRequest != null ? getRoutingPolicy(routeInfo.getRealtimeTableConfig()) : null; + String offlineRoutingPolicy = + offlineBrokerRequest != null ? getRoutingPolicy(routeInfo.getOfflineTableConfig()) : null; + errorMessage = addRoutingPolicyInErrMsg(errorMessage, realtimeRoutingPolicy, offlineRoutingPolicy); + errorMsgs.add(new QueryProcessingException(QueryErrorCode.BROKER_SEGMENT_UNAVAILABLE, errorMessage)); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_UNAVAILABLE_SEGMENTS, 1); + } + + return errorMsgs; + } + + /** + * Sets the query timeout on the offline and/or realtime broker requests and returns the remaining time in + * milliseconds for server execution. + * + * @throws TimeoutException when the query has already exceeded its timeout budget + */ + private long configureQueryTimeouts(@Nullable BrokerRequest offlineBrokerRequest, + @Nullable BrokerRequest realtimeBrokerRequest, @Nullable String offlineTableName, + @Nullable String realtimeTableName, @Nullable LogicalTableConfig logicalTableConfig, long timeSpentMs) + throws TimeoutException { + Long logicalTableQueryTimeout = logicalTableConfig != null && logicalTableConfig.getQueryConfig() != null + && logicalTableConfig.getQueryConfig().getTimeoutMs() != null + ? logicalTableConfig.getQueryConfig().getTimeoutMs() : null; + long remainingTimeMs = 0; + if (offlineBrokerRequest != null) { + remainingTimeMs = setQueryTimeout(offlineTableName, logicalTableQueryTimeout, + offlineBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs); + } + if (realtimeBrokerRequest != null) { + remainingTimeMs = Math.max(remainingTimeMs, setQueryTimeout(realtimeTableName, logicalTableQueryTimeout, + realtimeBrokerRequest.getPinotQuery().getQueryOptions(), timeSpentMs)); + } + return remainingTimeMs; + } + + /** + * Sets the maximum serialized response size per server and, when only a single server is queried, opts into the + * single-server final-result shortcut. + */ + private void configureServerResponseOptions(@Nullable BrokerRequest offlineBrokerRequest, + @Nullable BrokerRequest realtimeBrokerRequest, @Nullable Set offlineExecutionServers, + @Nullable Set realtimeExecutionServers, @Nullable QueryConfig offlineTableQueryConfig, + @Nullable QueryConfig realtimeTableQueryConfig, BrokerRequest serverBrokerRequest) { + int numServers = 0; + if (offlineExecutionServers != null) { + numServers += offlineExecutionServers.size(); + } + if (realtimeExecutionServers != null) { + numServers += realtimeExecutionServers.size(); + } + if (offlineBrokerRequest != null) { + Map queryOptions = offlineBrokerRequest.getPinotQuery().getQueryOptions(); + setMaxServerResponseSizeBytes(numServers, queryOptions, offlineTableQueryConfig); + // Set the query option to directly return final result for single server query unless it is explicitly disabled + if (numServers == 1) { + // Set the same flag in the original server request to be used in the reduce phase for hybrid table + if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null + && offlineBrokerRequest != serverBrokerRequest) { + serverBrokerRequest.getPinotQuery().getQueryOptions() + .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); + } + } + } + if (realtimeBrokerRequest != null) { + Map queryOptions = realtimeBrokerRequest.getPinotQuery().getQueryOptions(); + setMaxServerResponseSizeBytes(numServers, queryOptions, realtimeTableQueryConfig); + // Set the query option to directly return final result for single server query unless it is explicitly disabled + if (numServers == 1) { + // Set the same flag in the original server request to be used in the reduce phase for hybrid table + if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null + && realtimeBrokerRequest != serverBrokerRequest) { + serverBrokerRequest.getPinotQuery().getQueryOptions() + .put(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true"); + } + } + } + } + + /** + * Populates metadata on {@code brokerResponse}, emits query-level metrics, and logs the query. Must be called + * immediately after query execution; {@code executionEndTimeNs} should be {@link System#nanoTime()} captured right + * after {@code processBrokerRequest} returns. + */ + private void finalizeResponse(BrokerResponseNative brokerResponse, List errorMsgs, + @Nullable Set offlineExecutionServers, @Nullable Set realtimeExecutionServers, + int numPrunedSegmentsTotal, PinotQuery pinotQuery, @Nullable Schema schema, String database, String query, + String rawTableName, String workloadName, String tableName, RequestContext requestContext, + @Nullable RequesterIdentity requesterIdentity, @Nullable ServerStats serverStats, + SqlNodeAndOptions sqlNodeAndOptions, long routingEndTimeNs, long executionEndTimeNs) { + brokerResponse.setTablesQueried(Set.of(rawTableName)); + brokerResponse.setReplicaGroups(Stream.concat( + offlineExecutionServers != null ? offlineExecutionServers.stream() : Stream.empty(), + realtimeExecutionServers != null ? realtimeExecutionServers.stream() : Stream.empty()) + .map(ServerInstance::getPool) + .collect(Collectors.toSet())); + + for (QueryProcessingException errorMsg : errorMsgs) { + brokerResponse.addException(errorMsg); + } + brokerResponse.setNumSegmentsPrunedByBroker(numPrunedSegmentsTotal); + _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_EXECUTION, + executionEndTimeNs - routingEndTimeNs); + + // Track number of queries with number of groups limit reached + if (brokerResponse.isNumGroupsLimitReached()) { + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_NUM_GROUPS_LIMIT_REACHED, + 1); + } + + // server returns STRING as default dataType for all columns in (some) scenarios where no rows are returned + // this is an attempt to return more faithful information based on other sources + fillEmptyResponseSchema(pinotQuery, brokerResponse, schema, database, query); + + // Set total query processing time + long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); + brokerResponse.setTimeUsedMs(totalTimeMs); + augmentStatistics(requestContext, brokerResponse); + // include both broker side errorMsgs and server side errorMsgs + List brokerExceptions = brokerResponse.getExceptions(); + brokerExceptions.stream() + .filter(exception -> exception.getErrorCode() == QueryErrorCode.QUERY_VALIDATION.getId()) + .findFirst() + .ifPresent(exception -> _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1)); + if (!pinotQuery.isExplain() && QueryOptionsUtils.shouldDropResults(pinotQuery.getQueryOptions())) { + brokerResponse.setResultTable(null); + } + if (QueryOptionsUtils.isSecondaryWorkload(pinotQuery.getQueryOptions())) { + _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.SECONDARY_WORKLOAD_QUERY_TOTAL_TIME_MS, totalTimeMs, + TimeUnit.MILLISECONDS); + _brokerMetrics.addTimedValue(BrokerTimer.SECONDARY_WORKLOAD_QUERY_TOTAL_TIME_MS, totalTimeMs, + TimeUnit.MILLISECONDS); + } else { + _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, + TimeUnit.MILLISECONDS); + _brokerMetrics.addTimedValue(BrokerTimer.QUERY_TOTAL_TIME_MS, totalTimeMs, TimeUnit.MILLISECONDS); + } + _brokerMetrics.addTimedValue(workloadName, BrokerTimer.WORKLOAD_TOTAL_QUERY_TIME_MS, totalTimeMs, + TimeUnit.MILLISECONDS, ImmutableMap.of(MetricAttributeConstants.WORKLOAD_NAME, workloadName)); + + for (int group : brokerResponse.getReplicaGroups()) { + String replicaGroupTag = BrokerMetrics.getTagForPreferredGroup(sqlNodeAndOptions.getOptions()); + String replicaGroupId = String.valueOf(group); + + List tags = ImmutableList.of(replicaGroupTag, replicaGroupId); + Map attributes = ImmutableMap.of( + MetricAttributeConstants.REPLICA_GROUP_TAG, replicaGroupTag, + MetricAttributeConstants.REPLICA_GROUP_ID, replicaGroupId + ); + _brokerMetrics.addMeteredValue(BrokerMeter.REPLICA_QUERIES, 1, tags, attributes); + } + + // Log query and stats + _queryLogger.log( + new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, + QueryLogger.QueryLogParams.QueryEngine.SINGLE_STAGE, requesterIdentity, serverStats, workloadName)); + } + private void throwAccessDeniedError(long requestId, String query, RequestContext requestContext, String tableName, AuthorizationResult authorizationResult) { _brokerMetrics.addMeteredTableValue(tableName, BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1);