Multi stage stats

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java

@@ -761,7 +761,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
           realtimeRoutingTable = null;
         }
       }
-      BrokerResponseNative brokerResponse;
+      BrokerResponse brokerResponse;
       if (_queriesById != null) {
         // Start to track the running query for cancellation just before sending it out to servers to avoid any
         // potential failures that could happen before sending it out, like failures to calculate the routing table etc.
@@ -798,7 +798,6 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
         _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.BROKER_RESPONSES_WITH_NUM_GROUPS_LIMIT_REACHED,
             1);
       }
-      brokerResponse.setPartialResult(isPartialResult(brokerResponse));
 
       // Set total query processing time
       long totalTimeMs = TimeUnit.NANOSECONDS.toMillis(executionEndTimeNs - compilationStartTimeNs);
@@ -1870,19 +1869,14 @@ private static void attachTimeBoundary(PinotQuery pinotQuery, TimeBoundaryInfo t
    * Processes the optimized broker requests for both OFFLINE and REALTIME table.
    * TODO: Directly take PinotQuery
    */
-  protected abstract BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
+  protected abstract BrokerResponse processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
       BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest,
       @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable,
       @Nullable BrokerRequest realtimeBrokerRequest,
       @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> realtimeRoutingTable, long timeoutMs,
       ServerStats serverStats, RequestContext requestContext)
       throws Exception;
 
-  protected static boolean isPartialResult(BrokerResponse brokerResponse) {
-    return brokerResponse.isNumGroupsLimitReached() || brokerResponse.isMaxRowsInJoinReached()
-        || brokerResponse.getExceptionsSize() > 0;
-  }
-
   protected static void augmentStatistics(RequestContext statistics, BrokerResponse response) {
     statistics.setTotalDocs(response.getTotalDocs());
     statistics.setNumDocsScanned(response.getNumDocsScanned());

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

@@ -19,14 +19,13 @@
 package org.apache.pinot.broker.requesthandler;
 
 import com.fasterxml.jackson.databind.JsonNode;
-import com.google.common.collect.Maps;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.HttpHeaders;
@@ -48,7 +47,6 @@
 import org.apache.pinot.common.response.BrokerResponse;
 import org.apache.pinot.common.response.broker.BrokerResponseNative;
 import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
-import org.apache.pinot.common.response.broker.BrokerResponseStats;
 import org.apache.pinot.common.response.broker.QueryProcessingException;
 import org.apache.pinot.common.response.broker.ResultTable;
 import org.apache.pinot.common.utils.DataSchema;
@@ -58,14 +56,17 @@
 import org.apache.pinot.common.utils.request.RequestUtils;
 import org.apache.pinot.core.auth.Actions;
 import org.apache.pinot.core.auth.TargetType;
-import org.apache.pinot.core.query.reduce.ExecutionStatsAggregator;
 import org.apache.pinot.core.transport.ServerInstance;
 import org.apache.pinot.query.QueryEnvironment;
 import org.apache.pinot.query.catalog.PinotCatalog;
 import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.planner.PlanFragment;
 import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
 import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
+import org.apache.pinot.query.planner.plannode.PlanNode;
 import org.apache.pinot.query.routing.WorkerManager;
+import org.apache.pinot.query.runtime.MultiStageStatsTreeBuilder;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
 import org.apache.pinot.query.service.dispatch.QueryDispatcher;
 import org.apache.pinot.query.type.TypeFactory;
 import org.apache.pinot.query.type.TypeSystem;
@@ -195,24 +196,12 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
     }
 
     Map<String, String> queryOptions = sqlNodeAndOptions.getOptions();
-    boolean traceEnabled = Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.TRACE));
-    Map<Integer, ExecutionStatsAggregator> stageIdStatsMap;
-    if (!traceEnabled) {
-      stageIdStatsMap = Collections.singletonMap(0, new ExecutionStatsAggregator(false));
-    } else {
-      List<DispatchablePlanFragment> stagePlans = dispatchableSubPlan.getQueryStageList();
-      int numStages = stagePlans.size();
-      stageIdStatsMap = Maps.newHashMapWithExpectedSize(numStages);
-      for (int stageId = 0; stageId < numStages; stageId++) {
-        stageIdStatsMap.put(stageId, new ExecutionStatsAggregator(true));
-      }
-    }
 
     long executionStartTimeNs = System.nanoTime();
-    ResultTable queryResults;
+    QueryDispatcher.QueryResult queryResults;
     try {
-      queryResults = _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions,
-          stageIdStatsMap);
+      queryResults =
+          _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions);
     } catch (TimeoutException e) {
       for (String table : tableNames) {
         _brokerMetrics.addMeteredTableValue(table, BrokerMeter.BROKER_RESPONSES_WITH_TIMEOUTS, 1);
@@ -231,7 +220,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
     updatePhaseTimingForTables(tableNames, BrokerQueryPhase.QUERY_EXECUTION, executionEndTimeNs - executionStartTimeNs);
 
     BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2();
-    brokerResponse.setResultTable(queryResults);
+    brokerResponse.setResultTable(queryResults.getResultTable());
 
     // Attach unavailable segments
     int numUnavailableSegments = 0;
@@ -243,35 +232,14 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
           String.format("Find unavailable segments: %s for table: %s", unavailableSegments, tableName)));
     }
 
-    for (Map.Entry<Integer, ExecutionStatsAggregator> entry : stageIdStatsMap.entrySet()) {
-      if (entry.getKey() == 0) {
-        // Root stats are aggregated and added separately to broker response for backward compatibility
-        entry.getValue().setStats(brokerResponse);
-        continue;
-      }
-
-      BrokerResponseStats brokerResponseStats = new BrokerResponseStats();
-      if (!tableNames.isEmpty()) {
-        //TODO: Only using first table to assign broker metrics
-        // find a way to split metrics in case of multiple table
-        String rawTableName = TableNameBuilder.extractRawTableName(tableNames.iterator().next());
-        entry.getValue().setStageLevelStats(rawTableName, brokerResponseStats, _brokerMetrics);
-      } else {
-        entry.getValue().setStageLevelStats(null, brokerResponseStats, null);
-      }
-      brokerResponse.addStageStat(entry.getKey(), brokerResponseStats);
-    }
-
-    // Set partial result flag
-    brokerResponse.setPartialResult(isPartialResult(brokerResponse));
+    fillOldBrokerResponseStats(brokerResponse, queryResults.getQueryStats(), dispatchableSubPlan);
 
     // Set total query processing time
     // TODO: Currently we don't emit metric for QUERY_TOTAL_TIME_MS
     long totalTimeMs = TimeUnit.NANOSECONDS.toMillis(
         sqlNodeAndOptions.getParseTimeNs() + (executionEndTimeNs - compilationStartTimeNs));
     brokerResponse.setTimeUsedMs(totalTimeMs);
     requestContext.setQueryProcessingTime(totalTimeMs);
-    requestContext.setTraceInfo(brokerResponse.getTraceInfo());
     augmentStatistics(requestContext, brokerResponse);
 
     // Log query and stats
@@ -282,6 +250,22 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
     return brokerResponse;
   }
 
+  private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse,
+      List<MultiStageQueryStats.StageStats.Closed> queryStats, DispatchableSubPlan dispatchableSubPlan) {
+    List<PlanNode> planNodes = dispatchableSubPlan.getQueryStageList().stream()
+            .map(DispatchablePlanFragment::getPlanFragment)
+            .map(PlanFragment::getFragmentRoot)
+            .collect(Collectors.toList());
+    MultiStageStatsTreeBuilder treeBuilder = new MultiStageStatsTreeBuilder(planNodes, queryStats);
+    brokerResponse.setStageStats(treeBuilder.jsonStatsByStage(0));
+
+    for (MultiStageQueryStats.StageStats.Closed stageStats : queryStats) {
+      if (stageStats != null) { // for example pipeline breaker may not have stats
+        stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats));
+      }
+    }
+  }
+
   /**
    * Validates whether the requester has access to all the tables.
    */
@@ -325,7 +309,7 @@ private void updatePhaseTimingForTables(Set<String> tableNames, BrokerQueryPhase
     }
   }
 
-  private BrokerResponseNative constructMultistageExplainPlan(String sql, String plan) {
+  private BrokerResponse constructMultistageExplainPlan(String sql, String plan) {
     BrokerResponseNative brokerResponse = BrokerResponseNative.empty();
     List<Object[]> rows = new ArrayList<>();
     rows.add(new Object[]{sql, plan});
@@ -336,7 +320,7 @@ private BrokerResponseNative constructMultistageExplainPlan(String sql, String p
   }
 
   @Override
-  protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
+  protected BrokerResponse processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
       BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest,
       @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable,
       @Nullable BrokerRequest realtimeBrokerRequest,

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java

@@ -40,6 +40,7 @@
 import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.common.metrics.BrokerQueryPhase;
 import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.response.BrokerResponse;
 import org.apache.pinot.common.response.broker.BrokerResponseNative;
 import org.apache.pinot.common.response.broker.QueryProcessingException;
 import org.apache.pinot.common.utils.HashUtil;
@@ -100,7 +101,7 @@ public synchronized void shutDown() {
   }
 
   @Override
-  protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
+  protected BrokerResponse processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
       BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest,
       @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable,
       @Nullable BrokerRequest realtimeBrokerRequest,

pinot-common/src/main/java/org/apache/pinot/common/datablock/BaseDataBlock.java

@@ -93,7 +93,6 @@ public abstract class BaseDataBlock implements DataBlock {
   protected ByteBuffer _fixedSizeData;
   protected byte[] _variableSizeDataBytes;
   protected ByteBuffer _variableSizeData;
-  protected Map<String, String> _metadata;
 
   /**
    * construct a base data block.
@@ -114,7 +113,6 @@ public BaseDataBlock(int numRows, @Nullable DataSchema dataSchema, String[] stri
     _fixedSizeData = ByteBuffer.wrap(fixedSizeDataBytes);
     _variableSizeDataBytes = variableSizeDataBytes;
     _variableSizeData = ByteBuffer.wrap(variableSizeDataBytes);
-    _metadata = new HashMap<>();
     _errCodeToExceptionMap = new HashMap<>();
   }
 
@@ -131,7 +129,6 @@ public BaseDataBlock() {
     _fixedSizeData = null;
     _variableSizeDataBytes = null;
     _variableSizeData = null;
-    _metadata = new HashMap<>();
     _errCodeToExceptionMap = new HashMap<>();
   }
 
@@ -195,10 +192,7 @@ public BaseDataBlock(ByteBuffer byteBuffer)
     _variableSizeData = ByteBuffer.wrap(_variableSizeDataBytes);
 
     // Read metadata.
-    int metadataLength = byteBuffer.getInt();
-    if (metadataLength != 0) {
-      _metadata = deserializeMetadata(byteBuffer);
-    }
+    deserializeMetadata(byteBuffer);
   }
 
   @Override
@@ -232,7 +226,7 @@ public int getVersion() {
 
   @Override
   public Map<String, String> getMetadata() {
-    return _metadata;
+    return Collections.emptyMap();
   }
 
   @Override
@@ -432,6 +426,11 @@ public Map<Integer, String> getExceptions() {
     return _errCodeToExceptionMap;
   }
 
+  /**
+   * Serialize this data block to a byte array.
+   * <p>
+   * In order to deserialize it, {@link DataBlockUtils#getDataBlock(ByteBuffer)} should be used.
+   */
   @Override
   public byte[] toBytes()
       throws IOException {
@@ -444,9 +443,7 @@ public byte[] toBytes()
     // Write metadata: length followed by actual metadata bytes.
     // NOTE: We ignore metadata serialization time in "responseSerializationCpuTimeNs" as it's negligible while
     // considering it will bring a lot code complexity.
-    byte[] metadataBytes = serializeMetadata();
-    dataOutputStream.writeInt(metadataBytes.length);
-    dataOutputStream.write(metadataBytes);
+    serializeMetadata(dataOutputStream);
 
     return byteArrayOutputStream.toByteArray();
   }
@@ -525,14 +522,30 @@ private void writeLeadingSections(DataOutputStream dataOutputStream)
     }
   }
 
-  private byte[] serializeMetadata()
+  /**
+   * Writes the metadata section to the given data output stream.
+   */
+  protected void serializeMetadata(DataOutputStream dataOutputStream)
       throws IOException {
-    return new byte[0];
+    dataOutputStream.writeInt(0);
   }
 
-  private Map<String, String> deserializeMetadata(ByteBuffer buffer)
+  /**
+   * Deserializes the metadata section from the given byte buffer.
+   * <p>
+   * This is the counterpart of {@link #serializeMetadata(DataOutputStream)} and it is guaranteed that the buffer will
+   * be positioned at the start of the metadata section when this method is called.
+   * <p>
+   * <strong>Important:</strong> It is mandatory for implementations to leave the cursor at the end of the metadata, in
+   * the exact same position as it was when {@link #serializeMetadata(DataOutputStream)} was called.
+   * <p>
+   * <strong>Important:</strong> This method will be called at the end of the BaseDataConstructor constructor to read
+   * the metadata section. This means that it will be called <strong>before</strong> the subclass have been constructor
+   * have been called. Therefore it is not possible to use any subclass fields in this method.
+   */
+  protected void deserializeMetadata(ByteBuffer buffer)
       throws IOException {
-    return Collections.emptyMap();
+    buffer.getInt();
   }
 
   private byte[] serializeExceptions()
@@ -572,14 +585,9 @@ private Map<Integer, String> deserializeExceptions(ByteBuffer buffer)
   @Override
   public String toString() {
     if (_dataSchema == null) {
-      return _metadata.toString();
+      return "{}";
     } else {
-      StringBuilder stringBuilder = new StringBuilder();
-      stringBuilder.append("resultSchema:").append('\n');
-      stringBuilder.append(_dataSchema).append('\n');
-      stringBuilder.append("numRows: ").append(_numRows).append('\n');
-      stringBuilder.append("metadata: ").append(_metadata.toString()).append('\n');
-      return stringBuilder.toString();
+      return "resultSchema:" + '\n' + _dataSchema + '\n' + "numRows: " + _numRows + '\n';
     }
   }
 }

pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java

@@ -20,14 +20,16 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
 import org.apache.pinot.common.utils.DataSchema;
 
 
 /**
  * Column-wise data table. It stores data in columnar-major format.
  */
 public class ColumnarDataBlock extends BaseDataBlock {
-  private static final int VERSION = 1;
+  private static final int VERSION = 2;
   protected int[] _cumulativeColumnOffsetSizeInBytes;
   protected int[] _columnSizeInBytes;
 
@@ -80,17 +82,21 @@ protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
   }
 
   @Override
-  public ColumnarDataBlock toMetadataOnlyDataTable() {
-    ColumnarDataBlock metadataOnlyDataTable = new ColumnarDataBlock();
-    metadataOnlyDataTable._metadata.putAll(_metadata);
-    metadataOnlyDataTable._errCodeToExceptionMap.putAll(_errCodeToExceptionMap);
-    return metadataOnlyDataTable;
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof ColumnarDataBlock)) {
+      return false;
+    }
+    ColumnarDataBlock that = (ColumnarDataBlock) o;
+    return Objects.deepEquals(_cumulativeColumnOffsetSizeInBytes, that._cumulativeColumnOffsetSizeInBytes)
+        && Objects.deepEquals(_columnSizeInBytes, that._columnSizeInBytes);
   }
 
   @Override
-  public ColumnarDataBlock toDataOnlyDataTable() {
-    return new ColumnarDataBlock(_numRows, _dataSchema, _stringDictionary, _fixedSizeDataBytes, _variableSizeDataBytes);
+  public int hashCode() {
+    return Objects.hash(Arrays.hashCode(_cumulativeColumnOffsetSizeInBytes), Arrays.hashCode(_columnSizeInBytes));
   }
-
-  // TODO: add whole-column access methods.
+// TODO: add whole-column access methods.
 }

pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java

@@ -80,10 +80,6 @@ byte[] toBytes()
   @Nullable
   RoaringBitmap getNullRowIds(int colId);
 
-  DataBlock toMetadataOnlyDataTable();
-
-  DataBlock toDataOnlyDataTable();
-
   enum Type {
     ROW(0),
     COLUMNAR(1),