Improve server query cancellation and timeout checking during execution

pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java

@@ -142,7 +142,7 @@ public void init(TableDataManagerConfig tableDataManagerConfig, String instanceI
     _isStreamSegmentDownloadUntar = tableDataManagerParams.isStreamSegmentDownloadUntar();
     if (_isStreamSegmentDownloadUntar) {
       LOGGER.info("Using streamed download-untar for segment download! "
-              + "The rate limit interval for streamed download-untar is {} ms",
+              + "The rate limit interval for streamed download-untar is {} bytes",
           _streamSegmentDownloadUntarRateLimitBytesPerSec);
     }
     int maxParallelSegmentDownloads = tableDataManagerParams.getMaxParallelSegmentDownloads();

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java

@@ -198,7 +198,10 @@ protected IntermediateResultsBlock mergeResults()
     while (numBlocksMerged < _numOperators) {
       IntermediateResultsBlock blockToMerge =
           _blockingQueue.poll(endTimeMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
-      if (blockToMerge == null) {
+      // Timeout has reached, shouldn't continue to process. `_blockingQueue.poll` will continue to return blocks even
+      // if negative timeout is provided; therefore an extra check is needed
+      boolean timeoutReached = endTimeMs < System.currentTimeMillis();
+      if (blockToMerge == null || timeoutReached) {
         // Query times out, skip merging the remaining results blocks
         LOGGER.error("Timed out while polling results block, numBlocksMerged: {} (query: {})", numBlocksMerged,
             _queryContext);

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByOrderByCombineOperator.java

@@ -40,11 +40,13 @@
 import org.apache.pinot.core.data.table.UnboundedConcurrentIndexedTable;
 import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
 import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock;
+import org.apache.pinot.core.plan.DocIdSetPlanNode;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
 import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.core.util.GroupByUtils;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -165,6 +167,8 @@ protected void processSegments(int taskIndex) {
         // Merge aggregation group-by result.
         // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
         Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
+        // Count the number of merged keys
+        long mergedKeys = 0;
         // For now, only GroupBy OrderBy query has pre-constructed intermediate records
         if (intermediateRecords == null) {
           // Merge aggregation group-by result.
@@ -181,12 +185,16 @@ protected void processSegments(int taskIndex) {
                 values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
               }
               _indexedTable.upsert(new Key(keys), new Record(values));
+              mergedKeys++;
+              checkMergePhaseInterruption(mergedKeys);
             }
           }
         } else {
           for (IntermediateRecord intermediateResult : intermediateRecords) {
             //TODO: change upsert api so that it accepts intermediateRecord directly
             _indexedTable.upsert(intermediateResult._key, intermediateResult._record);
+            mergedKeys++;
+            checkMergePhaseInterruption(mergedKeys);
           }
         }
       } finally {
@@ -197,6 +205,13 @@ protected void processSegments(int taskIndex) {
     }
   }
 
+  // Check for thread interruption, every time after merging 10_000 keys
+  private void checkMergePhaseInterruption(long mergedKeys) {
+    if (mergedKeys % DocIdSetPlanNode.MAX_DOC_PER_CALL == 0 && Thread.interrupted()) {
+      throw new EarlyTerminationException();
+    }
+  }
+
   @Override
   protected void onException(Exception e) {
     _mergedProcessingExceptions.add(QueryException.getException(QueryException.QUERY_EXECUTION_ERROR, e));