You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2024/05/03 13:38:42 UTC
(pinot) branch master updated: Multi stage stats (#12704)
This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new fb9e6c51e4 Multi stage stats (#12704)
fb9e6c51e4 is described below
commit fb9e6c51e484270c3cf31d39fb49f7f6ba82ccb7
Author: Gonzalo Ortiz Jaureguizar <go...@users.noreply.github.com>
AuthorDate: Fri May 3 15:38:35 2024 +0200
Multi stage stats (#12704)
Change the multi-stage stats to return a tree like structure with stats per operator.
---
.../requesthandler/BaseBrokerRequestHandler.java | 10 +-
.../MultiStageBrokerRequestHandler.java | 72 +--
.../SingleConnectionBrokerRequestHandler.java | 3 +-
.../pinot/common/datablock/BaseDataBlock.java | 52 +-
.../pinot/common/datablock/ColumnarDataBlock.java | 26 +-
.../apache/pinot/common/datablock/DataBlock.java | 4 -
.../pinot/common/datablock/DataBlockUtils.java | 38 +-
.../pinot/common/datablock/MetadataBlock.java | 212 ++++---
.../pinot/common/datablock/RowDataBlock.java | 31 +-
.../apache/pinot/common/datatable/DataTable.java | 3 +
.../org/apache/pinot/common/datatable/StatMap.java | 498 ++++++++++++++++
.../pinot/common/response/BrokerResponse.java | 116 +---
.../response/broker/BrokerResponseNative.java | 59 +-
.../response/broker/BrokerResponseNativeV2.java | 383 +++++++++++-
.../response/broker/BrokerResponseStats.java | 7 -
.../common/datablock/BaseDataBlockContract.java | 42 ++
.../pinot/common/datablock/MetadataBlockTest.java | 187 +++++-
.../pinot/common/datablock/RowDataBlockTest.java | 35 +-
.../pinot/common/datablock/V1MetadataBlock.java} | 51 +-
.../apache/pinot/common/datatable/StatMapTest.java | 209 +++++++
.../core/operator/InstanceResponseOperator.java | 34 +-
.../StreamingInstanceResponseOperator.java | 20 +-
.../query/reduce/ExecutionStatsAggregator.java | 28 +-
.../tests/BaseClusterIntegrationTestSet.java | 2 +-
.../tests/HybridClusterIntegrationTest.java | 3 +-
.../rel/rules/PinotJoinToDynamicBroadcastRule.java | 15 +-
.../pinot/query/mailbox/GrpcSendingMailbox.java | 36 +-
.../query/mailbox/InMemorySendingMailbox.java | 9 +-
.../apache/pinot/query/mailbox/MailboxService.java | 9 +-
.../pinot/query/mailbox/ReceivingMailbox.java | 89 ++-
.../mailbox/channel/MailboxContentObserver.java | 25 +-
.../query/runtime/InStageStatsTreeBuilder.java | 207 +++++++
.../query/runtime/MultiStageStatsTreeBuilder.java | 43 ++
.../apache/pinot/query/runtime/QueryRunner.java | 5 +-
.../query/runtime/blocks/TransferableBlock.java | 94 ++-
.../runtime/blocks/TransferableBlockUtils.java | 12 +-
.../runtime/executor/OpChainSchedulerService.java | 7 +-
.../query/runtime/operator/AggregateOperator.java | 72 ++-
.../operator/BaseMailboxReceiveOperator.java | 104 +++-
.../query/runtime/operator/FilterOperator.java | 52 +-
.../query/runtime/operator/HashJoinOperator.java | 77 ++-
.../query/runtime/operator/IntersectOperator.java | 16 +-
.../LeafStageTransferableBlockOperator.java | 303 ++++++++--
.../runtime/operator/LiteralValueOperator.java | 41 +-
.../runtime/operator/MailboxReceiveOperator.java | 11 +
.../runtime/operator/MailboxSendOperator.java | 142 ++++-
.../query/runtime/operator/MinusOperator.java | 13 +
.../query/runtime/operator/MultiStageOperator.java | 252 +++++++-
.../pinot/query/runtime/operator/OpChain.java | 7 -
.../pinot/query/runtime/operator/OpChainStats.java | 109 ----
.../query/runtime/operator/OperatorStats.java | 119 ----
.../pinot/query/runtime/operator/SetOperator.java | 53 +-
.../pinot/query/runtime/operator/SortOperator.java | 70 ++-
.../operator/SortedMailboxReceiveOperator.java | 13 +-
.../query/runtime/operator/TransformOperator.java | 47 +-
.../query/runtime/operator/UnionOperator.java | 52 +-
.../runtime/operator/WindowAggregateOperator.java | 52 +-
.../runtime/operator/exchange/BlockExchange.java | 1 +
.../utils/BlockingMultiStreamConsumer.java | 28 +-
.../runtime/operator/utils/OperatorUtils.java | 77 ---
.../query/runtime/plan/MultiStageQueryStats.java | 651 +++++++++++++++++++++
.../runtime/plan/OpChainExecutionContext.java | 15 +-
.../plan/pipeline/PipelineBreakerExecutor.java | 2 +-
.../plan/pipeline/PipelineBreakerOperator.java | 69 ++-
.../plan/pipeline/PipelineBreakerResult.java | 12 +-
.../plan/server/ServerPlanRequestContext.java | 5 +-
.../query/service/dispatch/QueryDispatcher.java | 79 +--
.../pinot/query/mailbox/MailboxServiceTest.java | 60 +-
.../runtime/blocks/TransferableBlockTest.java | 44 ++
.../runtime/blocks/TransferableBlockTestUtils.java | 45 ++
.../runtime/blocks/TransferableBlockUtilsTest.java | 2 +-
.../executor/OpChainSchedulerServiceTest.java | 7 +-
.../runtime/operator/AggregateOperatorTest.java | 45 +-
.../query/runtime/operator/FilterOperatorTest.java | 32 +-
.../runtime/operator/HashJoinOperatorTest.java | 125 ++--
.../runtime/operator/IntersectOperatorTest.java | 14 +-
.../LeafStageTransferableBlockOperatorTest.java | 12 +-
.../runtime/operator/LiteralValueOperatorTest.java | 4 +-
.../operator/MailboxReceiveOperatorTest.java | 19 +-
.../runtime/operator/MailboxSendOperatorTest.java | 19 +-
.../query/runtime/operator/MinusOperatorTest.java | 14 +-
.../pinot/query/runtime/operator/OpChainTest.java | 217 +++----
.../query/runtime/operator/OperatorTestUtil.java | 51 +-
.../query/runtime/operator/SortOperatorTest.java | 92 +--
.../operator/SortedMailboxReceiveOperatorTest.java | 22 +-
.../runtime/operator/TransformOperatorTest.java | 16 +-
.../query/runtime/operator/UnionOperatorTest.java | 8 +-
.../operator/WindowAggregateOperatorTest.java | 69 +--
.../operator/exchange/BlockExchangeTest.java | 4 +-
.../runtime/plan/MultiStageQueryStatsTest.java | 100 ++++
.../plan/pipeline/PipelineBreakerExecutorTest.java | 34 +-
.../query/runtime/queries/QueryRunnerTest.java | 6 +-
.../query/runtime/queries/QueryRunnerTestBase.java | 16 +-
.../runtime/queries/ResourceBasedQueriesTest.java | 110 ++--
.../service/dispatch/QueryDispatcherTest.java | 4 +-
.../testutils/MockDataBlockOperatorFactory.java | 4 +-
.../pinot/tools/MultistageEngineQuickStart.java | 8 +
.../org/apache/pinot/tools/RealtimeQuickStart.java | 10 +
98 files changed, 4862 insertions(+), 1571 deletions(-)
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index bfe63b1882..51316e943c 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -768,7 +768,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
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.
@@ -805,7 +805,6 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
_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);
@@ -1874,7 +1873,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
* 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,
@@ -1882,11 +1881,6 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
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());
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index e1dc00a406..f954af9021 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++ b/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.request.BrokerRequest;
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.config.QueryOptionsUtils;
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;
@@ -194,24 +195,12 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
}
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);
@@ -230,7 +219,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
updatePhaseTimingForTables(tableNames, BrokerQueryPhase.QUERY_EXECUTION, executionEndTimeNs - executionStartTimeNs);
BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2();
- brokerResponse.setResultTable(queryResults);
+ brokerResponse.setResultTable(queryResults.getResultTable());
// Attach unavailable segments
int numUnavailableSegments = 0;
@@ -242,27 +231,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
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
@@ -270,7 +239,6 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
sqlNodeAndOptions.getParseTimeNs() + (executionEndTimeNs - compilationStartTimeNs));
brokerResponse.setTimeUsedMs(totalTimeMs);
requestContext.setQueryProcessingTime(totalTimeMs);
- requestContext.setTraceInfo(brokerResponse.getTraceInfo());
augmentStatistics(requestContext, brokerResponse);
// Log query and stats
@@ -281,6 +249,22 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
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.
*/
@@ -324,7 +308,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
}
}
- 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});
@@ -335,7 +319,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
}
@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,
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java
index 68ae70f9eb..b305f060fa 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java
@@ -40,6 +40,7 @@ import org.apache.pinot.common.metrics.BrokerMeter;
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 class SingleConnectionBrokerRequestHandler extends BaseBrokerRequestHandl
}
@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,
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/BaseDataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/BaseDataBlock.java
index 8c6410b75d..333ac673b6 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/BaseDataBlock.java
+++ b/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 abstract class BaseDataBlock implements DataBlock {
_fixedSizeData = ByteBuffer.wrap(fixedSizeDataBytes);
_variableSizeDataBytes = variableSizeDataBytes;
_variableSizeData = ByteBuffer.wrap(variableSizeDataBytes);
- _metadata = new HashMap<>();
_errCodeToExceptionMap = new HashMap<>();
}
@@ -131,7 +129,6 @@ public abstract class BaseDataBlock implements DataBlock {
_fixedSizeData = null;
_variableSizeDataBytes = null;
_variableSizeData = null;
- _metadata = new HashMap<>();
_errCodeToExceptionMap = new HashMap<>();
}
@@ -195,10 +192,7 @@ public abstract class BaseDataBlock implements DataBlock {
_variableSizeData = ByteBuffer.wrap(_variableSizeDataBytes);
// Read metadata.
- int metadataLength = byteBuffer.getInt();
- if (metadataLength != 0) {
- _metadata = deserializeMetadata(byteBuffer);
- }
+ deserializeMetadata(byteBuffer);
}
@Override
@@ -232,7 +226,7 @@ public abstract class BaseDataBlock implements DataBlock {
@Override
public Map<String, String> getMetadata() {
- return _metadata;
+ return Collections.emptyMap();
}
@Override
@@ -432,6 +426,11 @@ public abstract class BaseDataBlock implements DataBlock {
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 abstract class BaseDataBlock implements DataBlock {
// 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 @@ public abstract class BaseDataBlock implements DataBlock {
}
}
- 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 @@ public abstract class BaseDataBlock implements DataBlock {
@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';
}
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
index 216f4d9d91..51b01e6b69 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/ColumnarDataBlock.java
@@ -20,6 +20,8 @@ package org.apache.pinot.common.datablock;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
import org.apache.pinot.common.utils.DataSchema;
@@ -27,7 +29,7 @@ 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 @@ public class ColumnarDataBlock extends BaseDataBlock {
}
@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.
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
index 418426b4ac..768d6cdb88 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlock.java
@@ -80,10 +80,6 @@ public interface DataBlock {
@Nullable
RoaringBitmap getNullRowIds(int colId);
- DataBlock toMetadataOnlyDataTable();
-
- DataBlock toDataOnlyDataTable();
-
enum Type {
ROW(0),
COLUMNAR(1),
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
index 27f1140328..1a67f980b1 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockUtils.java
@@ -26,15 +26,20 @@ import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.response.ProcessingException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
public final class DataBlockUtils {
+ private static final Logger LOGGER = LoggerFactory.getLogger(DataBlockUtils.class);
+
private DataBlockUtils() {
}
static final int VERSION_TYPE_SHIFT = 5;
public static MetadataBlock getErrorDataBlock(Exception e) {
+ LOGGER.info("Caught exception while processing query", e);
if (e instanceof ProcessingException) {
return getErrorDataBlock(Collections.singletonMap(((ProcessingException) e).getErrorCode(), extractErrorMsg(e)));
} else {
@@ -51,33 +56,40 @@ public final class DataBlockUtils {
}
public static MetadataBlock getErrorDataBlock(Map<Integer, String> exceptions) {
- MetadataBlock errorBlock = new MetadataBlock(MetadataBlock.MetadataBlockType.ERROR);
- for (Map.Entry<Integer, String> exception : exceptions.entrySet()) {
- errorBlock.addException(exception.getKey(), exception.getValue());
- }
- return errorBlock;
+ return MetadataBlock.newError(exceptions);
+ }
+
+ /**
+ * Reads an integer from the given byte buffer.
+ * <p>
+ * The returned integer contains both the version and the type of the data block.
+ * {@link #getVersion(int)} and {@link #getType(int)} can be used to extract the version and the type.
+ * @param byteBuffer byte buffer to read from. A single int will be read
+ */
+ public static int readVersionType(ByteBuffer byteBuffer) {
+ return byteBuffer.getInt();
}
- public static MetadataBlock getEndOfStreamDataBlock() {
- return new MetadataBlock(MetadataBlock.MetadataBlockType.EOS);
+ public static int getVersion(int versionType) {
+ return versionType & ((1 << VERSION_TYPE_SHIFT) - 1);
}
- public static MetadataBlock getEndOfStreamDataBlock(Map<String, String> stats) {
- return new MetadataBlock(MetadataBlock.MetadataBlockType.EOS, stats);
+ public static DataBlock.Type getType(int versionType) {
+ return DataBlock.Type.fromOrdinal(versionType >> VERSION_TYPE_SHIFT);
}
public static DataBlock getDataBlock(ByteBuffer byteBuffer)
throws IOException {
- int versionType = byteBuffer.getInt();
- int version = versionType & ((1 << VERSION_TYPE_SHIFT) - 1);
- DataBlock.Type type = DataBlock.Type.fromOrdinal(versionType >> VERSION_TYPE_SHIFT);
+ int versionType = readVersionType(byteBuffer);
+ int version = getVersion(versionType);
+ DataBlock.Type type = getType(versionType);
switch (type) {
case COLUMNAR:
return new ColumnarDataBlock(byteBuffer);
case ROW:
return new RowDataBlock(byteBuffer);
case METADATA:
- return new MetadataBlock(byteBuffer);
+ return MetadataBlock.deserialize(byteBuffer, version);
default:
throw new UnsupportedOperationException("Unsupported data table version: " + version + " with type: " + type);
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
index 41c2e9d0ea..164ca51577 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
@@ -18,16 +18,19 @@
*/
package org.apache.pinot.common.datablock;
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
+import java.io.DataOutputStream;
import java.io.IOException;
+import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
import java.util.Map;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -36,108 +39,116 @@ import java.util.Map;
*/
public class MetadataBlock extends BaseDataBlock {
- private static final ObjectMapper JSON = new ObjectMapper();
-
+ private static final Logger LOGGER = LoggerFactory.getLogger(MetadataBlock.class);
@VisibleForTesting
- static final int VERSION = 1;
-
- public enum MetadataBlockType {
- /**
- * Indicates that this block is the final block to be sent
- * (End Of Stream) as part of an operator chain computation.
- */
- EOS,
+ static final int VERSION = 2;
+ @Nullable
+ private List<ByteBuffer> _statsByStage;
- /**
- * An {@code ERROR} metadata block indicates that there was
- * some error during computation. To retrieve the error that
- * occurred, use {@link MetadataBlock#getExceptions()}
- */
- ERROR
+ private MetadataBlock() {
+ this(Collections.emptyList());
}
- /**
- * Used to serialize the contents of the metadata block conveniently and in
- * a backwards compatible way. Use JSON because the performance of metadata block
- * SerDe should not be a bottleneck.
- */
- @JsonIgnoreProperties(ignoreUnknown = true)
- @VisibleForTesting
- static class Contents {
-
- private String _type;
- private Map<String, String> _stats;
-
- @JsonCreator
- public Contents(@JsonProperty("type") String type, @JsonProperty("stats") Map<String, String> stats) {
- _type = type;
- _stats = stats;
- }
-
- @JsonCreator
- public Contents() {
- this(null, new HashMap<>());
- }
-
- public String getType() {
- return _type;
- }
-
- public void setType(String type) {
- _type = type;
- }
-
- public Map<String, String> getStats() {
- return _stats;
- }
+ public static MetadataBlock newEos() {
+ return new MetadataBlock();
+ }
- public void setStats(Map<String, String> stats) {
- _stats = stats;
+ public static MetadataBlock newError(Map<Integer, String> exceptions) {
+ MetadataBlock errorBlock = new MetadataBlock();
+ for (Map.Entry<Integer, String> exception : exceptions.entrySet()) {
+ errorBlock.addException(exception.getKey(), exception.getValue());
}
+ return errorBlock;
}
- private final Contents _contents;
+ public MetadataBlock(List<ByteBuffer> statsByStage) {
+ super(0, null, new String[0], new byte[0], new byte[0]);
+ _statsByStage = statsByStage;
+ }
- public MetadataBlock(MetadataBlockType type) {
- this(type, new HashMap<>());
+ MetadataBlock(ByteBuffer byteBuffer)
+ throws IOException {
+ super(byteBuffer);
}
- public MetadataBlock(MetadataBlockType type, Map<String, String> stats) {
- super(0, null, new String[0], new byte[]{0}, toContents(new Contents(type.name(), stats)));
- _contents = new Contents(type.name(), stats);
+ @Override
+ protected void serializeMetadata(DataOutputStream output)
+ throws IOException {
+ if (_statsByStage == null) {
+ output.writeInt(0);
+ return;
+ }
+ int size = _statsByStage.size();
+ output.writeInt(size);
+ if (size > 0) {
+ byte[] bytes = new byte[4096];
+ for (ByteBuffer stat : _statsByStage) {
+ if (stat == null) {
+ output.writeBoolean(false);
+ } else {
+ output.writeBoolean(true);
+ output.writeInt(stat.remaining());
+ ByteBuffer duplicate = stat.duplicate();
+ while (duplicate.hasRemaining()) {
+ int length = Math.min(duplicate.remaining(), bytes.length);
+ duplicate.get(bytes, 0, length);
+ output.write(bytes, 0, length);
+ }
+ }
+ }
+ }
}
- private static byte[] toContents(Contents type) {
- try {
- return JSON.writeValueAsBytes(type);
- } catch (JsonProcessingException e) {
- throw new RuntimeException(e);
+ public static MetadataBlock deserialize(ByteBuffer byteBuffer, int version)
+ throws IOException {
+ switch (version) {
+ case 1:
+ case 2:
+ return new MetadataBlock(byteBuffer);
+ default:
+ throw new IOException("Unsupported metadata block version: " + version);
}
}
- public MetadataBlock(ByteBuffer byteBuffer)
+ @Override
+ protected void deserializeMetadata(ByteBuffer buffer)
throws IOException {
- super(byteBuffer);
- if (_variableSizeDataBytes != null && _variableSizeDataBytes.length > 0) {
- _contents = JSON.readValue(_variableSizeDataBytes, Contents.class);
- } else {
- _contents = new Contents();
+ try {
+ int statsSize = buffer.getInt();
+
+ List<ByteBuffer> stats = new ArrayList<>(statsSize);
+
+ for (int i = 0; i < statsSize; i++) {
+ if (buffer.get() != 0) {
+ int length = buffer.getInt();
+ buffer.limit(buffer.position() + length);
+ stats.add(buffer.slice());
+ buffer.position(buffer.limit());
+ buffer.limit(buffer.capacity());
+ } else {
+ stats.add(null);
+ }
+ }
+ _statsByStage = stats;
+ } catch (BufferUnderflowException e) {
+ LOGGER.info("Failed to read stats from metadata block. Considering it empty", e);;
+ } catch (RuntimeException e) {
+ LOGGER.warn("Failed to read stats from metadata block. Considering it empty", e);;
}
}
public MetadataBlockType getType() {
- String type = _contents.getType();
-
- // if type is null, then we're reading a legacy block where we didn't encode any
- // data. assume that it is an EOS block if there's no exceptions and an ERROR block
- // otherwise
- return type == null
- ? (getExceptions().isEmpty() ? MetadataBlockType.EOS : MetadataBlockType.ERROR)
- : MetadataBlockType.valueOf(type);
+ return _errCodeToExceptionMap.isEmpty() ? MetadataBlockType.EOS : MetadataBlockType.ERROR;
}
- public Map<String, String> getStats() {
- return _contents.getStats() != null ? _contents.getStats() : new HashMap<>();
+ /**
+ * Returns the list of serialized stats.
+ * <p>
+ * The returned list may contain nulls, which would mean that no stats were available for that stage.
+ */
+ @Nullable
+ public List<ByteBuffer> getStatsByStage() {
+ return _statsByStage;
}
@Override
@@ -156,12 +167,35 @@ public class MetadataBlock extends BaseDataBlock {
}
@Override
- public MetadataBlock toMetadataOnlyDataTable() {
- return this;
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof MetadataBlock)) {
+ return false;
+ }
+ MetadataBlock that = (MetadataBlock) o;
+ return Objects.equals(_statsByStage, that._statsByStage)
+ && _errCodeToExceptionMap.equals(that._errCodeToExceptionMap);
}
@Override
- public MetadataBlock toDataOnlyDataTable() {
- throw new UnsupportedOperationException();
+ public int hashCode() {
+ return Objects.hash(_statsByStage, _errCodeToExceptionMap);
+ }
+
+ public enum MetadataBlockType {
+ /**
+ * Indicates that this block is the final block to be sent
+ * (End Of Stream) as part of an operator chain computation.
+ */
+ EOS,
+
+ /**
+ * An {@code ERROR} metadata block indicates that there was
+ * some error during computation. To retrieve the error that
+ * occurred, use {@link MetadataBlock#getExceptions()}
+ */
+ ERROR
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
index f5aa806489..33a5e77cf7 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/RowDataBlock.java
@@ -20,6 +20,8 @@ package org.apache.pinot.common.datablock;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
import org.apache.pinot.common.utils.DataSchema;
@@ -27,7 +29,7 @@ import org.apache.pinot.common.utils.DataSchema;
* Wrapper for row-wise data table. It stores data in row-major format.
*/
public class RowDataBlock extends BaseDataBlock {
- private static final int VERSION = 1;
+ private static final int VERSION = 2;
protected int[] _columnOffsets;
protected int _rowSizeInBytes;
@@ -72,22 +74,25 @@ public class RowDataBlock extends BaseDataBlock {
return _fixedSizeData.getInt(offset + 4);
}
- @Override
- public RowDataBlock toMetadataOnlyDataTable() {
- RowDataBlock metadataOnlyDataTable = new RowDataBlock();
- metadataOnlyDataTable._metadata.putAll(_metadata);
- metadataOnlyDataTable._errCodeToExceptionMap.putAll(_errCodeToExceptionMap);
- return metadataOnlyDataTable;
+ public int getRowSizeInBytes() {
+ return _rowSizeInBytes;
}
@Override
- public RowDataBlock toDataOnlyDataTable() {
- return new RowDataBlock(_numRows, _dataSchema, _stringDictionary, _fixedSizeDataBytes, _variableSizeDataBytes);
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof RowDataBlock)) {
+ return false;
+ }
+ RowDataBlock that = (RowDataBlock) o;
+ return _rowSizeInBytes == that._rowSizeInBytes && Objects.deepEquals(_columnOffsets, that._columnOffsets);
}
- public int getRowSizeInBytes() {
- return _rowSizeInBytes;
+ @Override
+ public int hashCode() {
+ return Objects.hash(Arrays.hashCode(_columnOffsets), _rowSizeInBytes);
}
-
- // TODO: add whole-row access methods.
+// TODO: add whole-row access methods.
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
index 72121611eb..5fb1018dce 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
@@ -110,6 +110,8 @@ public interface DataTable {
NUM_SEGMENTS_PROCESSED(6, "numSegmentsProcessed", MetadataValueType.INT),
NUM_SEGMENTS_MATCHED(7, "numSegmentsMatched", MetadataValueType.INT),
NUM_CONSUMING_SEGMENTS_QUERIED(8, "numConsumingSegmentsQueried", MetadataValueType.INT),
+ // the timestamp indicating the freshness of the data queried in consuming segments.
+ // This can be ingestion timestamp if provided by the stream, or the last index time
MIN_CONSUMING_FRESHNESS_TIME_MS(9, "minConsumingFreshnessTimeMs", MetadataValueType.LONG),
TOTAL_DOCS(10, "totalDocs", MetadataValueType.LONG),
NUM_GROUPS_LIMIT_REACHED(11, "numGroupsLimitReached", MetadataValueType.STRING),
@@ -135,6 +137,7 @@ public interface DataTable {
OPERATOR_ID(31, "operatorId", MetadataValueType.STRING),
OPERATOR_EXEC_START_TIME_MS(32, "operatorExecStartTimeMs", MetadataValueType.LONG),
OPERATOR_EXEC_END_TIME_MS(33, "operatorExecEndTimeMs", MetadataValueType.LONG),
+ // Not actually used
MAX_ROWS_IN_JOIN_REACHED(34, "maxRowsInJoinReached", MetadataValueType.STRING);
// We keep this constant to track the max id added so far for backward compatibility.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java b/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java
new file mode 100644
index 0000000000..9d2818958b
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/common/datatable/StatMap.java
@@ -0,0 +1,498 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.datatable;
+
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.base.Preconditions;
+import java.io.DataInput;
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.utils.JsonUtils;
+
+
+/**
+ * A map that stores statistics.
+ * <p>
+ * Statistics must be keyed by an enum that implements {@link StatMap.Key}.
+ * <p>
+ * A stat map efficiently store, serialize and deserialize these statistics.
+ * <p>
+ * Serialization and deserialization is backward and forward compatible as long as the only change in the keys are:
+ * <ul>
+ * <li>Adding new keys</li>
+ * <li>Change the name of the keys</li>
+ * </ul>
+ *
+ * Any other change (like changing the type of key, changing their literal order are not supported or removing keys)
+ * are backward incompatible changes.
+ * @param <K>
+ */
+public class StatMap<K extends Enum<K> & StatMap.Key> {
+ private final Class<K> _keyClass;
+ private final Map<K, Object> _map;
+
+ private static final ConcurrentHashMap<Class<?>, Object[]> KEYS_BY_CLASS = new ConcurrentHashMap<>();
+
+ public StatMap(Class<K> keyClass) {
+ _keyClass = keyClass;
+ // TODO: Study whether this is fine or we should impose a single thread policy in StatMaps
+ _map = Collections.synchronizedMap(new EnumMap<>(keyClass));
+ }
+
+ public int getInt(K key) {
+ Preconditions.checkArgument(key.getType() == Type.INT, "Key %s is of type %s, not INT", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? 0 : (Integer) o;
+ }
+
+ public StatMap<K> merge(K key, int value) {
+ if (key.getType() == Type.LONG) {
+ merge(key, (long) value);
+ return this;
+ }
+ int oldValue = getInt(key);
+ int newValue = key.merge(oldValue, value);
+ if (newValue == 0) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ public long getLong(K key) {
+ if (key.getType() == Type.INT) {
+ return getInt(key);
+ }
+ Preconditions.checkArgument(key.getType() == Type.LONG, "Key %s is of type %s, not LONG", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? 0L : (Long) o;
+ }
+
+ public StatMap<K> merge(K key, long value) {
+ Preconditions.checkArgument(key.getType() == Type.LONG, "Key %s is of type %s, not LONG", key, key.getType());
+ long oldValue = getLong(key);
+ long newValue = key.merge(oldValue, value);
+ if (newValue == 0) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ public boolean getBoolean(K key) {
+ Preconditions.checkArgument(key.getType() == Type.BOOLEAN, "Key %s is of type %s, not BOOLEAN",
+ key, key.getType());
+ Object o = _map.get(key);
+ return o != null && (Boolean) o;
+ }
+
+ public StatMap<K> merge(K key, boolean value) {
+ boolean oldValue = getBoolean(key);
+ boolean newValue = key.merge(oldValue, value);
+ if (!newValue) {
+ _map.remove(key);
+ } else {
+ _map.put(key, Boolean.TRUE);
+ }
+ return this;
+ }
+
+ public String getString(K key) {
+ Preconditions.checkArgument(key.getType() == Type.STRING, "Key %s is of type %s, not STRING", key, key.getType());
+ Object o = _map.get(key);
+ return o == null ? null : (String) o;
+ }
+
+ public StatMap<K> merge(K key, String value) {
+ String oldValue = getString(key);
+ String newValue = key.merge(oldValue, value);
+ if (newValue == null) {
+ _map.remove(key);
+ } else {
+ _map.put(key, newValue);
+ }
+ return this;
+ }
+
+ /**
+ * Returns the value associated with the key.
+ * <p>
+ * Primitives will be boxed, so it is recommended to use the specific methods for each type.
+ */
+ public Object getAny(K key) {
+ switch (key.getType()) {
+ case BOOLEAN:
+ return getBoolean(key);
+ case INT:
+ return getInt(key);
+ case LONG:
+ return getLong(key);
+ case STRING:
+ return getString(key);
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+
+ /**
+ * Modifies this object to merge the values of the other object.
+ *
+ * @param other The object to merge with. This argument will not be modified.
+ * @return this object once it is modified.
+ */
+ public StatMap<K> merge(StatMap<K> other) {
+ Preconditions.checkState(_keyClass.equals(other._keyClass),
+ "Different key classes %s and %s", _keyClass, other._keyClass);
+ for (Map.Entry<K, Object> entry : other._map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ if (value == null) {
+ continue;
+ }
+ switch (key.getType()) {
+ case BOOLEAN:
+ merge(key, (boolean) value);
+ break;
+ case INT:
+ merge(key, (int) value);
+ break;
+ case LONG:
+ merge(key, (long) value);
+ break;
+ case STRING:
+ merge(key, (String) value);
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+ return this;
+ }
+
+ public StatMap<K> merge(DataInput input)
+ throws IOException {
+ byte serializedKeys = input.readByte();
+
+ K[] keys = (K[]) KEYS_BY_CLASS.computeIfAbsent(_keyClass, k -> k.getEnumConstants());
+ for (byte i = 0; i < serializedKeys; i++) {
+ int ordinal = input.readByte();
+ K key = keys[ordinal];
+ switch (key.getType()) {
+ case BOOLEAN:
+ merge(key, true);
+ break;
+ case INT:
+ merge(key, input.readInt());
+ break;
+ case LONG:
+ merge(key, input.readLong());
+ break;
+ case STRING:
+ merge(key, input.readUTF());
+ break;
+ default:
+ throw new IllegalStateException("Unknown type " + key.getType());
+ }
+ }
+ return this;
+ }
+
+ public ObjectNode asJson() {
+ ObjectNode node = JsonUtils.newObjectNode();
+
+ for (Map.Entry<K, Object> entry : _map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ switch (key.getType()) {
+ case BOOLEAN:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), false);
+ }
+ } else {
+ node.put(key.getStatName(), (boolean) value);
+ }
+ break;
+ case INT:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), 0);
+ }
+ } else {
+ node.put(key.getStatName(), (int) value);
+ }
+ break;
+ case LONG:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), 0L);
+ }
+ } else {
+ node.put(key.getStatName(), (long) value);
+ }
+ break;
+ case STRING:
+ if (value == null) {
+ if (key.includeDefaultInJson()) {
+ node.put(key.getStatName(), "");
+ }
+ } else {
+ node.put(key.getStatName(), (String) value);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+
+ return node;
+ }
+
+ public void serialize(DataOutput output)
+ throws IOException {
+
+ assert checkContainsNoDefault() : "No default value should be stored in the map";
+ output.writeByte(_map.size());
+
+ // We use written keys just to fail fast in tests if the number of keys written
+ // is not the same as the number of keys
+ int writtenKeys = 0;
+ K[] keys = (K[]) KEYS_BY_CLASS.computeIfAbsent(_keyClass, k -> k.getEnumConstants());
+ for (int ordinal = 0; ordinal < keys.length; ordinal++) {
+ K key = keys[ordinal];
+ switch (key.getType()) {
+ case BOOLEAN: {
+ if (getBoolean(key)) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ }
+ break;
+ }
+ case INT: {
+ int value = getInt(key);
+ if (value != 0) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeInt(value);
+ }
+ break;
+ }
+ case LONG: {
+ long value = getLong(key);
+ if (value != 0) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeLong(value);
+ }
+ break;
+ }
+ case STRING: {
+ String value = getString(key);
+ if (value != null) {
+ writtenKeys++;
+ output.writeByte(ordinal);
+ output.writeUTF(value);
+ }
+ break;
+ }
+ default:
+ throw new IllegalStateException("Unknown type " + key.getType());
+ }
+ }
+ assert writtenKeys == _map.size() : "Written keys " + writtenKeys + " but map size " + _map.size();
+ }
+
+ private boolean checkContainsNoDefault() {
+ for (Map.Entry<K, Object> entry : _map.entrySet()) {
+ K key = entry.getKey();
+ Object value = entry.getValue();
+ switch (key.getType()) {
+ case BOOLEAN:
+ if (value == null || !(boolean) value) {
+ throw new IllegalStateException("Boolean value must be true but " + value + " is stored for key " + key);
+ }
+ break;
+ case INT:
+ if (value == null || (int) value == 0) {
+ throw new IllegalStateException("Int value must be non-zero but " + value + " is stored for key " + key);
+ }
+ break;
+ case LONG:
+ if (value == null || (long) value == 0) {
+ throw new IllegalStateException("Long value must be non-zero but " + value + " is stored for key " + key);
+ }
+ break;
+ case STRING:
+ if (value == null) {
+ throw new IllegalStateException("String value must be non-null but null is stored for key " + key);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException("Unsupported type: " + key.getType());
+ }
+ }
+ return true;
+ }
+
+ public static String getDefaultStatName(Key key) {
+ String name = key.name();
+ StringBuilder result = new StringBuilder();
+ boolean capitalizeNext = false;
+
+ for (char c : name.toCharArray()) {
+ if (c == '_') {
+ capitalizeNext = true;
+ } else {
+ if (capitalizeNext) {
+ result.append(c);
+ capitalizeNext = false;
+ } else {
+ result.append(Character.toLowerCase(c));
+ }
+ }
+ }
+
+ return result.toString();
+ }
+
+ public static <K extends Enum<K> & Key> StatMap<K> deserialize(DataInput input, Class<K> keyClass)
+ throws IOException {
+ StatMap<K> result = new StatMap<>(keyClass);
+ result.merge(input);
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ StatMap<?> statMap = (StatMap<?>) o;
+ return Objects.equals(_map, statMap._map);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(_map);
+ }
+
+ @Override
+ public String toString() {
+ return asJson().toString();
+ }
+
+ public Class<K> getKeyClass() {
+ return _keyClass;
+ }
+
+ public boolean isEmpty() {
+ return _map.isEmpty();
+ }
+
+ public Iterable<K> keySet() {
+ return _map.keySet();
+ }
+
+ public interface Key {
+ String name();
+
+ /**
+ * The name of the stat used to report it. Names must be unique on the same key family.
+ */
+ default String getStatName() {
+ return getDefaultStatName(this);
+ }
+
+ default int merge(int value1, int value2) {
+ return value1 + value2;
+ }
+
+ default long merge(long value1, long value2) {
+ return value1 + value2;
+ }
+
+ default boolean merge(boolean value1, boolean value2) {
+ return value1 || value2;
+ }
+
+ default String merge(@Nullable String value1, @Nullable String value2) {
+ return value2 != null ? value2 : value1;
+ }
+
+ /**
+ * The type of the values associated to this key.
+ */
+ Type getType();
+
+ default boolean includeDefaultInJson() {
+ return false;
+ }
+
+ static int minPositive(int value1, int value2) {
+ if (value1 == 0 && value2 >= 0) {
+ return value2;
+ }
+ if (value2 == 0 && value1 >= 0) {
+ return value1;
+ }
+ return Math.min(value1, value2);
+ }
+
+ static long minPositive(long value1, long value2) {
+ if (value1 == 0 && value2 >= 0) {
+ return value2;
+ }
+ if (value2 == 0 && value1 >= 0) {
+ return value1;
+ }
+ return Math.min(value1, value2);
+ }
+
+ static int eqNotZero(int value1, int value2) {
+ if (value1 != value2) {
+ if (value1 == 0) {
+ return value2;
+ } else if (value2 == 0) {
+ return value1;
+ } else {
+ throw new IllegalStateException("Cannot merge non-zero values: " + value1 + " and " + value2);
+ }
+ }
+ return value1;
+ }
+ }
+
+ public enum Type {
+ BOOLEAN,
+ INT,
+ LONG,
+ STRING
+ }
+}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java b/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
index 06f65fc2aa..0d7ec268a7 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/BrokerResponse.java
@@ -18,9 +18,12 @@
*/
package org.apache.pinot.common.response;
+import java.io.IOException;
import java.util.List;
+import javax.annotation.Nullable;
import org.apache.pinot.common.response.broker.QueryProcessingException;
import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.spi.utils.JsonUtils;
/**
@@ -33,6 +36,8 @@ public interface BrokerResponse {
*/
void setExceptions(List<ProcessingException> exceptions);
+ void addToExceptions(QueryProcessingException processingException);
+
/**
* Set the number of servers got queried by the broker.
*
@@ -47,21 +52,20 @@ public interface BrokerResponse {
*/
void setNumServersResponded(int numServersResponded);
+ long getTimeUsedMs();
+
/**
* Set the total time used in request handling, into the broker response.
*/
void setTimeUsedMs(long timeUsedMs);
- /**
- * Set the total number of rows in result set
- */
- void setNumRowsResultSet(int numRowsResultSet);
-
/**
* Convert the broker response to JSON String.
*/
- String toJsonString()
- throws Exception;
+ default String toJsonString()
+ throws IOException {
+ return JsonUtils.objectToString(this);
+ }
/**
* Returns the number of servers queried.
@@ -147,12 +151,13 @@ public interface BrokerResponse {
* set the result table.
* @param resultTable result table to be set.
*/
- void setResultTable(ResultTable resultTable);
+ void setResultTable(@Nullable ResultTable resultTable);
/**
* Get the result table.
* @return result table.
*/
+ @Nullable
ResultTable getResultTable();
/**
@@ -163,12 +168,10 @@ public interface BrokerResponse {
/**
* Get the total number of rows in result set
*/
- int getNumRowsResultSet();
-
- /**
- * Set the total thread cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineThreadCpuTimeNs(long offlineThreadCpuTimeNs);
+ default int getNumRowsResultSet() {
+ ResultTable resultTable = getResultTable();
+ return resultTable == null ? 0 : resultTable.getRows().size();
+ }
/**
* Get the thread cpu time used against offline table in request handling, from the broker response.
@@ -180,74 +183,43 @@ public interface BrokerResponse {
*/
long getRealtimeThreadCpuTimeNs();
- /**
- * Set the total thread cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeThreadCpuTimeNs(long realtimeThreadCpuTimeNs);
-
/**
* Get the system activities cpu time used against offline table in request handling, from the broker response.
*/
long getOfflineSystemActivitiesCpuTimeNs();
- /**
- * Set the system activities cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineSystemActivitiesCpuTimeNs(long offlineSystemActivitiesCpuTimeNs);
-
/**
* Get the system activities cpu time used against realtime table in request handling, from the broker response.
*/
long getRealtimeSystemActivitiesCpuTimeNs();
- /**
- * Set the system activities cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeSystemActivitiesCpuTimeNs(long realtimeSystemActivitiesCpuTimeNs);
-
/**
* Get the response serialization cpu time used against offline table in request handling, from the broker response.
*/
long getOfflineResponseSerializationCpuTimeNs();
- /**
- * Set the response serialization cpu time used against offline table in request handling, into the broker response.
- */
- void setOfflineResponseSerializationCpuTimeNs(long offlineResponseSerializationCpuTimeNs);
-
/**
* Get the response serialization cpu time used against realtime table in request handling, from the broker response.
*/
long getRealtimeResponseSerializationCpuTimeNs();
- /**
- * Set the response serialization cpu time used against realtime table in request handling, into the broker response.
- */
- void setRealtimeResponseSerializationCpuTimeNs(long realtimeResponseSerializationCpuTimeNs);
-
/**
* Get the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
* against offline table in request handling, from the broker response.
*/
- long getOfflineTotalCpuTimeNs();
-
- /**
- * Set the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
- * against offline table in request handling, into the broker response.
- */
- void setOfflineTotalCpuTimeNs(long offlineTotalCpuTimeNs);
+ default long getOfflineTotalCpuTimeNs() {
+ return getOfflineThreadCpuTimeNs() + getOfflineSystemActivitiesCpuTimeNs()
+ + getOfflineResponseSerializationCpuTimeNs();
+ }
/**
* Get the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
* against realtime table in request handling, from the broker response.
*/
- long getRealtimeTotalCpuTimeNs();
-
- /**
- * Set the total cpu time(thread cpu time + system activities cpu time + response serialization cpu time) used
- * against realtime table in request handling, into the broker response.
- */
- void setRealtimeTotalCpuTimeNs(long realtimeTotalCpuTimeNs);
+ default long getRealtimeTotalCpuTimeNs() {
+ return getRealtimeThreadCpuTimeNs() + getRealtimeSystemActivitiesCpuTimeNs()
+ + getRealtimeResponseSerializationCpuTimeNs();
+ }
/**
* Get the total number of segments pruned on the Broker side
@@ -264,11 +236,6 @@ public interface BrokerResponse {
*/
long getNumSegmentsPrunedByServer();
- /**
- * Set the total number of segments pruned on the Server side
- */
- void setNumSegmentsPrunedByServer(long numSegmentsPrunedByServer);
-
/**
* Get the total number of segments pruned due to invalid data or schema.
*
@@ -276,13 +243,6 @@ public interface BrokerResponse {
*/
long getNumSegmentsPrunedInvalid();
- /**
- * Set the total number of segments pruned due to invalid data or schema.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedInvalid(long numSegmentsPrunedInvalid);
-
/**
* Get the total number of segments pruned by applying the limit optimization.
*
@@ -290,13 +250,6 @@ public interface BrokerResponse {
*/
long getNumSegmentsPrunedByLimit();
- /**
- * Set the total number of segments pruned by applying the limit optimization.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedByLimit(long numSegmentsPrunedByLimit);
-
/**
* Get the total number of segments pruned applying value optimizations, like bloom filters.
*
@@ -304,33 +257,16 @@ public interface BrokerResponse {
*/
long getNumSegmentsPrunedByValue();
- /**
- * Set the total number of segments pruned applying value optimizations, like bloom filters.
- *
- * This value is always lower or equal than {@link #getNumSegmentsPrunedByServer()}
- */
- void setNumSegmentsPrunedByValue(long numSegmentsPrunedByValue);
-
/**
* Get the total number of segments with an EmptyFilterOperator when Explain Plan is called
*/
long getExplainPlanNumEmptyFilterSegments();
- /**
- * Set the total number of segments with an EmptyFilterOperator when Explain Plan is called
- */
- void setExplainPlanNumEmptyFilterSegments(long explainPlanNumEmptyFilterSegments);
-
/**
* Get the total number of segments with a MatchAllFilterOperator when Explain Plan is called
*/
long getExplainPlanNumMatchAllFilterSegments();
- /**
- * Set the total number of segments with a MatchAllFilterOperator when Explain Plan is called
- */
- void setExplainPlanNumMatchAllFilterSegments(long explainPlanNumMatchAllFilterSegments);
-
/**
* get request ID for the query
*/
@@ -354,4 +290,6 @@ public interface BrokerResponse {
long getBrokerReduceTimeMs();
void setBrokerReduceTimeMs(long brokerReduceTimeMs);
+
+ boolean isPartialResult();
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
index 9fe098e26d..d22ee51c89 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNative.java
@@ -22,6 +22,7 @@ import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
+import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
@@ -37,7 +38,7 @@ import org.apache.pinot.spi.utils.JsonUtils;
/**
* This class implements pinot-broker's response format for any given query.
* All fields either primitive data types, or native objects (as opposed to JSONObjects).
- *
+ * <p>
* Supports serialization via JSON.
*/
@JsonPropertyOrder({
@@ -78,8 +79,6 @@ public class BrokerResponseNative implements BrokerResponse {
private long _totalDocs = 0L;
private boolean _numGroupsLimitReached = false;
- private boolean _maxRowsInJoinReached = false;
- private boolean _partialResult = false;
private long _timeUsedMs = 0L;
private long _offlineThreadCpuTimeNs = 0L;
private long _realtimeThreadCpuTimeNs = 0L;
@@ -131,6 +130,7 @@ public class BrokerResponseNative implements BrokerResponse {
return new BrokerResponseNative();
}
+ @VisibleForTesting
public static BrokerResponseNative fromJsonString(String jsonString)
throws IOException {
return JsonUtils.stringToObject(jsonString, BrokerResponseNative.class);
@@ -143,7 +143,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("offlineSystemActivitiesCpuTimeNs")
- @Override
public void setOfflineSystemActivitiesCpuTimeNs(long offlineSystemActivitiesCpuTimeNs) {
_offlineSystemActivitiesCpuTimeNs = offlineSystemActivitiesCpuTimeNs;
}
@@ -155,7 +154,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("realtimeSystemActivitiesCpuTimeNs")
- @Override
public void setRealtimeSystemActivitiesCpuTimeNs(long realtimeSystemActivitiesCpuTimeNs) {
_realtimeSystemActivitiesCpuTimeNs = realtimeSystemActivitiesCpuTimeNs;
}
@@ -167,7 +165,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("offlineThreadCpuTimeNs")
- @Override
public void setOfflineThreadCpuTimeNs(long timeUsedMs) {
_offlineThreadCpuTimeNs = timeUsedMs;
}
@@ -179,7 +176,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("realtimeThreadCpuTimeNs")
- @Override
public void setRealtimeThreadCpuTimeNs(long timeUsedMs) {
_realtimeThreadCpuTimeNs = timeUsedMs;
}
@@ -191,7 +187,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("offlineResponseSerializationCpuTimeNs")
- @Override
public void setOfflineResponseSerializationCpuTimeNs(long offlineResponseSerializationCpuTimeNs) {
_offlineResponseSerializationCpuTimeNs = offlineResponseSerializationCpuTimeNs;
}
@@ -203,7 +198,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("realtimeResponseSerializationCpuTimeNs")
- @Override
public void setRealtimeResponseSerializationCpuTimeNs(long realtimeResponseSerializationCpuTimeNs) {
_realtimeResponseSerializationCpuTimeNs = realtimeResponseSerializationCpuTimeNs;
}
@@ -215,7 +209,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("offlineTotalCpuTimeNs")
- @Override
public void setOfflineTotalCpuTimeNs(long offlineTotalCpuTimeNs) {
_offlineTotalCpuTimeNs = offlineTotalCpuTimeNs;
}
@@ -227,7 +220,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("realtimeTotalCpuTimeNs")
- @Override
public void setRealtimeTotalCpuTimeNs(long realtimeTotalCpuTimeNs) {
_realtimeTotalCpuTimeNs = realtimeTotalCpuTimeNs;
}
@@ -251,7 +243,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("numSegmentsPrunedByServer")
- @Override
public void setNumSegmentsPrunedByServer(long numSegmentsPrunedByServer) {
_numSegmentsPrunedByServer = numSegmentsPrunedByServer;
}
@@ -263,7 +254,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("numSegmentsPrunedInvalid")
- @Override
public void setNumSegmentsPrunedInvalid(long numSegmentsPrunedInvalid) {
_numSegmentsPrunedInvalid = numSegmentsPrunedInvalid;
}
@@ -275,7 +265,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("numSegmentsPrunedByLimit")
- @Override
public void setNumSegmentsPrunedByLimit(long numSegmentsPrunedByLimit) {
_numSegmentsPrunedByLimit = numSegmentsPrunedByLimit;
}
@@ -287,7 +276,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("numSegmentsPrunedByValue")
- @Override
public void setNumSegmentsPrunedByValue(long numSegmentsPrunedByValue) {
_numSegmentsPrunedByValue = numSegmentsPrunedByValue;
}
@@ -299,7 +287,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("explainPlanNumEmptyFilterSegments")
- @Override
public void setExplainPlanNumEmptyFilterSegments(long explainPlanNumEmptyFilterSegments) {
_explainPlanNumEmptyFilterSegments = explainPlanNumEmptyFilterSegments;
}
@@ -311,7 +298,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("explainPlanNumMatchAllFilterSegments")
- @Override
public void setExplainPlanNumMatchAllFilterSegments(long explainPlanNumMatchAllFilterSegments) {
_explainPlanNumMatchAllFilterSegments = explainPlanNumMatchAllFilterSegments;
}
@@ -350,7 +336,6 @@ public class BrokerResponseNative implements BrokerResponse {
}
@JsonProperty("numServersQueried")
- @Override
public void setNumServersQueried(int numServersQueried) {
_numServersQueried = numServersQueried;
}
@@ -448,6 +433,7 @@ public class BrokerResponseNative implements BrokerResponse {
public long getNumConsumingSegmentsProcessed() {
return _numConsumingSegmentsProcessed;
}
+
@JsonProperty("numConsumingSegmentsProcessed")
public void setNumConsumingSegmentsProcessed(long numConsumingSegmentsProcessed) {
_numConsumingSegmentsProcessed = numConsumingSegmentsProcessed;
@@ -497,27 +483,18 @@ public class BrokerResponseNative implements BrokerResponse {
_numGroupsLimitReached = numGroupsLimitReached;
}
- @JsonProperty("maxRowsInJoinReached")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
public boolean isMaxRowsInJoinReached() {
- return _maxRowsInJoinReached;
- }
-
- @JsonProperty("maxRowsInJoinReached")
- public void setMaxRowsInJoinReached(boolean maxRowsInJoinReached) {
- _maxRowsInJoinReached = maxRowsInJoinReached;
+ return false;
}
- @JsonProperty("partialResult")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
public boolean isPartialResult() {
- return _partialResult;
- }
-
- @JsonProperty("partialResult")
- public void setPartialResult(boolean partialResult) {
- _partialResult = partialResult;
+ return isNumGroupsLimitReached() || getExceptionsSize() > 0 || isMaxRowsInJoinReached();
}
@JsonProperty("timeUsedMs")
+ @Override
public long getTimeUsedMs() {
return _timeUsedMs;
}
@@ -528,16 +505,10 @@ public class BrokerResponseNative implements BrokerResponse {
_timeUsedMs = timeUsedMs;
}
- @JsonProperty("numRowsResultSet")
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
@Override
public int getNumRowsResultSet() {
- return _numRowsResultSet;
- }
-
- @JsonProperty("numRowsResultSet")
- @Override
- public void setNumRowsResultSet(int numRowsResultSet) {
- _numRowsResultSet = numRowsResultSet;
+ return BrokerResponse.super.getNumRowsResultSet();
}
@JsonProperty("segmentStatistics")
@@ -560,15 +531,11 @@ public class BrokerResponseNative implements BrokerResponse {
_traceInfo = traceInfo;
}
- @Override
- public String toJsonString()
- throws IOException {
- return JsonUtils.objectToString(this);
- }
-
@JsonIgnore
@Override
public void setExceptions(List<ProcessingException> exceptions) {
+ // TODO: This is incorrect. It is adding and not setting the exceptions
+ // But there is some code that seems to depend on this.
for (ProcessingException exception : exceptions) {
_processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
index 60a22460f3..40f521e251 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
@@ -18,16 +18,17 @@
*/
package org.apache.pinot.common.response.broker;
+import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
-import java.io.IOException;
+import com.fasterxml.jackson.databind.node.ObjectNode;
import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Collections;
import java.util.List;
-import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.common.response.BrokerResponse;
import org.apache.pinot.common.response.ProcessingException;
-import org.apache.pinot.common.utils.DataSchema;
-import org.apache.pinot.spi.utils.JsonUtils;
/**
@@ -46,28 +47,36 @@ import org.apache.pinot.spi.utils.JsonUtils;
"realtimeResponseSerializationCpuTimeNs", "offlineTotalCpuTimeNs", "realtimeTotalCpuTimeNs", "brokerReduceTimeMs",
"segmentStatistics", "traceInfo", "partialResult"
})
-public class BrokerResponseNativeV2 extends BrokerResponseNative {
+public class BrokerResponseNativeV2 implements BrokerResponse {
- private final Map<Integer, BrokerResponseStats> _stageIdStats = new HashMap<>();
+ private ObjectNode _stageStats = null;
+ /**
+ * The max number of rows seen at runtime.
+ * <p>
+ * In single-stage this doesn't make sense given it is the max number of rows read from the table. But in multi-stage
+ * virtual rows can be generated. For example, in a join query, the number of rows can be more than the number of rows
+ * in the table.
+ */
+ private long _maxRowsInOperator = 0;
+ private final StatMap<StatKey> _brokerStats = new StatMap<>(StatKey.class);
+ private final List<QueryProcessingException> _processingExceptions;
+ private ResultTable _resultTable;
+ private String _requestId;
+ private String _brokerId;
public BrokerResponseNativeV2() {
+ _processingExceptions = new ArrayList<>();
}
public BrokerResponseNativeV2(ProcessingException exception) {
- super(exception);
+ this(Collections.singletonList(exception));
}
public BrokerResponseNativeV2(List<ProcessingException> exceptions) {
- super(exceptions);
- }
-
- /** Generate EXPLAIN PLAN output when queries are evaluated by Broker without going to the Server. */
- private static BrokerResponseNativeV2 getBrokerResponseExplainPlanOutput() {
- BrokerResponseNativeV2 brokerResponse = BrokerResponseNativeV2.empty();
- List<Object[]> rows = new ArrayList<>();
- rows.add(new Object[]{"BROKER_EVALUATE", 0, -1});
- brokerResponse.setResultTable(new ResultTable(DataSchema.EXPLAIN_RESULT_SCHEMA, rows));
- return brokerResponse;
+ _processingExceptions = new ArrayList<>(exceptions.size());
+ for (ProcessingException exception : exceptions) {
+ _processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
+ }
}
/**
@@ -77,20 +86,338 @@ public class BrokerResponseNativeV2 extends BrokerResponseNative {
return new BrokerResponseNativeV2();
}
- public static BrokerResponseNativeV2 fromJsonString(String jsonString)
- throws IOException {
- return JsonUtils.stringToObject(jsonString, BrokerResponseNativeV2.class);
+ public void setStageStats(ObjectNode stageStats) {
+ _stageStats = stageStats;
+ }
+
+ @JsonProperty
+ public ObjectNode getStageStats() {
+ return _stageStats;
+ }
+
+ /**
+ * Get the max number of rows seen by a single operator in the query processing chain.
+ */
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ public long getMaxRowsInOperator() {
+ return _maxRowsInOperator;
+ }
+
+ public void mergeMaxRowsInOperator(long maxRows) {
+ _maxRowsInOperator = Math.max(_maxRowsInOperator, maxRows);
+ }
+
+ @Override
+ public long getTimeUsedMs() {
+ return _brokerStats.getLong(StatKey.TIME_USED_MS);
+ }
+
+ @Override
+ public void setTimeUsedMs(long timeUsedMs) {
+ _brokerStats.merge(StatKey.TIME_USED_MS, timeUsedMs);
+ }
+
+ @Override
+ public long getNumDocsScanned() {
+ return _brokerStats.getLong(StatKey.NUM_DOCS_SCANNED);
+ }
+
+ @Override
+ public long getNumEntriesScannedInFilter() {
+ return _brokerStats.getLong(StatKey.NUM_ENTRIES_SCANNED_IN_FILTER);
+ }
+
+ @Override
+ public long getNumEntriesScannedPostFilter() {
+ return _brokerStats.getLong(StatKey.NUM_ENTRIES_SCANNED_POST_FILTER);
+ }
+
+ @Override
+ public long getNumSegmentsQueried() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_QUERIED);
+ }
+
+ @Override
+ public long getNumSegmentsProcessed() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PROCESSED);
+ }
+
+ @Override
+ public long getNumSegmentsMatched() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_MATCHED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsQueried() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_QUERIED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsProcessed() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_PROCESSED);
+ }
+
+ @Override
+ public long getNumConsumingSegmentsMatched() {
+ return _brokerStats.getLong(StatKey.NUM_CONSUMING_SEGMENTS_MATCHED);
+ }
+
+ @Override
+ public long getMinConsumingFreshnessTimeMs() {
+ return _brokerStats.getLong(StatKey.MIN_CONSUMING_FRESHNESS_TIME_MS);
+ }
+
+ @Override
+ public long getTotalDocs() {
+ return _brokerStats.getLong(StatKey.TOTAL_DOCS);
+ }
+
+ @Override
+ public boolean isNumGroupsLimitReached() {
+ return _brokerStats.getBoolean(StatKey.NUM_GROUPS_LIMIT_REACHED);
+ }
+
+ public void mergeNumGroupsLimitReached(boolean numGroupsLimitReached) {
+ _brokerStats.merge(StatKey.NUM_GROUPS_LIMIT_REACHED, numGroupsLimitReached);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByServer() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_SERVER);
}
- public void addStageStat(Integer stageId, BrokerResponseStats brokerResponseStats) {
- // StageExecutionWallTime will always be there, other stats are optional such as OperatorStats
- if (brokerResponseStats.getStageExecWallTimeMs() != -1) {
- _stageIdStats.put(stageId, brokerResponseStats);
+ @Override
+ public long getNumSegmentsPrunedInvalid() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_INVALID);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByLimit() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_LIMIT);
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByValue() {
+ return _brokerStats.getLong(StatKey.NUM_SEGMENTS_PRUNED_BY_VALUE);
+ }
+
+ @Override
+ public long getExplainPlanNumEmptyFilterSegments() {
+ return _brokerStats.getLong(StatKey.EXPLAIN_PLAN_NUM_EMPTY_FILTER_SEGMENTS);
+ }
+
+ @Override
+ public long getExplainPlanNumMatchAllFilterSegments() {
+ return _brokerStats.getLong(StatKey.EXPLAIN_PLAN_NUM_MATCH_ALL_FILTER_SEGMENTS);
+ }
+
+ @Override
+ public long getOfflineTotalCpuTimeNs() {
+ return getOfflineThreadCpuTimeNs() + getOfflineSystemActivitiesCpuTimeNs()
+ + getOfflineResponseSerializationCpuTimeNs();
+ }
+
+ @Override
+ public long getRealtimeTotalCpuTimeNs() {
+ return getRealtimeThreadCpuTimeNs() + getRealtimeSystemActivitiesCpuTimeNs()
+ + getRealtimeResponseSerializationCpuTimeNs();
+ }
+
+ @Override
+ public void setExceptions(List<ProcessingException> exceptions) {
+ for (ProcessingException exception : exceptions) {
+ _processingExceptions.add(new QueryProcessingException(exception.getErrorCode(), exception.getMessage()));
}
}
- @JsonProperty("stageStats")
- public Map<Integer, BrokerResponseStats> getStageIdStats() {
- return _stageIdStats;
+ public void addToExceptions(QueryProcessingException processingException) {
+ _processingExceptions.add(processingException);
+ }
+
+ @Override
+ public int getNumServersQueried() {
+ return _brokerStats.getInt(StatKey.NUM_SERVERS_QUERIED);
+ }
+
+ @Override
+ public void setNumServersQueried(int numServersQueried) {
+ _brokerStats.merge(StatKey.NUM_SERVERS_QUERIED, numServersQueried);
+ }
+
+ @Override
+ public int getNumServersResponded() {
+ return _brokerStats.getInt(StatKey.NUM_SERVERS_RESPONDED);
+ }
+
+ @Override
+ public void setNumServersResponded(int numServersResponded) {
+ _brokerStats.merge(StatKey.NUM_SERVERS_RESPONDED, numServersResponded);
+ }
+
+ @JsonProperty("maxRowsInJoinReached")
+ public boolean isMaxRowsInJoinReached() {
+ return _brokerStats.getBoolean(StatKey.MAX_ROWS_IN_JOIN_REACHED);
+ }
+
+ @JsonProperty("maxRowsInJoinReached")
+ public void mergeMaxRowsInJoinReached(boolean maxRowsInJoinReached) {
+ _brokerStats.merge(StatKey.MAX_ROWS_IN_JOIN_REACHED, maxRowsInJoinReached);
+ }
+
+ @Override
+ public int getExceptionsSize() {
+ return _processingExceptions.size();
+ }
+
+ @Override
+ public void setResultTable(@Nullable ResultTable resultTable) {
+ _resultTable = resultTable;
+ }
+
+ @Nullable
+ @Override
+ @JsonInclude(JsonInclude.Include.NON_NULL)
+ public ResultTable getResultTable() {
+ return _resultTable;
+ }
+
+ @JsonProperty("exceptions")
+ @Override
+ public List<QueryProcessingException> getProcessingExceptions() {
+ return List.of();
+ }
+
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ @Override
+ public int getNumRowsResultSet() {
+ return BrokerResponse.super.getNumRowsResultSet();
+ }
+
+ @Override
+ public long getOfflineThreadCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeThreadCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getOfflineSystemActivitiesCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeSystemActivitiesCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getOfflineResponseSerializationCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getRealtimeResponseSerializationCpuTimeNs() {
+ return 0;
+ }
+
+ @Override
+ public long getNumSegmentsPrunedByBroker() {
+ return _brokerStats.getInt(StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER);
+ }
+
+ @Override
+ public void setNumSegmentsPrunedByBroker(long numSegmentsPrunedByBroker) {
+ _brokerStats.merge(StatKey.NUM_SEGMENTS_PRUNED_BY_BROKER, (int) numSegmentsPrunedByBroker);
+ }
+
+ @Override
+ public String getRequestId() {
+ return _requestId;
+ }
+
+ @Override
+ public void setRequestId(String requestId) {
+ _requestId = requestId;
+ }
+
+ @Override
+ public String getBrokerId() {
+ return _brokerId;
+ }
+
+ @Override
+ public void setBrokerId(String requestId) {
+ _brokerId = requestId;
+ }
+
+ @Override
+ public long getBrokerReduceTimeMs() {
+ return _brokerStats.getLong(StatKey.BROKER_REDUCE_TIME_MS);
+ }
+
+ @Override
+ public void setBrokerReduceTimeMs(long brokerReduceTimeMs) {
+ _brokerStats.merge(StatKey.BROKER_REDUCE_TIME_MS, brokerReduceTimeMs);
+ }
+
+ @JsonProperty(access = JsonProperty.Access.READ_ONLY)
+ @Override
+ public boolean isPartialResult() {
+ return isNumGroupsLimitReached() || getExceptionsSize() > 0 || isMaxRowsInJoinReached();
+ }
+
+ public void addServerStats(StatMap<StatKey> serverStats) {
+ // Set execution statistics.
+ _brokerStats.merge(serverStats);
+ }
+
+ public enum StatKey implements StatMap.Key {
+ TIME_USED_MS(StatMap.Type.LONG),
+ NUM_DOCS_SCANNED(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_IN_FILTER(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_POST_FILTER(StatMap.Type.LONG),
+ NUM_SEGMENTS_QUERIED(StatMap.Type.INT),
+ NUM_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_SEGMENTS_MATCHED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_QUERIED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_MATCHED(StatMap.Type.INT),
+ MIN_CONSUMING_FRESHNESS_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public long merge(long value1, long value2) {
+ return StatMap.Key.minPositive(value1, value2);
+ }
+ },
+ TOTAL_DOCS(StatMap.Type.LONG),
+ NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
+ NUM_SEGMENTS_PRUNED_BY_SERVER(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_INVALID(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_LIMIT(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_VALUE(StatMap.Type.INT),
+ NUM_SERGMENTS_PRUNED_BY_BROKER(StatMap.Type.INT),
+ EXPLAIN_PLAN_NUM_EMPTY_FILTER_SEGMENTS(StatMap.Type.INT),
+ EXPLAIN_PLAN_NUM_MATCH_ALL_FILTER_SEGMENTS(StatMap.Type.INT),
+ THREAD_CPU_TIME_NS(StatMap.Type.LONG),
+ SYSTEM_ACTIVITIES_CPU_TIME_NS(StatMap.Type.LONG),
+ RESPONSE_SER_CPU_TIME_NS(StatMap.Type.LONG),
+ NUM_SEGMENTS_PRUNED_BY_BROKER(StatMap.Type.INT),
+ MAX_ROWS_IN_JOIN_REACHED(StatMap.Type.BOOLEAN),
+ NUM_SERVERS_RESPONDED(StatMap.Type.INT),
+ NUM_SERVERS_QUERIED(StatMap.Type.INT),
+ BROKER_REDUCE_TIME_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
index 3acc9f3349..9763187600 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseStats.java
@@ -21,12 +21,10 @@ package org.apache.pinot.common.response.broker;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
-import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import org.apache.pinot.spi.utils.JsonUtils;
// TODO: Decouple the execution stats aggregator logic and make it into a util that can aggregate 2 values with the
@@ -109,11 +107,6 @@ public class BrokerResponseStats extends BrokerResponseNative {
_stageExecutionUnit = stageExecutionUnit;
}
- public String toJsonString()
- throws IOException {
- return JsonUtils.objectToString(this);
- }
-
@JsonProperty("operatorStats")
public Map<String, Map<String, String>> getOperatorStats() {
return _operatorStats;
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java
new file mode 100644
index 0000000000..763fd7344c
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/BaseDataBlockContract.java
@@ -0,0 +1,42 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.datablock;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.testng.Assert.*;
+
+
+public abstract class BaseDataBlockContract {
+
+ protected abstract BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException;
+
+ public void testSerdeCorrectness(BaseDataBlock dataBlock)
+ throws IOException {
+ byte[] bytes = dataBlock.toBytes();
+ ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
+ int versionType = DataBlockUtils.readVersionType(byteBuffer);
+ BaseDataBlock deserialize = deserialize(byteBuffer, versionType);
+
+ assertEquals(byteBuffer.position(), bytes.length, "Buffer position should be at the end of the buffer");
+ assertEquals(deserialize, dataBlock, "Deserialized data block should be the same as the original data block");
+ }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
index 0f85f4a69e..cfceddc65d 100644
--- a/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/MetadataBlockTest.java
@@ -18,43 +18,185 @@
*/
package org.apache.pinot.common.datablock;
-import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Lists;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Collections;
import java.util.HashMap;
+import java.util.List;
import org.testng.annotations.Test;
import static org.testng.Assert.*;
-public class MetadataBlockTest {
+public class MetadataBlockTest extends BaseDataBlockContract {
+ @Override
+ protected BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException {
+ return new MetadataBlock(byteBuffer);
+ }
@Test
- public void shouldEncodeContentsAsJSON()
+ public void emptyMetadataBlock()
throws Exception {
// Given:
- MetadataBlock.MetadataBlockType type = MetadataBlock.MetadataBlockType.EOS;
-
// When:
- MetadataBlock metadataBlock = new MetadataBlock(type);
+ MetadataBlock metadataBlock = MetadataBlock.newEos();
// Then:
- byte[] expected = new ObjectMapper().writeValueAsBytes(new MetadataBlock.Contents("EOS", new HashMap<>()));
- assertEquals(metadataBlock._variableSizeDataBytes, expected);
+ byte[] expectedFixed = new byte[0];
+ assertEquals(metadataBlock._fixedSizeDataBytes, expectedFixed);
+ byte[] expectedVariable = new byte[0];
+ assertEquals(metadataBlock._variableSizeDataBytes, expectedVariable);
+
+ List<ByteBuffer> statsByStage = metadataBlock.getStatsByStage();
+ assertNotNull(statsByStage, "Expected stats by stage to be non-null");
+ assertEquals(statsByStage.size(), 0, "Expected no stats by stage");
+ }
+
+ @Test
+ public void emptyDataBlockCorrectness()
+ throws IOException {
+ testSerdeCorrectness(MetadataBlock.newEos());
+ }
+
+ @Test
+ public void v1ErrorWithExceptionsIsDecodedAsV2ErrorWithSameExceptions()
+ throws IOException {
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.ERROR);
+ v1MetadataBlock.addException(250, "timeout");
+ v1MetadataBlock.addException(500, "server error");
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR, "Expected error type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), v1MetadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V2 EOS with empty stats is read in V1 as EOS without stats
+ */
+ @Test
+ public void v2EosWithoutStatsIsReadInV1AsEosWithoutStats()
+ throws IOException {
+ ByteBuffer stats = ByteBuffer.wrap(new byte[]{0, 0, 0, 0});
+ MetadataBlock metadataBlock = new MetadataBlock(Lists.newArrayList(stats));
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected no exceptions");
}
+ /**
+ * Verifies that a V2 EOS with stats is read in V1 as EOS without stats
+ */
@Test
- public void shouldDefaultToEosWithNoErrorsOnLegacyMetadataBlock()
+ public void v2EosWithStatsIsReadInV1AsEosWithoutStats()
+ throws IOException {
+ MetadataBlock metadataBlock = MetadataBlock.newEos();
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected no exceptions");
+ }
+
+ /**
+ * Verifies that a V2 error code is read in V1 as error with the same exceptions
+ */
+ @Test
+ public void v2ErrorIsReadInV1AsErrorWithSameExceptions()
+ throws IOException {
+ HashMap<Integer, String> errorMap = new HashMap<>();
+ errorMap.put(250, "timeout");
+ MetadataBlock metadataBlock = MetadataBlock.newError(errorMap);
+
+ byte[] bytes = metadataBlock.toBytes();
+
+ // This is how V1 blocks were deserialized
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(buff);
+
+ assertEquals(v1MetadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR, "Expected error type");
+ assertEquals(v1MetadataBlock.getStats(), Collections.emptyMap(), "Expected no stats by stage");
+ assertEquals(v1MetadataBlock.getExceptions(), metadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V1 error code is decoded as a V2 error code with the same exceptions
+ */
+ @Test
+ public void v1EosWithoutStatsIsDecodedAsV2EosWithoutStats()
+ throws IOException {
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.EOS, new HashMap<>());
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), v1MetadataBlock.getExceptions(), "Expected exceptions");
+ }
+
+ /**
+ * Verifies that a V1 EOS with stats is decoded as a V2 EOS without stats
+ */
+ @Test
+ public void v1EosWithStatsIsDecodedAsV2EosWithoutStats()
+ throws IOException {
+ HashMap<String, String> stats = new HashMap<>();
+ stats.put("foo", "bar");
+ stats.put("baz", "qux");
+ V1MetadataBlock v1MetadataBlock = new V1MetadataBlock(MetadataBlock.MetadataBlockType.EOS, stats);
+
+ byte[] bytes = v1MetadataBlock.toBytes();
+ ByteBuffer buff = ByteBuffer.wrap(bytes);
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(buff);
+
+ assertTrue(dataBlock instanceof MetadataBlock, "V1MetadataBlock should be always decoded as MetadataBlock");
+ MetadataBlock metadataBlock = (MetadataBlock) dataBlock;
+ assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.EOS, "Expected EOS type");
+ assertEquals(metadataBlock.getStatsByStage(), Collections.emptyList(), "Expected no stats by stage");
+ assertEquals(metadataBlock.getExceptions(), Collections.emptyMap(), "Expected no exceptions");
+ }
+
+ /**
+ * Verifies that a V0 EOS without exceptions is decoded as a V2 EOS without stats
+ */
+ @Test
+ public void v0EosWithoutExceptionsIsDecodedAsV2EosWithoutStats()
throws IOException {
// Given:
// MetadataBlock used to be encoded without any data, we should make sure that
// during rollout or if server versions are mismatched that we can still handle
// the old format
- OldMetadataBlock legacyBlock = new OldMetadataBlock();
+ V0MetadataBlock legacyBlock = new V0MetadataBlock();
byte[] bytes = legacyBlock.toBytes();
// When:
ByteBuffer buff = ByteBuffer.wrap(bytes);
- buff.getInt(); // consume the version information before decoding
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
MetadataBlock metadataBlock = new MetadataBlock(buff);
// Then:
@@ -62,29 +204,30 @@ public class MetadataBlockTest {
}
@Test
- public void shouldDefaultToErrorOnLegacyMetadataBlockWithErrors()
+ public void v0EosWithExceptionsIsDecodedAsV2ErrorWithSameExceptions()
throws IOException {
// Given:
// MetadataBlock used to be encoded without any data, we should make sure that
// during rollout or if server versions are mismatched that we can still handle
// the old format
- OldMetadataBlock legacyBlock = new OldMetadataBlock();
+ V0MetadataBlock legacyBlock = new V0MetadataBlock();
legacyBlock.addException(250, "timeout");
byte[] bytes = legacyBlock.toBytes();
// When:
ByteBuffer buff = ByteBuffer.wrap(bytes);
- buff.getInt(); // consume the version information before decoding
+ DataBlockUtils.readVersionType(buff); // consume the version information before decoding
MetadataBlock metadataBlock = new MetadataBlock(buff);
// Then:
assertEquals(metadataBlock.getType(), MetadataBlock.MetadataBlockType.ERROR);
+ assertEquals(metadataBlock.getExceptions(), legacyBlock.getExceptions(), "Expected exceptions");
}
@Test(expectedExceptions = UnsupportedOperationException.class)
public void shouldThrowExceptionWhenUsingReadMethods() {
// Given:
- MetadataBlock block = new MetadataBlock(MetadataBlock.MetadataBlockType.EOS);
+ MetadataBlock block = MetadataBlock.newEos();
// When:
// (should through exception)
@@ -94,9 +237,9 @@ public class MetadataBlockTest {
/**
* This is mostly just used as an internal serialization tool
*/
- private static class OldMetadataBlock extends BaseDataBlock {
+ private static class V0MetadataBlock extends BaseDataBlock {
- public OldMetadataBlock() {
+ public V0MetadataBlock() {
super(0, null, new String[0], new byte[0], new byte[0]);
}
@@ -114,15 +257,5 @@ public class MetadataBlockTest {
protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
return 0;
}
-
- @Override
- public DataBlock toMetadataOnlyDataTable() {
- return null;
- }
-
- @Override
- public DataBlock toDataOnlyDataTable() {
- return null;
- }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java
similarity index 53%
copy from pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java
copy to pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java
index 3415bfb3fc..b3bbe26ff6 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/RowDataBlockTest.java
@@ -16,34 +16,25 @@
* specific language governing permissions and limitations
* under the License.
*/
-package org.apache.pinot.query.runtime.operator;
+package org.apache.pinot.common.datablock;
-import java.util.List;
-import javax.annotation.Nullable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
import org.apache.pinot.common.utils.DataSchema;
-import org.apache.pinot.core.data.table.Record;
-import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.testng.annotations.Test;
-/**
- * Minus/Except operator.
- */
-public class MinusOperator extends SetOperator {
- private static final String EXPLAIN_NAME = "MINUS";
-
- public MinusOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
- DataSchema dataSchema) {
- super(opChainExecutionContext, upstreamOperators, dataSchema);
- }
-
- @Nullable
+public class RowDataBlockTest extends BaseDataBlockContract {
@Override
- public String toExplainString() {
- return EXPLAIN_NAME;
+ protected BaseDataBlock deserialize(ByteBuffer byteBuffer, int versionType)
+ throws IOException {
+ return new RowDataBlock(byteBuffer);
}
- @Override
- protected boolean handleRowMatched(Object[] row) {
- return _rightRowSet.add(new Record(row));
+ @Test
+ public void emptyDataBlockCorrectness()
+ throws IOException {
+ DataSchema dataSchema = new DataSchema(new String[0], new DataSchema.ColumnDataType[0]);
+ testSerdeCorrectness(new RowDataBlock(0, dataSchema, new String[0], new byte[0], new byte[0]));
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java b/pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java
similarity index 77%
copy from pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
copy to pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java
index 41c2e9d0ea..990d9489cc 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/MetadataBlock.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datablock/V1MetadataBlock.java
@@ -31,31 +31,24 @@ import java.util.Map;
/**
- * A block type to indicate some metadata about the current processing state.
- * For the different types of metadata blocks see {@link MetadataBlockType}.
+ * The datablock used before the introduction of {@link org.apache.pinot.common.datatable.StatMap}.
+ * <p>
+ * This version stored the metadata in a {@code Map<String, String>} which was encoded as JSON and stored in the
+ * variable size data buffer.
+ * <p>
+ * Instances of this class are not actually seen by the operators. Instead, they are converted to {@link MetadataBlock}
+ * in {@link MetadataBlock#deserialize(ByteBuffer, int)}.
+ * <p>
+ * The reason to keep it here is mostly for backwards compatibility and testing. In order to simplify the code, the
+ * stats engine just ignores the metadata of these objects, but we need to be able to deserialize them anyway.
*/
-public class MetadataBlock extends BaseDataBlock {
+public class V1MetadataBlock extends BaseDataBlock {
private static final ObjectMapper JSON = new ObjectMapper();
@VisibleForTesting
static final int VERSION = 1;
- public enum MetadataBlockType {
- /**
- * Indicates that this block is the final block to be sent
- * (End Of Stream) as part of an operator chain computation.
- */
- EOS,
-
- /**
- * An {@code ERROR} metadata block indicates that there was
- * some error during computation. To retrieve the error that
- * occurred, use {@link MetadataBlock#getExceptions()}
- */
- ERROR
- }
-
/**
* Used to serialize the contents of the metadata block conveniently and in
* a backwards compatible way. Use JSON because the performance of metadata block
@@ -98,11 +91,11 @@ public class MetadataBlock extends BaseDataBlock {
private final Contents _contents;
- public MetadataBlock(MetadataBlockType type) {
+ public V1MetadataBlock(MetadataBlock.MetadataBlockType type) {
this(type, new HashMap<>());
}
- public MetadataBlock(MetadataBlockType type, Map<String, String> stats) {
+ public V1MetadataBlock(MetadataBlock.MetadataBlockType type, Map<String, String> stats) {
super(0, null, new String[0], new byte[]{0}, toContents(new Contents(type.name(), stats)));
_contents = new Contents(type.name(), stats);
}
@@ -115,7 +108,7 @@ public class MetadataBlock extends BaseDataBlock {
}
}
- public MetadataBlock(ByteBuffer byteBuffer)
+ public V1MetadataBlock(ByteBuffer byteBuffer)
throws IOException {
super(byteBuffer);
if (_variableSizeDataBytes != null && _variableSizeDataBytes.length > 0) {
@@ -125,15 +118,15 @@ public class MetadataBlock extends BaseDataBlock {
}
}
- public MetadataBlockType getType() {
+ public MetadataBlock.MetadataBlockType getType() {
String type = _contents.getType();
// if type is null, then we're reading a legacy block where we didn't encode any
// data. assume that it is an EOS block if there's no exceptions and an ERROR block
// otherwise
return type == null
- ? (getExceptions().isEmpty() ? MetadataBlockType.EOS : MetadataBlockType.ERROR)
- : MetadataBlockType.valueOf(type);
+ ? (getExceptions().isEmpty() ? MetadataBlock.MetadataBlockType.EOS : MetadataBlock.MetadataBlockType.ERROR)
+ : MetadataBlock.MetadataBlockType.valueOf(type);
}
public Map<String, String> getStats() {
@@ -154,14 +147,4 @@ public class MetadataBlock extends BaseDataBlock {
protected int positionOffsetInVariableBufferAndGetLength(int rowId, int colId) {
throw new UnsupportedOperationException("Metadata block uses JSON encoding for field access");
}
-
- @Override
- public MetadataBlock toMetadataOnlyDataTable() {
- return this;
- }
-
- @Override
- public MetadataBlock toDataOnlyDataTable() {
- throw new UnsupportedOperationException();
- }
}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java b/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java
new file mode 100644
index 0000000000..25f5ef8a07
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/datatable/StatMapTest.java
@@ -0,0 +1,209 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.datatable;
+
+import com.google.common.io.ByteArrayDataInput;
+import com.google.common.io.ByteArrayDataOutput;
+import com.google.common.io.ByteStreams;
+import java.io.IOException;
+import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
+import org.testng.Assert;
+import org.testng.SkipException;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+
+public class StatMapTest {
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckWhenAddInt(MyStats stat) {
+ if (stat.getType() == StatMap.Type.INT) {
+ throw new SkipException("Skipping INT test");
+ }
+ if (stat.getType() == StatMap.Type.LONG) {
+ throw new SkipException("Skipping LONG test");
+ }
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, 1);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckWhenAddLong(MyStats stat) {
+ if (stat.getType() == StatMap.Type.LONG) {
+ throw new SkipException("Skipping LONG test");
+ }
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, 1L);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckPutBoolean(MyStats stat) {
+ if (stat.getType() == StatMap.Type.BOOLEAN) {
+ throw new SkipException("Skipping BOOLEAN test");
+ }
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, true);
+ }
+
+ @Test(dataProvider = "allTypeStats", expectedExceptions = IllegalArgumentException.class)
+ public void dynamicTypeCheckPutString(MyStats stat) {
+ if (stat.getType() == StatMap.Type.STRING) {
+ throw new SkipException("Skipping STRING test");
+ }
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ statMap.merge(stat, "foo");
+ }
+
+ @Test(dataProvider = "allTypeStats")
+ public void singleEncodeDecode(MyStats stat)
+ throws IOException {
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ switch (stat.getType()) {
+ case BOOLEAN:
+ statMap.merge(stat, true);
+ break;
+ case INT:
+ statMap.merge(stat, 1);
+ break;
+ case LONG:
+ statMap.merge(stat, 1L);
+ break;
+ case STRING:
+ statMap.merge(stat, "foo");
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ testSerializeDeserialize(statMap);
+ }
+
+ @Test
+ public void encodeDecodeAll()
+ throws IOException {
+ StatMap<MyStats> statMap = new StatMap<>(MyStats.class);
+ for (MyStats stat : MyStats.values()) {
+ switch (stat.getType()) {
+ case BOOLEAN:
+ statMap.merge(stat, true);
+ break;
+ case INT:
+ statMap.merge(stat, 1);
+ break;
+ case LONG:
+ statMap.merge(stat, 1L);
+ break;
+ case STRING:
+ statMap.merge(stat, "foo");
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ }
+ testSerializeDeserialize(statMap);
+ }
+
+ private <K extends Enum<K> & StatMap.Key> void testSerializeDeserialize(StatMap<K> statMap)
+ throws IOException {
+ ByteArrayDataOutput output = ByteStreams.newDataOutput();
+ statMap.serialize(output);
+
+ ByteArrayDataInput input = ByteStreams.newDataInput(output.toByteArray());
+ StatMap<MyStats> deserializedStatMap = StatMap.deserialize(input, MyStats.class);
+
+ Assert.assertEquals(statMap, deserializedStatMap);
+ }
+
+ @Test(dataProvider = "complexStats")
+ public void mergeEquivalence(StatMap<?> statMap)
+ throws IOException {
+ StatMap mergedOnHeap = new StatMap<>(statMap.getKeyClass());
+ mergedOnHeap.merge(statMap);
+
+ ByteArrayDataOutput output = ByteStreams.newDataOutput();
+ statMap.serialize(output);
+
+ ByteArrayDataInput input = ByteStreams.newDataInput(output.toByteArray());
+ StatMap mergedSerialized = new StatMap<>(statMap.getKeyClass());
+ mergedSerialized.merge(input);
+
+ Assert.assertEquals(mergedOnHeap, mergedSerialized,
+ "Merging objects should be equal to merging serialized buffers");
+ }
+
+ @DataProvider(name = "complexStats")
+ static StatMap<?>[] complexStats() {
+ return new StatMap<?>[] {
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, true)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, true)
+ .merge(MyStats.LONG_KEY, 0L)
+ .merge(MyStats.INT_KEY, 1)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 0)
+ .merge(MyStats.STR_KEY, "foo"),
+ new StatMap<>(MyStats.class)
+ .merge(MyStats.BOOL_KEY, false)
+ .merge(MyStats.LONG_KEY, 1L)
+ .merge(MyStats.INT_KEY, 1),
+ new StatMap<>(BrokerResponseNativeV2.StatKey.class)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_QUERIED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_PROCESSED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_SEGMENTS_MATCHED, 1)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_DOCS_SCANNED, 10)
+ .merge(BrokerResponseNativeV2.StatKey.NUM_ENTRIES_SCANNED_POST_FILTER, 5)
+ .merge(BrokerResponseNativeV2.StatKey.TOTAL_DOCS, 5)
+ .merge(BrokerResponseNativeV2.StatKey.TIME_USED_MS, 95)
+ };
+ }
+
+ @DataProvider(name = "allTypeStats")
+ static MyStats[] allTypeStats() {
+ return MyStats.values();
+ }
+
+ public enum MyStats implements StatMap.Key {
+ BOOL_KEY(StatMap.Type.BOOLEAN),
+ LONG_KEY(StatMap.Type.LONG),
+ INT_KEY(StatMap.Type.INT),
+ STR_KEY(StatMap.Type.STRING);
+
+ private final StatMap.Type _type;
+
+ MyStats(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
index 0ef9bb2207..5ee75e5bf1 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/InstanceResponseOperator.java
@@ -45,6 +45,9 @@ public class InstanceResponseOperator extends BaseOperator<InstanceResponseBlock
protected final int _fetchContextSize;
protected final QueryContext _queryContext;
+ protected long _threadCpuTimeNs;
+ protected long _systemActivitiesCpuTimeNs;
+
public InstanceResponseOperator(BaseCombineOperator<?> combineOperator, List<SegmentContext> segmentContexts,
List<FetchContext> fetchContexts, QueryContext queryContext) {
_combineOperator = combineOperator;
@@ -81,12 +84,24 @@ public class InstanceResponseOperator extends BaseOperator<InstanceResponseBlock
@Override
protected InstanceResponseBlock getNextBlock() {
+ BaseResultsBlock baseResultsBlock = getBaseBlock();
+ return buildInstanceResponseBlock(baseResultsBlock);
+ }
+
+ protected InstanceResponseBlock buildInstanceResponseBlock(BaseResultsBlock baseResultsBlock) {
+ InstanceResponseBlock instanceResponseBlock = new InstanceResponseBlock(baseResultsBlock);
+ instanceResponseBlock.addMetadata(MetadataKey.THREAD_CPU_TIME_NS.getName(), String.valueOf(_threadCpuTimeNs));
+ instanceResponseBlock.addMetadata(MetadataKey.SYSTEM_ACTIVITIES_CPU_TIME_NS.getName(),
+ String.valueOf(_systemActivitiesCpuTimeNs));
+ return instanceResponseBlock;
+ }
+
+ protected BaseResultsBlock getBaseBlock() {
if (ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled()) {
long startWallClockTimeNs = System.nanoTime();
ThreadResourceUsageProvider mainThreadResourceUsageProvider = new ThreadResourceUsageProvider();
BaseResultsBlock resultsBlock = getCombinedResults();
- InstanceResponseBlock instanceResponseBlock = new InstanceResponseBlock(resultsBlock);
long mainThreadCpuTimeNs = mainThreadResourceUsageProvider.getThreadTimeNs();
long totalWallClockTimeNs = System.nanoTime() - startWallClockTimeNs;
@@ -97,22 +112,17 @@ public class InstanceResponseOperator extends BaseOperator<InstanceResponseBlock
*/
long multipleThreadCpuTimeNs = resultsBlock.getExecutionThreadCpuTimeNs();
int numServerThreads = resultsBlock.getNumServerThreads();
- long systemActivitiesCpuTimeNs =
- calSystemActivitiesCpuTimeNs(totalWallClockTimeNs, multipleThreadCpuTimeNs, mainThreadCpuTimeNs,
- numServerThreads);
-
- long threadCpuTimeNs = mainThreadCpuTimeNs + multipleThreadCpuTimeNs;
- instanceResponseBlock.addMetadata(MetadataKey.THREAD_CPU_TIME_NS.getName(), String.valueOf(threadCpuTimeNs));
- instanceResponseBlock.addMetadata(MetadataKey.SYSTEM_ACTIVITIES_CPU_TIME_NS.getName(),
- String.valueOf(systemActivitiesCpuTimeNs));
+ _systemActivitiesCpuTimeNs = calSystemActivitiesCpuTimeNs(totalWallClockTimeNs, multipleThreadCpuTimeNs,
+ mainThreadCpuTimeNs, numServerThreads);
+ _threadCpuTimeNs = mainThreadCpuTimeNs + multipleThreadCpuTimeNs;
- return instanceResponseBlock;
+ return resultsBlock;
} else {
- return new InstanceResponseBlock(getCombinedResults());
+ return getCombinedResults();
}
}
- private BaseResultsBlock getCombinedResults() {
+ protected BaseResultsBlock getCombinedResults() {
try {
prefetchAll();
return _combineOperator.nextBlock();
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
index 61a5b8cc83..1130efb4d7 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingInstanceResponseOperator.java
@@ -36,6 +36,12 @@ import org.apache.pinot.spi.exception.QueryCancelledException;
import org.apache.pinot.spi.trace.Tracing;
+/**
+ * Like {@link InstanceResponseOperator}, but instead of sending all the data to the broker at once, it streams the data
+ * to a given {@link ResultsBlockStreamer}.
+ *
+ * This is used in multi-stage to stream data to the receiving mailboxes.
+ */
public class StreamingInstanceResponseOperator extends InstanceResponseOperator {
private static final String EXPLAIN_NAME = "STREAMING_INSTANCE_RESPONSE";
@@ -58,7 +64,7 @@ public class StreamingInstanceResponseOperator extends InstanceResponseOperator
prefetchAll();
if (_streamingCombineOperator != null) {
_streamingCombineOperator.start();
- BaseResultsBlock resultsBlock = _streamingCombineOperator.nextBlock();
+ BaseResultsBlock resultsBlock = getBaseBlock();
while (!(resultsBlock instanceof MetadataResultsBlock)) {
if (resultsBlock instanceof ExceptionResultsBlock) {
return new InstanceResponseBlock(resultsBlock);
@@ -66,20 +72,20 @@ public class StreamingInstanceResponseOperator extends InstanceResponseOperator
if (resultsBlock.getNumRows() > 0) {
_streamer.send(resultsBlock);
}
- resultsBlock = _streamingCombineOperator.nextBlock();
+ resultsBlock = getBaseBlock();
}
// Return a metadata-only block in the end
- return new InstanceResponseBlock(resultsBlock);
+ return buildInstanceResponseBlock(resultsBlock);
} else {
// Handle single block combine operator in streaming fashion
- BaseResultsBlock resultsBlock = _combineOperator.nextBlock();
+ BaseResultsBlock resultsBlock = getBaseBlock();
if (resultsBlock instanceof ExceptionResultsBlock) {
return new InstanceResponseBlock(resultsBlock);
}
if (resultsBlock.getNumRows() > 0) {
_streamer.send(resultsBlock);
}
- return new InstanceResponseBlock(resultsBlock).toMetadataOnlyResponseBlock();
+ return buildInstanceResponseBlock(resultsBlock).toMetadataOnlyResponseBlock();
}
} catch (EarlyTerminationException e) {
Exception killedErrorMsg = Tracing.getThreadAccountant().getErrorStatus();
@@ -96,6 +102,10 @@ public class StreamingInstanceResponseOperator extends InstanceResponseOperator
}
}
+ protected BaseResultsBlock getCombinedResults() {
+ return _combineOperator.nextBlock();
+ }
+
@Override
public String toExplainString() {
return EXPLAIN_NAME;
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
index c05644af5f..c40133baab 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/ExecutionStatsAggregator.java
@@ -34,7 +34,6 @@ import org.apache.pinot.common.metrics.BrokerMeter;
import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.metrics.BrokerTimer;
import org.apache.pinot.common.response.broker.BrokerResponseNative;
-import org.apache.pinot.common.response.broker.BrokerResponseStats;
import org.apache.pinot.common.response.broker.QueryProcessingException;
import org.apache.pinot.core.transport.ServerRoutingInstance;
import org.apache.pinot.spi.config.table.TableType;
@@ -74,7 +73,6 @@ public class ExecutionStatsAggregator {
private long _explainPlanNumEmptyFilterSegments = 0L;
private long _explainPlanNumMatchAllFilterSegments = 0L;
private boolean _numGroupsLimitReached = false;
- private boolean _maxRowsInJoinReached = false;
private int _numBlocks = 0;
private int _numRows = 0;
private long _stageExecutionTimeMs = 0;
@@ -248,12 +246,10 @@ public class ExecutionStatsAggregator {
if (numTotalDocsString != null) {
_numTotalDocs += Long.parseLong(numTotalDocsString);
}
+
_numGroupsLimitReached |=
Boolean.parseBoolean(metadata.get(DataTable.MetadataKey.NUM_GROUPS_LIMIT_REACHED.getName()));
- _maxRowsInJoinReached |=
- Boolean.parseBoolean(metadata.get(DataTable.MetadataKey.MAX_ROWS_IN_JOIN_REACHED.getName()));
-
String numBlocksString = metadata.get(DataTable.MetadataKey.NUM_BLOCKS.getName());
if (numBlocksString != null) {
_numBlocks += Long.parseLong(numBlocksString);
@@ -309,15 +305,12 @@ public class ExecutionStatsAggregator {
brokerResponseNative.setNumSegmentsMatched(_numSegmentsMatched);
brokerResponseNative.setTotalDocs(_numTotalDocs);
brokerResponseNative.setNumGroupsLimitReached(_numGroupsLimitReached);
- brokerResponseNative.setMaxRowsInJoinReached(_maxRowsInJoinReached);
brokerResponseNative.setOfflineThreadCpuTimeNs(_offlineThreadCpuTimeNs);
brokerResponseNative.setRealtimeThreadCpuTimeNs(_realtimeThreadCpuTimeNs);
brokerResponseNative.setOfflineSystemActivitiesCpuTimeNs(_offlineSystemActivitiesCpuTimeNs);
brokerResponseNative.setRealtimeSystemActivitiesCpuTimeNs(_realtimeSystemActivitiesCpuTimeNs);
brokerResponseNative.setOfflineResponseSerializationCpuTimeNs(_offlineResponseSerializationCpuTimeNs);
brokerResponseNative.setRealtimeResponseSerializationCpuTimeNs(_realtimeResponseSerializationCpuTimeNs);
- brokerResponseNative.setOfflineTotalCpuTimeNs(_offlineTotalCpuTimeNs);
- brokerResponseNative.setRealtimeTotalCpuTimeNs(_realtimeTotalCpuTimeNs);
brokerResponseNative.setNumSegmentsPrunedByServer(_numSegmentsPrunedByServer);
brokerResponseNative.setNumSegmentsPrunedInvalid(_numSegmentsPrunedInvalid);
brokerResponseNative.setNumSegmentsPrunedByLimit(_numSegmentsPrunedByLimit);
@@ -364,25 +357,6 @@ public class ExecutionStatsAggregator {
}
}
- public void setStageLevelStats(@Nullable String rawTableName, BrokerResponseStats brokerResponseStats,
- @Nullable BrokerMetrics brokerMetrics) {
- if (_enableTrace) {
- setStats(rawTableName, brokerResponseStats, brokerMetrics);
- brokerResponseStats.setOperatorStats(_operatorStats);
- }
-
- brokerResponseStats.setNumBlocks(_numBlocks);
- brokerResponseStats.setNumRows(_numRows);
- brokerResponseStats.setMaxRowsInJoinReached(_maxRowsInJoinReached);
- brokerResponseStats.setNumGroupsLimitReached(_numGroupsLimitReached);
- brokerResponseStats.setStageExecutionTimeMs(_stageExecutionTimeMs);
- brokerResponseStats.setStageExecutionUnit(_stageExecutionUnit);
- brokerResponseStats.setTableNames(new ArrayList<>(_tableNames));
- if (_stageExecStartTimeMs >= 0 && _stageExecEndTimeMs >= 0) {
- brokerResponseStats.setStageExecWallTimeMs(_stageExecEndTimeMs - _stageExecStartTimeMs);
- }
- }
-
private void withNotNullLongMetadata(Map<String, String> metadata, DataTable.MetadataKey key, LongConsumer consumer) {
String strValue = metadata.get(key.getName());
if (strValue != null) {
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
index e7c5af0dab..6bc1f0cbda 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
@@ -406,7 +406,7 @@ public abstract class BaseClusterIntegrationTestSet extends BaseClusterIntegrati
for (String query : queries) {
JsonNode response = postQuery(query);
for (String statName : statNames) {
- assertTrue(response.has(statName));
+ assertTrue(response.has(statName), "Response does not contain stat: " + statName);
}
}
}
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
index f66c194d85..a54e5314f7 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/HybridClusterIntegrationTest.java
@@ -317,7 +317,8 @@ public class HybridClusterIntegrationTest extends BaseClusterIntegrationTestSet
final String resultTag = "resultTable";
// dropResults=true - resultTable must not be in the response
- Assert.assertFalse(postQueryWithOptions(query, "dropResults=true").has(resultTag));
+ JsonNode jsonNode = postQueryWithOptions(query, "dropResults=true");
+ Assert.assertFalse(jsonNode.has(resultTag));
// dropResults=TrUE (case insensitive match) - resultTable must not be in the response
Assert.assertFalse(postQueryWithOptions(query, "dropResults=TrUE").has(resultTag));
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
index 5398e58b77..0dc2dc5884 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotJoinToDynamicBroadcastRule.java
@@ -24,6 +24,7 @@ import java.util.List;
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.plan.hep.HepRelVertex;
+import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelDistributions;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Exchange;
@@ -164,11 +165,15 @@ public class PinotJoinToDynamicBroadcastRule extends RelOptRule {
// 2. when hash key are the same but hash functions are different, it can be done via normal hash shuffle.
boolean isColocatedJoin = PinotHintStrategyTable.isHintOptionTrue(join.getHints(),
PinotHintOptions.JOIN_HINT_OPTIONS, PinotHintOptions.JoinHintOptions.IS_COLOCATED_BY_JOIN_KEYS);
- PinotLogicalExchange dynamicBroadcastExchange = isColocatedJoin
- ? PinotLogicalExchange.create(right.getInput(), RelDistributions.hash(join.analyzeCondition().rightKeys),
- PinotRelExchangeType.PIPELINE_BREAKER)
- : PinotLogicalExchange.create(right.getInput(), RelDistributions.BROADCAST_DISTRIBUTED,
- PinotRelExchangeType.PIPELINE_BREAKER);
+ PinotLogicalExchange dynamicBroadcastExchange;
+ RelNode rightInput = right.getInput();
+ if (isColocatedJoin) {
+ RelDistribution dist = RelDistributions.hash(join.analyzeCondition().rightKeys);
+ dynamicBroadcastExchange = PinotLogicalExchange.create(rightInput, dist, PinotRelExchangeType.PIPELINE_BREAKER);
+ } else {
+ RelDistribution dist = RelDistributions.BROADCAST_DISTRIBUTED;
+ dynamicBroadcastExchange = PinotLogicalExchange.create(rightInput, dist, PinotRelExchangeType.PIPELINE_BREAKER);
+ }
Join dynamicFilterJoin =
new LogicalJoin(join.getCluster(), join.getTraitSet(), left.getInput(), dynamicBroadcastExchange,
join.getCondition(), join.getVariablesSet(), join.getJoinType(), join.isSemiJoinDone(),
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
index 87a8bda603..e2bb53faf5 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java
@@ -24,12 +24,14 @@ import io.grpc.stub.StreamObserver;
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.proto.Mailbox.MailboxContent;
import org.apache.pinot.common.proto.PinotMailboxGrpc;
import org.apache.pinot.query.mailbox.channel.ChannelManager;
import org.apache.pinot.query.mailbox.channel.MailboxStatusObserver;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -45,16 +47,19 @@ public class GrpcSendingMailbox implements SendingMailbox {
private final String _hostname;
private final int _port;
private final long _deadlineMs;
+ private final StatMap<MailboxSendOperator.StatKey> _statMap;
private final MailboxStatusObserver _statusObserver = new MailboxStatusObserver();
private StreamObserver<MailboxContent> _contentObserver;
- public GrpcSendingMailbox(String id, ChannelManager channelManager, String hostname, int port, long deadlineMs) {
+ public GrpcSendingMailbox(String id, ChannelManager channelManager, String hostname, int port, long deadlineMs,
+ StatMap<MailboxSendOperator.StatKey> statMap) {
_id = id;
_channelManager = channelManager;
_hostname = hostname;
_port = port;
_deadlineMs = deadlineMs;
+ _statMap = statMap;
}
@Override
@@ -64,25 +69,31 @@ public class GrpcSendingMailbox implements SendingMailbox {
return;
}
if (LOGGER.isDebugEnabled()) {
- LOGGER.debug("==[GRPC SEND]== sending data to: " + _id);
+ LOGGER.debug("==[GRPC SEND]== sending message " + block + " to: " + _id);
}
if (_contentObserver == null) {
_contentObserver = getContentObserver();
}
_contentObserver.onNext(toMailboxContent(block));
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("==[GRPC SEND]== message " + block + " sent to: " + _id);
+ }
}
@Override
public void complete() {
if (isTerminated()) {
+ LOGGER.debug("Already terminated mailbox: {}", _id);
return;
}
+ LOGGER.debug("Completing mailbox: {}", _id);
_contentObserver.onCompleted();
}
@Override
public void cancel(Throwable t) {
if (isTerminated()) {
+ LOGGER.debug("Already terminated mailbox: {}", _id);
return;
}
LOGGER.debug("Cancelling mailbox: {}", _id);
@@ -118,9 +129,22 @@ public class GrpcSendingMailbox implements SendingMailbox {
private MailboxContent toMailboxContent(TransferableBlock block)
throws IOException {
- DataBlock dataBlock = block.getDataBlock();
- byte[] bytes = dataBlock.toBytes();
- ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
- return MailboxContent.newBuilder().setMailboxId(_id).setPayload(byteString).build();
+ _statMap.merge(MailboxSendOperator.StatKey.RAW_MESSAGES, 1);
+ long start = System.currentTimeMillis();
+ try {
+ DataBlock dataBlock = block.getDataBlock();
+ byte[] bytes = dataBlock.toBytes();
+ ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Serialized block: {} to {} bytes", block, bytes.length);
+ }
+ _statMap.merge(MailboxSendOperator.StatKey.SERIALIZED_BYTES, bytes.length);
+ return MailboxContent.newBuilder().setMailboxId(_id).setPayload(byteString).build();
+ } catch (Throwable t) {
+ LOGGER.warn("Caught exception while serializing block: {}", block, t);
+ throw t;
+ } finally {
+ _statMap.merge(MailboxSendOperator.StatKey.SERIALIZATION_TIME_MS, System.currentTimeMillis() - start);
+ }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
index d1452712bc..8adf8db073 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java
@@ -19,8 +19,10 @@
package org.apache.pinot.query.mailbox;
import java.util.concurrent.TimeoutException;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.spi.exception.QueryCancelledException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -36,11 +38,14 @@ public class InMemorySendingMailbox implements SendingMailbox {
private ReceivingMailbox _receivingMailbox;
private volatile boolean _isTerminated;
private volatile boolean _isEarlyTerminated;
+ private final StatMap<MailboxSendOperator.StatKey> _statMap;
- public InMemorySendingMailbox(String id, MailboxService mailboxService, long deadlineMs) {
+ public InMemorySendingMailbox(String id, MailboxService mailboxService, long deadlineMs,
+ StatMap<MailboxSendOperator.StatKey> statMap) {
_id = id;
_mailboxService = mailboxService;
_deadlineMs = deadlineMs;
+ _statMap = statMap;
}
@Override
@@ -52,8 +57,10 @@ public class InMemorySendingMailbox implements SendingMailbox {
if (_receivingMailbox == null) {
_receivingMailbox = _mailboxService.getReceivingMailbox(_id);
}
+ _statMap.merge(MailboxSendOperator.StatKey.IN_MEMORY_MESSAGES, 1);
long timeoutMs = _deadlineMs - System.currentTimeMillis();
ReceivingMailbox.ReceivingMailboxStatus status = _receivingMailbox.offer(block, timeoutMs);
+
switch (status) {
case SUCCESS:
break;
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
index 7c171f1bf4..c71af0386d 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
@@ -23,8 +23,10 @@ import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalListener;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.mailbox.channel.ChannelManager;
import org.apache.pinot.query.mailbox.channel.GrpcMailboxServer;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -99,11 +101,12 @@ public class MailboxService {
* not open the underlying channel or acquire any additional resources. Instead, it will initialize lazily when the
* data is sent for the first time.
*/
- public SendingMailbox getSendingMailbox(String hostname, int port, String mailboxId, long deadlineMs) {
+ public SendingMailbox getSendingMailbox(String hostname, int port, String mailboxId, long deadlineMs,
+ StatMap<MailboxSendOperator.StatKey> statMap) {
if (_hostname.equals(hostname) && _port == port) {
- return new InMemorySendingMailbox(mailboxId, this, deadlineMs);
+ return new InMemorySendingMailbox(mailboxId, this, deadlineMs, statMap);
} else {
- return new GrpcSendingMailbox(mailboxId, _channelManager, hostname, port, deadlineMs);
+ return new GrpcSendingMailbox(mailboxId, _channelManager, hostname, port, deadlineMs, statMap);
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
index 97c8731f95..e58b5f7dd5 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
@@ -19,12 +19,19 @@
package org.apache.pinot.query.mailbox;
import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
+import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datablock.DataBlockUtils;
+import org.apache.pinot.common.datablock.MetadataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.slf4j.Logger;
@@ -50,9 +57,11 @@ public class ReceivingMailbox {
private final BlockingQueue<TransferableBlock> _blocks = new ArrayBlockingQueue<>(DEFAULT_MAX_PENDING_BLOCKS);
private final AtomicReference<TransferableBlock> _errorBlock = new AtomicReference<>();
private volatile boolean _isEarlyTerminated = false;
+ private long _lastArriveTime = System.currentTimeMillis();
@Nullable
private volatile Reader _reader;
+ private final StatMap<StatKey> _stats = new StatMap<>(StatKey.class);
public ReceivingMailbox(String id) {
_id = id;
@@ -72,11 +81,53 @@ public class ReceivingMailbox {
return _id;
}
+ /**
+ * Offers a raw block into the mailbox within the timeout specified, returns whether the block is successfully added.
+ * If the block is not added, an error block is added to the mailbox.
+ * <p>
+ * Contrary to {@link #offer(TransferableBlock, long)}, the block may be an
+ * {@link TransferableBlock#isErrorBlock() error block}.
+ */
+ public ReceivingMailboxStatus offerRaw(ByteBuffer byteBuffer, long timeoutMs)
+ throws IOException {
+ TransferableBlock block;
+ long now = System.currentTimeMillis();
+ _stats.merge(StatKey.WAIT_CPU_TIME_MS, now - _lastArriveTime);
+ _lastArriveTime = now;
+ _stats.merge(StatKey.DESERIALIZED_BYTES, byteBuffer.remaining());
+ _stats.merge(StatKey.DESERIALIZED_MESSAGES, 1);
+
+ now = System.currentTimeMillis();
+ DataBlock dataBlock = DataBlockUtils.getDataBlock(byteBuffer);
+ _stats.merge(StatKey.DESERIALIZATION_TIME_MS, System.currentTimeMillis() - now);
+
+ if (dataBlock instanceof MetadataBlock) {
+ Map<Integer, String> exceptions = dataBlock.getExceptions();
+ if (exceptions.isEmpty()) {
+ block = TransferableBlockUtils.wrap(dataBlock);
+ } else {
+ setErrorBlock(TransferableBlockUtils.getErrorTransferableBlock(exceptions));
+ return ReceivingMailboxStatus.FIRST_ERROR;
+ }
+ } else {
+ block = TransferableBlockUtils.wrap(dataBlock);
+ }
+ return offerPrivate(block, timeoutMs);
+ }
+
+ public ReceivingMailboxStatus offer(TransferableBlock block, long timeoutMs) {
+ long now = System.currentTimeMillis();
+ _stats.merge(StatKey.WAIT_CPU_TIME_MS, now - _lastArriveTime);
+ _lastArriveTime = now;
+ _stats.merge(StatKey.IN_MEMORY_MESSAGES, 1);
+ return offerPrivate(block, timeoutMs);
+ }
+
/**
* Offers a non-error block into the mailbox within the timeout specified, returns whether the block is successfully
* added. If the block is not added, an error block is added to the mailbox.
*/
- public ReceivingMailboxStatus offer(TransferableBlock block, long timeoutMs) {
+ private ReceivingMailboxStatus offerPrivate(TransferableBlock block, long timeoutMs) {
TransferableBlock errorBlock = _errorBlock.get();
if (errorBlock != null) {
LOGGER.debug("Mailbox: {} is already cancelled or errored out, ignoring the late block", _id);
@@ -90,7 +141,10 @@ public class ReceivingMailbox {
return ReceivingMailboxStatus.TIMEOUT;
}
try {
- if (_blocks.offer(block, timeoutMs, TimeUnit.MILLISECONDS)) {
+ long now = System.currentTimeMillis();
+ boolean accepted = _blocks.offer(block, timeoutMs, TimeUnit.MILLISECONDS);
+ _stats.merge(StatKey.OFFER_CPU_TIME_MS, System.currentTimeMillis() - now);
+ if (accepted) {
errorBlock = _errorBlock.get();
if (errorBlock == null) {
if (LOGGER.isDebugEnabled()) {
@@ -170,11 +224,40 @@ public class ReceivingMailbox {
}
}
+ public StatMap<StatKey> getStatMap() {
+ return _stats;
+ }
+
public interface Reader {
void blockReadyToRead();
}
public enum ReceivingMailboxStatus {
- SUCCESS, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ SUCCESS, FIRST_ERROR, ERROR, TIMEOUT, CANCELLED, EARLY_TERMINATED
+ }
+
+ public enum StatKey implements StatMap.Key {
+ DESERIALIZED_MESSAGES(StatMap.Type.INT),
+ DESERIALIZED_BYTES(StatMap.Type.LONG),
+ DESERIALIZATION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ IN_MEMORY_MESSAGES(StatMap.Type.INT),
+ OFFER_CPU_TIME_MS(StatMap.Type.LONG),
+ WAIT_CPU_TIME_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
index ce91701c3a..1d817f91aa 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/channel/MailboxContentObserver.java
@@ -20,16 +20,12 @@ package org.apache.pinot.query.mailbox.channel;
import io.grpc.Context;
import io.grpc.stub.StreamObserver;
-import java.util.Map;
+import java.nio.ByteBuffer;
import java.util.concurrent.TimeUnit;
-import org.apache.pinot.common.datablock.DataBlock;
-import org.apache.pinot.common.datablock.DataBlockUtils;
-import org.apache.pinot.common.datablock.MetadataBlock;
import org.apache.pinot.common.proto.Mailbox.MailboxContent;
import org.apache.pinot.common.proto.Mailbox.MailboxStatus;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
-import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -62,22 +58,9 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
_mailbox = _mailboxService.getReceivingMailbox(mailboxId);
}
try {
- TransferableBlock block;
- DataBlock dataBlock = DataBlockUtils.getDataBlock(mailboxContent.getPayload().asReadOnlyByteBuffer());
- if (dataBlock instanceof MetadataBlock) {
- Map<Integer, String> exceptions = dataBlock.getExceptions();
- if (exceptions.isEmpty()) {
- block = TransferableBlockUtils.getEndOfStreamTransferableBlock(((MetadataBlock) dataBlock).getStats());
- } else {
- _mailbox.setErrorBlock(TransferableBlockUtils.getErrorTransferableBlock(exceptions));
- return;
- }
- } else {
- block = new TransferableBlock(dataBlock);
- }
-
long timeoutMs = Context.current().getDeadline().timeRemaining(TimeUnit.MILLISECONDS);
- ReceivingMailbox.ReceivingMailboxStatus status = _mailbox.offer(block, timeoutMs);
+ ByteBuffer buffer = mailboxContent.getPayload().asReadOnlyByteBuffer();
+ ReceivingMailbox.ReceivingMailboxStatus status = _mailbox.offerRaw(buffer, timeoutMs);
switch (status) {
case SUCCESS:
_responseObserver.onNext(MailboxStatus.newBuilder().setMailboxId(mailboxId)
@@ -88,6 +71,8 @@ public class MailboxContentObserver implements StreamObserver<MailboxContent> {
LOGGER.warn("Mailbox: {} already cancelled from upstream", mailboxId);
cancelStream();
break;
+ case FIRST_ERROR:
+ return;
case ERROR:
LOGGER.warn("Mailbox: {} already errored out (received error block before)", mailboxId);
cancelStream();
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
new file mode 100644
index 0000000000..16cdda026e
--- /dev/null
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
@@ -0,0 +1,207 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.function.IntFunction;
+import org.apache.pinot.query.planner.plannode.AbstractPlanNode;
+import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.ExchangeNode;
+import org.apache.pinot.query.planner.plannode.FilterNode;
+import org.apache.pinot.query.planner.plannode.JoinNode;
+import org.apache.pinot.query.planner.plannode.MailboxReceiveNode;
+import org.apache.pinot.query.planner.plannode.MailboxSendNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.plannode.PlanNodeVisitor;
+import org.apache.pinot.query.planner.plannode.ProjectNode;
+import org.apache.pinot.query.planner.plannode.SetOpNode;
+import org.apache.pinot.query.planner.plannode.SortNode;
+import org.apache.pinot.query.planner.plannode.TableScanNode;
+import org.apache.pinot.query.planner.plannode.ValueNode;
+import org.apache.pinot.query.planner.plannode.WindowNode;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class InStageStatsTreeBuilder implements PlanNodeVisitor<ObjectNode, Void> {
+ private static final Logger LOGGER = LoggerFactory.getLogger(InStageStatsTreeBuilder.class);
+
+ private final MultiStageQueryStats.StageStats _stageStats;
+ private int _index;
+ private static final String CHILDREN_KEY = "children";
+ private final IntFunction<ObjectNode> _jsonStatsByStage;
+
+ public InStageStatsTreeBuilder(MultiStageQueryStats.StageStats stageStats, IntFunction<ObjectNode> jsonStatsByStage) {
+ _stageStats = stageStats;
+ _index = stageStats.getLastOperatorIndex();
+ _jsonStatsByStage = jsonStatsByStage;
+ }
+
+ private ObjectNode selfNode(MultiStageOperator.Type type) {
+ ObjectNode json = JsonUtils.newObjectNode();
+ json.put("type", type.toString());
+ Iterator<Map.Entry<String, JsonNode>> statsIt = _stageStats.getOperatorStats(_index).asJson().fields();
+ while (statsIt.hasNext()) {
+ Map.Entry<String, JsonNode> entry = statsIt.next();
+ json.set(entry.getKey(), entry.getValue());
+ }
+ return json;
+ }
+
+ private ObjectNode recursiveCase(AbstractPlanNode node, MultiStageOperator.Type expectedType) {
+ MultiStageOperator.Type type = _stageStats.getOperatorType(_index);
+ /*
+ Sometimes the operator type is not what we expect, but we can still build the tree
+ This always happen in stage 0, in which case we have two operators but we only have stats for the receive
+ operator.
+ This may also happen leaf stages, in which case the all the stage but the send operator will be compiled into
+ a single leaf node.
+ */
+ if (type != expectedType) {
+ if (type == MultiStageOperator.Type.LEAF) {
+ return selfNode(MultiStageOperator.Type.LEAF);
+ }
+ List<PlanNode> inputs = node.getInputs();
+ int childrenSize = inputs.size();
+ switch (childrenSize) {
+ case 0:
+ return JsonUtils.newObjectNode();
+ case 1:
+ return inputs.get(0).visit(this, null);
+ default:
+ ObjectNode json = JsonUtils.newObjectNode();
+ ArrayNode children = JsonUtils.newArrayNode();
+ for (int i = 0; i < childrenSize; i++) {
+ _index--;
+ if (inputs.size() > i) {
+ children.add(inputs.get(i).visit(this, null));
+ }
+ }
+ json.set(CHILDREN_KEY, children);
+ return json;
+ }
+ }
+ ObjectNode json = selfNode(type);
+ List<PlanNode> inputs = node.getInputs();
+ int size = inputs.size();
+ JsonNode[] childrenArr = new JsonNode[size];
+ if (size > _index) {
+ LOGGER.warn("Operator {} has {} inputs but only {} stats are left", type, size,
+ _index);
+ return json;
+ }
+ for (int i = size - 1; i >= 0; i--) {
+ PlanNode planNode = inputs.get(i);
+ _index--;
+ JsonNode child = planNode.visit(this, null);
+
+ childrenArr[i] = child;
+ }
+ json.set(CHILDREN_KEY, JsonUtils.objectToJsonNode(childrenArr));
+ return json;
+ }
+
+ @Override
+ public ObjectNode visitAggregate(AggregateNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.AGGREGATE);
+ }
+
+ @Override
+ public ObjectNode visitFilter(FilterNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.FILTER);
+ }
+
+ @Override
+ public ObjectNode visitJoin(JoinNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.HASH_JOIN);
+ }
+
+ @Override
+ public ObjectNode visitMailboxReceive(MailboxReceiveNode node, Void context) {
+ ObjectNode json = selfNode(MultiStageOperator.Type.MAILBOX_RECEIVE);
+
+ ArrayNode children = JsonUtils.newArrayNode();
+ int senderStageId = node.getSenderStageId();
+ children.add(_jsonStatsByStage.apply(senderStageId));
+ json.set(CHILDREN_KEY, children);
+ return json;
+ }
+
+ @Override
+ public ObjectNode visitMailboxSend(MailboxSendNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.MAILBOX_SEND);
+ }
+
+ @Override
+ public ObjectNode visitProject(ProjectNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.TRANSFORM);
+ }
+
+ @Override
+ public ObjectNode visitSort(SortNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.SORT_OR_LIMIT);
+ }
+
+ @Override
+ public ObjectNode visitTableScan(TableScanNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.LEAF);
+ }
+
+ @Override
+ public ObjectNode visitValue(ValueNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.LITERAL);
+ }
+
+ @Override
+ public ObjectNode visitWindow(WindowNode node, Void context) {
+ return recursiveCase(node, MultiStageOperator.Type.WINDOW);
+ }
+
+ @Override
+ public ObjectNode visitSetOp(SetOpNode node, Void context) {
+ MultiStageOperator.Type type;
+ switch (node.getSetOpType()) {
+ case UNION:
+ type = MultiStageOperator.Type.UNION;
+ break;
+ case INTERSECT:
+ type = MultiStageOperator.Type.INTERSECT;
+ break;
+ case MINUS:
+ type = MultiStageOperator.Type.MINUS;
+ break;
+ default:
+ throw new IllegalStateException("Unexpected set op type: " + node.getSetOpType());
+ }
+ return recursiveCase(node, type);
+ }
+
+ @Override
+ public ObjectNode visitExchange(ExchangeNode node, Void context) {
+ throw new UnsupportedOperationException("ExchangeNode should not be visited");
+ }
+}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/MultiStageStatsTreeBuilder.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/MultiStageStatsTreeBuilder.java
new file mode 100644
index 0000000000..74c81ed820
--- /dev/null
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/MultiStageStatsTreeBuilder.java
@@ -0,0 +1,43 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime;
+
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.util.List;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+
+
+public class MultiStageStatsTreeBuilder {
+ private final List<PlanNode> _planNodes;
+ private final List<? extends MultiStageQueryStats.StageStats> _queryStats;
+
+ public MultiStageStatsTreeBuilder(List<PlanNode> planNodes,
+ List<? extends MultiStageQueryStats.StageStats> queryStats) {
+ _planNodes = planNodes;
+ _queryStats = queryStats;
+ }
+
+ public ObjectNode jsonStatsByStage(int stage) {
+ MultiStageQueryStats.StageStats stageStats = _queryStats.get(stage);
+ PlanNode planNode = _planNodes.get(stage);
+ InStageStatsTreeBuilder treeBuilder = new InStageStatsTreeBuilder(stageStats, this::jsonStatsByStage);
+ return planNode.visit(treeBuilder, null);
+ }
+}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
index fe078e87c6..0032545c0f 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
@@ -26,6 +26,7 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeoutException;
import javax.annotation.Nullable;
import org.apache.helix.HelixManager;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
import org.apache.pinot.core.data.manager.InstanceDataManager;
@@ -42,6 +43,7 @@ import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.executor.ExecutorServiceUtils;
import org.apache.pinot.query.runtime.executor.OpChainSchedulerService;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.query.runtime.operator.OpChain;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.query.runtime.plan.PhysicalPlanVisitor;
@@ -177,8 +179,9 @@ public class QueryRunner {
receiverMailboxInfos);
for (RoutingInfo routingInfo : routingInfos) {
try {
+ StatMap<MailboxSendOperator.StatKey> statMap = new StatMap<>(MailboxSendOperator.StatKey.class);
_mailboxService.getSendingMailbox(routingInfo.getHostname(), routingInfo.getPort(),
- routingInfo.getMailboxId(), deadlineMs).send(errorBlock);
+ routingInfo.getMailboxId(), deadlineMs, statMap).send(errorBlock);
} catch (TimeoutException e) {
LOGGER.warn("Timed out sending error block to mailbox: {} for request: {}, stage: {}",
routingInfo.getMailboxId(), requestId, stageId, e);
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlock.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlock.java
index e38b1e2723..f42fed4ecc 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlock.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlock.java
@@ -19,22 +19,24 @@
package org.apache.pinot.query.runtime.blocks;
import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.ColumnarDataBlock;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.datablock.DataBlockUtils;
import org.apache.pinot.common.datablock.MetadataBlock;
import org.apache.pinot.common.datablock.RowDataBlock;
-import org.apache.pinot.common.response.ProcessingException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Block;
import org.apache.pinot.core.common.datablock.DataBlockBuilder;
import org.apache.pinot.core.util.DataBlockExtractUtils;
-import org.apache.pinot.query.runtime.operator.OperatorStats;
-import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
-
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
/**
* A {@code TransferableBlock} is a wrapper around {@link DataBlock} for transferring data using
@@ -42,12 +44,15 @@ import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
*/
public class TransferableBlock implements Block {
private final DataBlock.Type _type;
+ @Nullable
private final DataSchema _dataSchema;
private final int _numRows;
private List<Object[]> _container;
private DataBlock _dataBlock;
private Map<Integer, String> _errCodeToExceptionMap;
+ @Nullable
+ private final MultiStageQueryStats _queryStats;
public TransferableBlock(List<Object[]> container, DataSchema dataSchema, DataBlock.Type type) {
_container = container;
@@ -57,6 +62,7 @@ public class TransferableBlock implements Block {
_type = type;
_numRows = _container.size();
_errCodeToExceptionMap = new HashMap<>();
+ _queryStats = null;
}
public TransferableBlock(DataBlock dataBlock) {
@@ -66,19 +72,49 @@ public class TransferableBlock implements Block {
: dataBlock instanceof RowDataBlock ? DataBlock.Type.ROW : DataBlock.Type.METADATA;
_numRows = _dataBlock.getNumberOfRows();
_errCodeToExceptionMap = null;
+ _queryStats = null;
+ }
+
+ public TransferableBlock(MultiStageQueryStats stats) {
+ _queryStats = stats;
+ _type = DataBlock.Type.METADATA;
+ _numRows = 0;
+ _dataSchema = null;
+ _errCodeToExceptionMap = null;
}
- public Map<String, OperatorStats> getResultMetadata() {
+ public List<ByteBuffer> getSerializedStatsByStage() {
if (isSuccessfulEndOfStreamBlock()) {
- return OperatorUtils.getOperatorStatsFromMetadata((MetadataBlock) _dataBlock);
+ List<ByteBuffer> statsByStage;
+ if (_dataBlock instanceof MetadataBlock) {
+ statsByStage = ((MetadataBlock) _dataBlock).getStatsByStage();
+ if (statsByStage == null) {
+ return new ArrayList<>();
+ }
+ } else {
+ Preconditions.checkArgument(_queryStats != null, "QueryStats is null for a successful EOS block");
+ try {
+ statsByStage = _queryStats.serialize();
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ return statsByStage;
}
- return new HashMap<>();
+ return new ArrayList<>();
+ }
+
+ @Nullable
+ public MultiStageQueryStats getQueryStats() {
+ return _queryStats;
}
public int getNumRows() {
return _numRows;
}
+ @Nullable
public DataSchema getDataSchema() {
return _dataSchema;
}
@@ -112,13 +148,6 @@ public class TransferableBlock implements Block {
return _container;
}
- /**
- * Returns whether the data block is already constructed.
- */
- public boolean isDataBlockConstructed() {
- return _dataBlock != null;
- }
-
/**
* Retrieve the binary-packed version of the data block.
* If not already constructed. It will use {@link DataBlockBuilder} to construct the binary-packed format from
@@ -129,13 +158,23 @@ public class TransferableBlock implements Block {
public DataBlock getDataBlock() {
if (_dataBlock == null) {
try {
- if (_type == DataBlock.Type.ROW) {
- _dataBlock = DataBlockBuilder.buildFromRows(_container, _dataSchema);
- } else {
- _dataBlock = DataBlockBuilder.buildFromColumns(_container, _dataSchema);
+ switch (_type) {
+ case ROW:
+ _dataBlock = DataBlockBuilder.buildFromRows(_container, _dataSchema);
+ break;
+ case COLUMNAR:
+ _dataBlock = DataBlockBuilder.buildFromColumns(_container, _dataSchema);
+ break;
+ case METADATA:
+ _dataBlock = new MetadataBlock(getSerializedStatsByStage());
+ break;
+ default:
+ throw new UnsupportedOperationException("Unable to construct block with type: " + _type);
+ }
+ if (_errCodeToExceptionMap != null) {
+ _dataBlock.getExceptions().putAll(_errCodeToExceptionMap);
+ _errCodeToExceptionMap = null;
}
- _dataBlock.getExceptions().putAll(_errCodeToExceptionMap);
- _errCodeToExceptionMap = null;
} catch (Exception e) {
throw new RuntimeException("Unable to create DataBlock", e);
}
@@ -147,18 +186,6 @@ public class TransferableBlock implements Block {
return _dataBlock != null ? _dataBlock.getExceptions() : _errCodeToExceptionMap;
}
- public void addException(ProcessingException processingException) {
- addException(processingException.getErrorCode(), processingException.getMessage());
- }
-
- public void addException(int errCode, String errMsg) {
- if (_dataBlock != null) {
- _dataBlock.addException(errCode, errMsg);
- } else {
- _errCodeToExceptionMap.put(errCode, errMsg);
- }
- }
-
/**
* Return the type of block (one of ROW, COLUMNAR, or METADATA).
*
@@ -208,6 +235,9 @@ public class TransferableBlock implements Block {
return false;
}
+ if (_queryStats != null) {
+ return MetadataBlock.MetadataBlockType.EOS == type;
+ }
MetadataBlock metadata = (MetadataBlock) _dataBlock;
return metadata.getType() == type;
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtils.java
index 355b6fe294..2f58711033 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtils.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtils.java
@@ -26,11 +26,13 @@ import java.util.List;
import java.util.Map;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.datablock.DataBlockUtils;
+import org.apache.pinot.common.datablock.MetadataBlock;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
public final class TransferableBlockUtils {
private static final int MEDIAN_COLUMN_SIZE_BYTES = 8;
- private static final TransferableBlock EMPTY_EOS = new TransferableBlock(DataBlockUtils.getEndOfStreamDataBlock());
+ private static final TransferableBlock EMPTY_EOS = new TransferableBlock(MetadataBlock.newEos());
private TransferableBlockUtils() {
// do not instantiate.
@@ -40,8 +42,12 @@ public final class TransferableBlockUtils {
return EMPTY_EOS;
}
- public static TransferableBlock getEndOfStreamTransferableBlock(Map<String, String> statsMap) {
- return new TransferableBlock(DataBlockUtils.getEndOfStreamDataBlock(statsMap));
+ public static TransferableBlock getEndOfStreamTransferableBlock(MultiStageQueryStats stats) {
+ return new TransferableBlock(stats);
+ }
+
+ public static TransferableBlock wrap(DataBlock dataBlock) {
+ return new TransferableBlock(dataBlock);
}
public static TransferableBlock getErrorTransferableBlock(Exception e) {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
index 440c79200b..ef0f9a7d12 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
@@ -51,7 +51,6 @@ public class OpChainSchedulerService {
Throwable thrown = null;
try {
LOGGER.trace("({}): Executing", operatorChain);
- operatorChain.getStats().executing();
TransferableBlock result = operatorChain.getRoot().nextBlock();
while (!result.isEndOfStreamBlock()) {
result = operatorChain.getRoot().nextBlock();
@@ -59,13 +58,13 @@ public class OpChainSchedulerService {
isFinished = true;
if (result.isErrorBlock()) {
returnedErrorBlock = result;
- LOGGER.error("({}): Completed erroneously {} {}", operatorChain, operatorChain.getStats(),
+ LOGGER.error("({}): Completed erroneously {} {}", operatorChain, result.getQueryStats(),
result.getExceptions());
} else {
- LOGGER.debug("({}): Completed {}", operatorChain, operatorChain.getStats());
+ LOGGER.debug("({}): Completed {}", operatorChain, result.getQueryStats());
}
} catch (Exception e) {
- LOGGER.error("({}): Failed to execute operator chain! {}", operatorChain, operatorChain.getStats(), e);
+ LOGGER.error("({}): Failed to execute operator chain!", operatorChain, e);
thrown = e;
} finally {
_submittedOpChainMap.remove(operatorChain.getId());
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
index 6b8bc3bcc4..239a2e5eb9 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
@@ -29,7 +29,7 @@ import javax.annotation.Nullable;
import org.apache.calcite.sql.SqlKind;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.common.datablock.DataBlock;
-import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.request.Literal;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FunctionContext;
@@ -48,9 +48,10 @@ import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.plannode.AbstractPlanNode;
import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.roaringbitmap.RoaringBitmap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -59,6 +60,7 @@ import org.roaringbitmap.RoaringBitmap;
* When the list of aggregation calls is empty, this class is used to calculate distinct result based on group by keys.
*/
public class AggregateOperator extends MultiStageOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(AggregateOperator.class);
private static final String EXPLAIN_NAME = "AGGREGATE_OPERATOR";
private static final CountAggregationFunction COUNT_STAR_AGG_FUNCTION =
new CountAggregationFunction(Collections.singletonList(ExpressionContext.forIdentifier("*")), false);
@@ -69,12 +71,15 @@ public class AggregateOperator extends MultiStageOperator {
private final AggType _aggType;
private final MultistageAggregationExecutor _aggregationExecutor;
private final MultistageGroupByExecutor _groupByExecutor;
+ @Nullable
+ private TransferableBlock _eosBlock;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
private boolean _hasConstructedAggregateBlock;
- public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, DataSchema resultSchema,
- List<RexExpression> aggCalls, List<RexExpression> groupSet, AggType aggType, List<Integer> filterArgIndices,
- @Nullable AbstractPlanNode.NodeHint nodeHint) {
+ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator,
+ DataSchema resultSchema, List<RexExpression> aggCalls, List<RexExpression> groupSet, AggType aggType,
+ List<Integer> filterArgIndices, @Nullable AbstractPlanNode.NodeHint nodeHint) {
super(context);
_inputOperator = inputOperator;
_resultSchema = resultSchema;
@@ -118,6 +123,22 @@ public class AggregateOperator extends MultiStageOperator {
}
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.AGGREGATE;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_inputOperator);
@@ -132,29 +153,31 @@ public class AggregateOperator extends MultiStageOperator {
@Override
protected TransferableBlock getNextBlock() {
if (_hasConstructedAggregateBlock) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _eosBlock != null;
+ return _eosBlock;
}
TransferableBlock finalBlock = _aggregationExecutor != null ? consumeAggregation() : consumeGroupBy();
// returning upstream error block if finalBlock contains error.
if (finalBlock.isErrorBlock()) {
return finalBlock;
}
- return produceAggregatedBlock();
+ assert finalBlock.isSuccessfulEndOfStreamBlock() : "Final block must be EOS block";
+ _eosBlock = updateEosBlock(finalBlock, _statMap);
+ return produceAggregatedBlock(finalBlock);
}
- private TransferableBlock produceAggregatedBlock() {
+ private TransferableBlock produceAggregatedBlock(TransferableBlock finalBlock) {
_hasConstructedAggregateBlock = true;
if (_aggregationExecutor != null) {
return new TransferableBlock(_aggregationExecutor.getResult(), _resultSchema, DataBlock.Type.ROW);
} else {
List<Object[]> rows = _groupByExecutor.getResult();
if (rows.isEmpty()) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return _eosBlock;
} else {
TransferableBlock dataBlock = new TransferableBlock(rows, _resultSchema, DataBlock.Type.ROW);
if (_groupByExecutor.isNumGroupsLimitReached()) {
- OperatorStats operatorStats = _opChainStats.getOperatorStats(_context, _operatorId);
- operatorStats.recordSingleStat(DataTable.MetadataKey.NUM_GROUPS_LIMIT_REACHED.getName(), "true");
+ _statMap.merge(StatKey.NUM_GROUPS_LIMIT_REACHED, true);
_inputOperator.earlyTerminate();
}
return dataBlock;
@@ -407,4 +430,31 @@ public class AggregateOperator extends MultiStageOperator {
return DataBlockExtractUtils.extractColumn(block.getDataBlock(), colId);
}
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
index a88e122739..d60a0e102d 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
@@ -24,6 +24,7 @@ import java.util.List;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.calcite.rel.RelDistribution;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
import org.apache.pinot.query.planner.physical.MailboxIdUtils;
@@ -48,6 +49,8 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
protected final RelDistribution.Type _exchangeType;
protected final List<String> _mailboxIds;
protected final BlockingMultiStreamConsumer.OfTransferableBlock _multiConsumer;
+ protected final List<StatMap<ReceivingMailbox.StatKey>> _receivingStats;
+ protected final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public BaseMailboxReceiveOperator(OpChainExecutionContext context, RelDistribution.Type exchangeType,
int senderStageId) {
@@ -69,8 +72,9 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
List<ReadMailboxAsyncStream> asyncStreams = _mailboxIds.stream()
.map(mailboxId -> new ReadMailboxAsyncStream(_mailboxService.getReceivingMailbox(mailboxId), this))
.collect(Collectors.toList());
- _multiConsumer =
- new BlockingMultiStreamConsumer.OfTransferableBlock(context.getId(), context.getDeadlineMs(), asyncStreams);
+ _receivingStats = asyncStreams.stream().map(stream -> stream._mailbox.getStatMap()).collect(Collectors.toList());
+ _multiConsumer = new BlockingMultiStreamConsumer.OfTransferableBlock(context, asyncStreams);
+ _statMap.merge(StatKey.FAN_IN, _mailboxIds.size());
}
@Override
@@ -79,6 +83,11 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
_multiConsumer.earlyTerminate();
}
+ @Override
+ public Type getOperatorType() {
+ return Type.MAILBOX_RECEIVE;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return Collections.emptyList();
@@ -96,6 +105,29 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
_multiConsumer.cancel(t);
}
+ @Override
+ protected TransferableBlock updateEosBlock(TransferableBlock upstreamEos, StatMap<?> statMap) {
+ for (StatMap<ReceivingMailbox.StatKey> receivingStats : _receivingStats) {
+ addReceivingStats(receivingStats);
+ }
+ return super.updateEosBlock(upstreamEos, statMap);
+ }
+
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ private void addReceivingStats(StatMap<ReceivingMailbox.StatKey> from) {
+ _statMap.merge(StatKey.RAW_MESSAGES, from.getInt(ReceivingMailbox.StatKey.DESERIALIZED_MESSAGES));
+ _statMap.merge(StatKey.DESERIALIZED_BYTES, from.getLong(ReceivingMailbox.StatKey.DESERIALIZED_BYTES));
+ _statMap.merge(StatKey.DESERIALIZATION_TIME_MS, from.getLong(ReceivingMailbox.StatKey.DESERIALIZATION_TIME_MS));
+ _statMap.merge(StatKey.IN_MEMORY_MESSAGES, from.getInt(ReceivingMailbox.StatKey.IN_MEMORY_MESSAGES));
+ _statMap.merge(StatKey.DOWNSTREAM_WAIT_MS, from.getLong(ReceivingMailbox.StatKey.OFFER_CPU_TIME_MS));
+ _statMap.merge(StatKey.UPSTREAM_WAIT_MS, from.getLong(ReceivingMailbox.StatKey.WAIT_CPU_TIME_MS));
+ }
+
private static class ReadMailboxAsyncStream implements AsyncStream<TransferableBlock> {
private final ReceivingMailbox _mailbox;
private final BaseMailboxReceiveOperator _operator;
@@ -116,7 +148,6 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
TransferableBlock block = _mailbox.poll();
if (block != null && block.isSuccessfulEndOfStreamBlock()) {
_operator._mailboxService.releaseReceivingMailbox(_mailbox);
- _operator._opChainStats.getOperatorStatsMap().putAll(block.getResultMetadata());
}
return block;
}
@@ -136,4 +167,71 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
_mailbox.cancel();
}
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ /**
+ * How many send mailboxes are being read by this receive operator.
+ * <p>
+ * Clock time will be proportional to this number and the parallelism of the stage.
+ */
+ FAN_IN(StatMap.Type.INT) {
+ @Override
+ public int merge(int value1, int value2) {
+ return Math.max(value1, value2);
+ }
+ },
+ /**
+ * How many messages have been received in heap format by this mailbox.
+ * <p>
+ * The lower the relation between RAW_MESSAGES and IN_MEMORY_MESSAGES, the more efficient the exchange is.
+ */
+ IN_MEMORY_MESSAGES(StatMap.Type.INT),
+ /**
+ * How many messages have been received in raw format and therefore deserialized by this mailbox.
+ * <p>
+ * The higher the relation between RAW_MESSAGES and IN_MEMORY_MESSAGES, the less efficient the exchange is.
+ */
+ RAW_MESSAGES(StatMap.Type.INT),
+ /**
+ * How many bytes have been deserialized by this mailbox.
+ * <p>
+ * A high number here indicates that the mailbox is receiving a lot of data from other servers.
+ */
+ DESERIALIZED_BYTES(StatMap.Type.LONG),
+ /**
+ * How long (in CPU time) it took to deserialize the raw messages received by this mailbox.
+ */
+ DESERIALIZATION_TIME_MS(StatMap.Type.LONG),
+ /**
+ * How long (in CPU time) it took to offer the messages to downstream operator.
+ */
+ DOWNSTREAM_WAIT_MS(StatMap.Type.LONG),
+ /**
+ * How long (in CPU time) it took to wait for the messages to be offered to downstream operator.
+ */
+ UPSTREAM_WAIT_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/FilterOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/FilterOperator.java
index 5747d6191d..8cab4899bc 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/FilterOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/FilterOperator.java
@@ -24,6 +24,7 @@ import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.logical.RexExpression;
@@ -32,6 +33,8 @@ import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.utils.BooleanUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/*
@@ -48,11 +51,14 @@ import org.apache.pinot.spi.utils.BooleanUtils;
Note: Scalar functions are the ones we have in v1 engine and only do function name and arg # matching.
*/
public class FilterOperator extends MultiStageOperator {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(FilterOperator.class);
private static final String EXPLAIN_NAME = "FILTER";
private final MultiStageOperator _upstreamOperator;
private final TransformOperand _filterOperand;
private final DataSchema _dataSchema;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public FilterOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, DataSchema dataSchema,
RexExpression filter) {
@@ -64,6 +70,22 @@ public class FilterOperator extends MultiStageOperator {
"Filter operand must return BOOLEAN, got: %s", _filterOperand.getResultType());
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.FILTER;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_upstreamOperator);
@@ -79,7 +101,10 @@ public class FilterOperator extends MultiStageOperator {
protected TransferableBlock getNextBlock() {
TransferableBlock block = _upstreamOperator.nextBlock();
if (block.isEndOfStreamBlock()) {
- return block;
+ if (block.isErrorBlock()) {
+ return block;
+ }
+ return updateEosBlock(block, _statMap);
}
List<Object[]> resultRows = new ArrayList<>();
for (Object[] row : block.getContainer()) {
@@ -90,4 +115,29 @@ public class FilterOperator extends MultiStageOperator {
}
return new TransferableBlock(resultRows, _dataSchema, DataBlock.Type.ROW);
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ };
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
index 92d50dc54e..aa2be1850c 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
@@ -31,7 +31,7 @@ import javax.annotation.Nullable;
import org.apache.calcite.rel.core.JoinRelType;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.common.datablock.DataBlock;
-import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.response.ProcessingException;
import org.apache.pinot.common.utils.DataSchema;
@@ -45,9 +45,12 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.utils.BooleanUtils;
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -65,6 +68,8 @@ import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOver
// TODO: Move inequi out of hashjoin. (https://github.com/apache/pinot/issues/9728)
// TODO: Support memory size based resource limit.
public class HashJoinOperator extends MultiStageOperator {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(HashJoinOperator.class);
private static final String EXPLAIN_NAME = "HASH_JOIN";
private static final int INITIAL_HEURISTIC_SIZE = 16;
private static final int DEFAULT_MAX_ROWS_IN_JOIN = 1024 * 1024; // 2^20, around 1MM rows
@@ -88,6 +93,7 @@ public class HashJoinOperator extends MultiStageOperator {
private final int _resultColumnSize;
private final List<TransformOperand> _joinClauseEvaluators;
private boolean _isHashTableBuilt;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
// Used by non-inner join.
// Needed to indicate we have finished processing all results after returning last block.
@@ -112,6 +118,11 @@ public class HashJoinOperator extends MultiStageOperator {
private int _currentRowsInHashTable = 0;
+ @Nullable
+ private MultiStageQueryStats _rightSideStats = null;
+ @Nullable
+ private MultiStageQueryStats _leftSideStats = null;
+
public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftTableOperator,
MultiStageOperator rightTableOperator, DataSchema leftSchema, JoinNode node) {
super(context);
@@ -146,6 +157,22 @@ public class HashJoinOperator extends MultiStageOperator {
_joinOverflowMode = getJoinOverflowMode(metadata, node.getJoinHints());
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.HASH_JOIN;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
private int getMaxRowInJoin(Map<String, String> opChainMetadata, @Nullable AbstractPlanNode.NodeHint nodeHint) {
if (nodeHint != null) {
Map<String, String> joinOptions = nodeHint._hintOptions.get(PinotHintOptions.JOIN_HINT_OPTIONS);
@@ -191,7 +218,8 @@ public class HashJoinOperator extends MultiStageOperator {
protected TransferableBlock getNextBlock()
throws ProcessingException {
if (_isTerminated) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _leftSideStats != null;
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_leftSideStats);
}
if (!_isHashTableBuilt) {
// Build JOIN hash table
@@ -207,6 +235,7 @@ public class HashJoinOperator extends MultiStageOperator {
private void buildBroadcastHashTable()
throws ProcessingException {
+ long startTime = System.currentTimeMillis();
TransferableBlock rightBlock = _rightTableOperator.nextBlock();
while (!TransferableBlockUtils.isEndOfStream(rightBlock)) {
List<Object[]> container = rightBlock.getContainer();
@@ -223,8 +252,7 @@ public class HashJoinOperator extends MultiStageOperator {
// Just fill up the buffer.
int remainingRows = _maxRowsInHashTable - _currentRowsInHashTable;
container = container.subList(0, remainingRows);
- OperatorStats operatorStats = _opChainStats.getOperatorStats(_context, _operatorId);
- operatorStats.recordSingleStat(DataTable.MetadataKey.MAX_ROWS_IN_JOIN_REACHED.getName(), "true");
+ _statMap.merge(StatKey.MAX_ROWS_IN_JOIN_REACHED, true);
// setting only the rightTableOperator to be early terminated and awaits EOS block next.
_rightTableOperator.earlyTerminate();
}
@@ -246,7 +274,10 @@ public class HashJoinOperator extends MultiStageOperator {
_upstreamErrorBlock = rightBlock;
} else {
_isHashTableBuilt = true;
+ _rightSideStats = rightBlock.getQueryStats();
+ assert _rightSideStats != null;
}
+ _statMap.merge(StatKey.TIME_BUILDING_HASH_TABLE_MS, System.currentTimeMillis() - startTime);
}
private TransferableBlock buildJoinedDataBlock(TransferableBlock leftBlock) {
@@ -255,8 +286,13 @@ public class HashJoinOperator extends MultiStageOperator {
return _upstreamErrorBlock;
}
if (leftBlock.isSuccessfulEndOfStreamBlock()) {
+ assert _rightSideStats != null;
+ _leftSideStats = leftBlock.getQueryStats();
+ assert _leftSideStats != null;
+ _leftSideStats.mergeInOrder(_rightSideStats, getOperatorType(), _statMap);
+
if (!needUnmatchedRightRows()) {
- return leftBlock;
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_leftSideStats);
}
// TODO: Moved to a different function.
// Return remaining non-matched rows for non-inner join.
@@ -390,4 +426,35 @@ public class HashJoinOperator extends MultiStageOperator {
private boolean needUnmatchedLeftRows() {
return _joinType == JoinRelType.LEFT || _joinType == JoinRelType.FULL;
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ MAX_ROWS_IN_JOIN_REACHED(StatMap.Type.BOOLEAN),
+ /**
+ * How long (CPU time) has been spent on building the hash table.
+ */
+ TIME_BUILDING_HASH_TABLE_MS(StatMap.Type.LONG);
+
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/IntersectOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/IntersectOperator.java
index b0ba3c6fa3..2ef003801d 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/IntersectOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/IntersectOperator.java
@@ -23,19 +23,33 @@ import javax.annotation.Nullable;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.data.table.Record;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Intersect operator.
*/
public class IntersectOperator extends SetOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(IntersectOperator.class);
private static final String EXPLAIN_NAME = "INTERSECT";
- public IntersectOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
+ public IntersectOperator(OpChainExecutionContext opChainExecutionContext,
+ List<MultiStageOperator> upstreamOperators,
DataSchema dataSchema) {
super(opChainExecutionContext, upstreamOperators, dataSchema);
}
+ @Override
+ public Type getOperatorType() {
+ return Type.INTERSECT;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Nullable
@Override
public String toExplainString() {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
index ea5a47df98..75dc55a150 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperator.java
@@ -33,10 +33,13 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
+import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.datablock.MetadataBlock;
import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
@@ -52,8 +55,11 @@ import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.utils.TypeUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionValue;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -70,6 +76,8 @@ import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionValu
* </ul>
*/
public class LeafStageTransferableBlockOperator extends MultiStageOperator {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(LeafStageTransferableBlockOperator.class);
private static final String EXPLAIN_NAME = "LEAF_STAGE_TRANSFER_OPERATOR";
// Use a special results block to indicate that this is the last results block
@@ -85,7 +93,7 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
private Future<Void> _executionFuture;
private volatile Map<Integer, String> _exceptions;
- private volatile Map<String, String> _executionStats;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public LeafStageTransferableBlockOperator(OpChainExecutionContext context, List<ServerQueryRequest> requests,
DataSchema dataSchema, QueryExecutor queryExecutor, ExecutorService executorService) {
@@ -99,6 +107,24 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
Integer maxStreamingPendingBlocks = QueryOptionsUtils.getMaxStreamingPendingBlocks(context.getOpChainMetadata());
_blockingQueue = new ArrayBlockingQueue<>(maxStreamingPendingBlocks != null ? maxStreamingPendingBlocks
: QueryOptionValue.DEFAULT_MAX_STREAMING_PENDING_BLOCKS);
+ String tableName = context.getLeafStageContext().getStagePlan().getStageMetadata().getTableName();
+ _statMap.merge(StatKey.TABLE, tableName);
+ }
+
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.LEAF;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
}
@Override
@@ -135,14 +161,125 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
}
}
- private TransferableBlock constructMetadataBlock() {
- // All data blocks have been returned. Record the stats and return EOS.
- Map<String, String> executionStats = _executionStats;
+ private void mergeExecutionStats(@Nullable Map<String, String> executionStats) {
if (executionStats != null) {
- OperatorStats operatorStats = _opChainStats.getOperatorStats(_context, getOperatorId());
- operatorStats.recordExecutionStats(executionStats);
+ for (Map.Entry<String, String> entry : executionStats.entrySet()) {
+ DataTable.MetadataKey key = DataTable.MetadataKey.getByName(entry.getKey());
+ if (key == null) {
+ LOGGER.debug("Skipping unknown execution stat: {}", entry.getKey());
+ continue;
+ }
+ switch (key) {
+ case UNKNOWN:
+ LOGGER.debug("Skipping unknown execution stat: {}", entry.getKey());
+ break;
+ case TABLE:
+ _statMap.merge(StatKey.TABLE, entry.getValue());
+ break;
+ case NUM_DOCS_SCANNED:
+ _statMap.merge(StatKey.NUM_DOCS_SCANNED, Long.parseLong(entry.getValue()));
+ break;
+ case NUM_ENTRIES_SCANNED_IN_FILTER:
+ _statMap.merge(StatKey.NUM_ENTRIES_SCANNED_IN_FILTER, Long.parseLong(entry.getValue()));
+ break;
+ case NUM_ENTRIES_SCANNED_POST_FILTER:
+ _statMap.merge(StatKey.NUM_ENTRIES_SCANNED_POST_FILTER, Long.parseLong(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_QUERIED:
+ _statMap.merge(StatKey.NUM_SEGMENTS_QUERIED, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_PROCESSED:
+ _statMap.merge(StatKey.NUM_SEGMENTS_PROCESSED, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_MATCHED:
+ _statMap.merge(StatKey.NUM_SEGMENTS_MATCHED, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_CONSUMING_SEGMENTS_QUERIED:
+ _statMap.merge(StatKey.NUM_CONSUMING_SEGMENTS_QUERIED, Integer.parseInt(entry.getValue()));
+ break;
+ case MIN_CONSUMING_FRESHNESS_TIME_MS:
+ _statMap.merge(StatKey.MIN_CONSUMING_FRESHNESS_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ case TOTAL_DOCS:
+ _statMap.merge(StatKey.TOTAL_DOCS, Long.parseLong(entry.getValue()));
+ break;
+ case NUM_GROUPS_LIMIT_REACHED:
+ _statMap.merge(StatKey.NUM_GROUPS_LIMIT_REACHED, Boolean.parseBoolean(entry.getValue()));
+ break;
+ case TIME_USED_MS:
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ case TRACE_INFO:
+ LOGGER.debug("Skipping trace info: {}", entry.getValue());
+ break;
+ case REQUEST_ID:
+ LOGGER.debug("Skipping request ID: {}", entry.getValue());
+ break;
+ case NUM_RESIZES:
+ _statMap.merge(StatKey.NUM_RESIZES, Integer.parseInt(entry.getValue()));
+ break;
+ case RESIZE_TIME_MS:
+ _statMap.merge(StatKey.RESIZE_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ case THREAD_CPU_TIME_NS:
+ _statMap.merge(StatKey.THREAD_CPU_TIME_NS, Long.parseLong(entry.getValue()));
+ break;
+ case SYSTEM_ACTIVITIES_CPU_TIME_NS:
+ _statMap.merge(StatKey.SYSTEM_ACTIVITIES_CPU_TIME_NS, Long.parseLong(entry.getValue()));
+ break;
+ case RESPONSE_SER_CPU_TIME_NS:
+ _statMap.merge(StatKey.RESPONSE_SER_CPU_TIME_NS, Long.parseLong(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_PRUNED_BY_SERVER:
+ _statMap.merge(StatKey.NUM_SEGMENTS_PRUNED_BY_SERVER, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_PRUNED_INVALID:
+ _statMap.merge(StatKey.NUM_SEGMENTS_PRUNED_INVALID, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_PRUNED_BY_LIMIT:
+ _statMap.merge(StatKey.NUM_SEGMENTS_PRUNED_BY_LIMIT, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_SEGMENTS_PRUNED_BY_VALUE:
+ _statMap.merge(StatKey.NUM_SEGMENTS_PRUNED_BY_VALUE, Integer.parseInt(entry.getValue()));
+ break;
+ case EXPLAIN_PLAN_NUM_EMPTY_FILTER_SEGMENTS:
+ LOGGER.debug("Skipping empty filter segments: {}", entry.getValue());
+ break;
+ case EXPLAIN_PLAN_NUM_MATCH_ALL_FILTER_SEGMENTS:
+ LOGGER.debug("Skipping match all filter segments: {}", entry.getValue());
+ break;
+ case NUM_CONSUMING_SEGMENTS_PROCESSED:
+ _statMap.merge(StatKey.NUM_CONSUMING_SEGMENTS_PROCESSED, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_CONSUMING_SEGMENTS_MATCHED:
+ _statMap.merge(StatKey.NUM_CONSUMING_SEGMENTS_MATCHED, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_BLOCKS:
+ _statMap.merge(StatKey.NUM_BLOCKS, Integer.parseInt(entry.getValue()));
+ break;
+ case NUM_ROWS:
+ _statMap.merge(StatKey.EMITTED_ROWS, Integer.parseInt(entry.getValue()));
+ break;
+ case OPERATOR_EXECUTION_TIME_MS:
+ _statMap.merge(StatKey.OPERATOR_EXECUTION_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ case OPERATOR_EXEC_START_TIME_MS:
+ _statMap.merge(StatKey.OPERATOR_EXEC_START_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ case OPERATOR_EXEC_END_TIME_MS:
+ _statMap.merge(StatKey.OPERATOR_EXEC_END_TIME_MS, Long.parseLong(entry.getValue()));
+ break;
+ default: {
+ throw new IllegalArgumentException("Unhandled V1 execution stat: " + entry.getKey());
+ }
+ }
+ }
}
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ }
+
+ private TransferableBlock constructMetadataBlock() {
+ MultiStageQueryStats multiStageQueryStats = MultiStageQueryStats.createLeaf(_context.getStageId(), _statMap);
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(multiStageQueryStats);
}
private Future<Void> startExecution() {
@@ -170,11 +307,14 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
addResultsBlock(resultsBlock);
}
// Collect the execution stats
- _executionStats = instanceResponseBlock.getResponseMetadata();
+ mergeExecutionStats(instanceResponseBlock.getResponseMetadata());
}
} else {
assert _requests.size() == 2;
- Future<?>[] futures = new Future[2];
+ Future<Map<String, String>>[] futures = new Future[2];
+ // TODO: this latch mechanism is not the most elegant. We should change it to use a CompletionService.
+ // In order to interrupt the execution in case of error, we could different mechanisms like throwing in the
+ // future, or using a shared volatile variable.
CountDownLatch latch = new CountDownLatch(2);
for (int i = 0; i < 2; i++) {
ServerQueryRequest request = _requests.get(i);
@@ -190,6 +330,7 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
// Drain the latch when receiving exception block and not wait for the other thread to finish
_exceptions = exceptions;
latch.countDown();
+ return Collections.emptyMap();
} else {
// NOTE: Instance response block might contain data (not metadata only) when all the segments are
// pruned. Add the results block if it contains data.
@@ -198,16 +339,8 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
addResultsBlock(resultsBlock);
}
// Collect the execution stats
- Map<String, String> executionStats = instanceResponseBlock.getResponseMetadata();
- synchronized (LeafStageTransferableBlockOperator.this) {
- if (_executionStats == null) {
- _executionStats = executionStats;
- } else {
- aggregateExecutionStats(_executionStats, executionStats);
- }
- }
+ return instanceResponseBlock.getResponseMetadata();
}
- return null;
} finally {
latch.countDown();
}
@@ -218,9 +351,17 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
throw new TimeoutException("Timed out waiting for leaf stage to finish");
}
// Propagate the exception thrown by the leaf stage
+ for (Future<Map<String, String>> future : futures) {
+ Map<String, String> stats =
+ future.get(_context.getDeadlineMs() - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+ mergeExecutionStats(stats);
+ }
+ } catch (TimeoutException e) {
+ // Cancel all the futures and throw the exception
for (Future<?> future : futures) {
- future.get();
+ future.cancel(true);
}
+ throw new TimeoutException("Timed out waiting for leaf stage to finish");
} finally {
for (Future<?> future : futures) {
future.cancel(true);
@@ -249,28 +390,16 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
for (Map.Entry<String, String> entry : stats2.entrySet()) {
String k2 = entry.getKey();
String v2 = entry.getValue();
- stats1.compute(k2, (k1, v1) -> {
- if (v1 == null) {
- return v2;
- }
+ stats1.merge(k2, v2, (val1, val2) -> {
try {
- return Long.toString(Long.parseLong(v1) + Long.parseLong(v2));
+ return Long.toString(Long.parseLong(val1) + Long.parseLong(val2));
} catch (Exception e) {
- return v1 + "\n" + v2;
+ return val1 + "\n" + val2;
}
});
}
}
- /**
- * Leaf stage operators should always collect stats for the tables used in queries
- * Otherwise the Broker response will just contain zeros for every stat value
- */
- @Override
- protected boolean shouldCollectStats() {
- return true;
- }
-
@Override
public void close() {
if (_executionFuture != null) {
@@ -419,4 +548,110 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator {
addResultsBlock(block);
}
}
+
+ public enum StatKey implements StatMap.Key {
+ TABLE(StatMap.Type.STRING, null),
+ EXECUTION_TIME_MS(StatMap.Type.LONG, BrokerResponseNativeV2.StatKey.TIME_USED_MS) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG, null) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ NUM_DOCS_SCANNED(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_IN_FILTER(StatMap.Type.LONG),
+ NUM_ENTRIES_SCANNED_POST_FILTER(StatMap.Type.LONG),
+ NUM_SEGMENTS_QUERIED(StatMap.Type.INT),
+ NUM_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_SEGMENTS_MATCHED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_QUERIED(StatMap.Type.INT),
+ // the timestamp indicating the freshness of the data queried in consuming segments.
+ // This can be ingestion timestamp if provided by the stream, or the last index time
+ MIN_CONSUMING_FRESHNESS_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public long merge(long value1, long value2) {
+ return StatMap.Key.minPositive(value1, value2);
+ }
+ },
+ TOTAL_DOCS(StatMap.Type.LONG),
+ NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
+ NUM_RESIZES(StatMap.Type.INT, null),
+ RESIZE_TIME_MS(StatMap.Type.LONG, null),
+ THREAD_CPU_TIME_NS(StatMap.Type.LONG, null),
+ SYSTEM_ACTIVITIES_CPU_TIME_NS(StatMap.Type.LONG, null),
+ RESPONSE_SER_CPU_TIME_NS(StatMap.Type.LONG, null) {
+ @Override
+ public String getStatName() {
+ return "responseSerializationCpuTimeNs";
+ }
+ },
+ NUM_SEGMENTS_PRUNED_BY_SERVER(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_INVALID(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_LIMIT(StatMap.Type.INT),
+ NUM_SEGMENTS_PRUNED_BY_VALUE(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_PROCESSED(StatMap.Type.INT),
+ NUM_CONSUMING_SEGMENTS_MATCHED(StatMap.Type.INT),
+ NUM_BLOCKS(StatMap.Type.INT, null),
+ OPERATOR_EXECUTION_TIME_MS(StatMap.Type.LONG, null),
+ OPERATOR_EXEC_START_TIME_MS(StatMap.Type.LONG, null) {
+ @Override
+ public long merge(long value1, long value2) {
+ return StatMap.Key.minPositive(value1, value2);
+ }
+ },
+ OPERATOR_EXEC_END_TIME_MS(StatMap.Type.LONG, null) {
+ @Override
+ public long merge(long value1, long value2) {
+ return Math.max(value1, value2);
+ }
+ },;
+ private final StatMap.Type _type;
+ @Nullable
+ private final BrokerResponseNativeV2.StatKey _brokerKey;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ _brokerKey = BrokerResponseNativeV2.StatKey.valueOf(name());
+ }
+
+ StatKey(StatMap.Type type, @Nullable BrokerResponseNativeV2.StatKey brokerKey) {
+ _type = type;
+ _brokerKey = brokerKey;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+
+ public void updateBrokerMetadata(StatMap<BrokerResponseNativeV2.StatKey> oldMetadata, StatMap<StatKey> stats) {
+ if (_brokerKey != null) {
+ switch (_type) {
+ case LONG:
+ if (_brokerKey.getType() == StatMap.Type.INT) {
+ oldMetadata.merge(_brokerKey, (int) stats.getLong(this));
+ } else {
+ oldMetadata.merge(_brokerKey, stats.getLong(this));
+ }
+ break;
+ case INT:
+ oldMetadata.merge(_brokerKey, stats.getInt(this));
+ break;
+ case BOOLEAN:
+ oldMetadata.merge(_brokerKey, stats.getBoolean(this));
+ break;
+ case STRING:
+ oldMetadata.merge(_brokerKey, stats.getString(this));
+ break;
+ default:
+ throw new IllegalStateException("Unsupported type: " + _type);
+ }
+ }
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LiteralValueOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LiteralValueOperator.java
index ad0cf2a129..3279cb4cd0 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LiteralValueOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LiteralValueOperator.java
@@ -23,10 +23,12 @@ import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -39,6 +41,7 @@ public class LiteralValueOperator extends MultiStageOperator {
private final DataSchema _dataSchema;
private final TransferableBlock _rexLiteralBlock;
private boolean _isLiteralBlockReturned;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public LiteralValueOperator(OpChainExecutionContext context, DataSchema dataSchema,
List<List<RexExpression>> rexLiteralRows) {
@@ -49,6 +52,17 @@ public class LiteralValueOperator extends MultiStageOperator {
_isLiteralBlockReturned = context.getId().getVirtualServerId() != 0;
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of();
@@ -66,10 +80,20 @@ public class LiteralValueOperator extends MultiStageOperator {
_isLiteralBlockReturned = true;
return _rexLiteralBlock;
} else {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return createEosBlock();
}
}
+ protected TransferableBlock createEosBlock() {
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(
+ MultiStageQueryStats.createLiteral(_context.getStageId(), _statMap));
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.LITERAL;
+ }
+
private TransferableBlock constructBlock(List<List<RexExpression>> rexLiteralRows) {
List<Object[]> blockContent = new ArrayList<>();
for (List<RexExpression> rexLiteralRow : rexLiteralRows) {
@@ -81,4 +105,19 @@ public class LiteralValueOperator extends MultiStageOperator {
}
return new TransferableBlock(blockContent, _dataSchema, DataBlock.Type.ROW);
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG),
+ EMITTED_ROWS(StatMap.Type.LONG);
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java
index 584b49640f..6fa7e3119b 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java
@@ -22,6 +22,8 @@ import org.apache.calcite.rel.RelDistribution;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -29,6 +31,7 @@ import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
* {@link MultiStageOperator#getNextBlock()} API.
*/
public class MailboxReceiveOperator extends BaseMailboxReceiveOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(MailboxReceiveOperator.class);
private static final String EXPLAIN_NAME = "MAILBOX_RECEIVE";
public MailboxReceiveOperator(OpChainExecutionContext context, RelDistribution.Type exchangeType, int senderStageId) {
@@ -40,6 +43,11 @@ public class MailboxReceiveOperator extends BaseMailboxReceiveOperator {
return EXPLAIN_NAME;
}
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
protected TransferableBlock getNextBlock() {
TransferableBlock block = _multiConsumer.readBlockBlocking();
@@ -50,6 +58,9 @@ public class MailboxReceiveOperator extends BaseMailboxReceiveOperator {
while (_isEarlyTerminated && !block.isEndOfStreamBlock()) {
block = _multiConsumer.readBlockBlocking();
}
+ if (block.isSuccessfulEndOfStreamBlock()) {
+ updateEosBlock(block, _statMap);
+ }
return block;
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java
index 5d66d6afac..fce214e7aa 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java
@@ -19,6 +19,7 @@
package org.apache.pinot.query.runtime.operator;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import java.util.Collections;
import java.util.EnumSet;
@@ -28,6 +29,7 @@ import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.mailbox.SendingMailbox;
import org.apache.pinot.query.planner.logical.RexExpression;
@@ -37,7 +39,7 @@ import org.apache.pinot.query.routing.RoutingInfo;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.exchange.BlockExchange;
-import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.exception.QueryCancelledException;
import org.slf4j.Logger;
@@ -62,29 +64,34 @@ public class MailboxSendOperator extends MultiStageOperator {
private final List<RexExpression> _collationKeys;
private final List<RelFieldCollation.Direction> _collationDirections;
private final boolean _isSortOnSender;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator sourceOperator,
RelDistribution.Type distributionType, @Nullable List<Integer> distributionKeys,
@Nullable List<RexExpression> collationKeys, @Nullable List<RelFieldCollation.Direction> collationDirections,
boolean isSortOnSender, int receiverStageId) {
- this(context, sourceOperator, getBlockExchange(context, distributionType, distributionKeys, receiverStageId),
+ this(context, sourceOperator,
+ statMap -> getBlockExchange(context, distributionType, distributionKeys, receiverStageId, statMap),
collationKeys, collationDirections, isSortOnSender);
+ _statMap.merge(StatKey.STAGE, context.getStageId());
+ _statMap.merge(StatKey.PARALLELISM, 1);
}
@VisibleForTesting
- MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator sourceOperator, BlockExchange exchange,
+ MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator sourceOperator,
+ Function<StatMap<StatKey>, BlockExchange> exchangeFactory,
@Nullable List<RexExpression> collationKeys, @Nullable List<RelFieldCollation.Direction> collationDirections,
boolean isSortOnSender) {
super(context);
_sourceOperator = sourceOperator;
- _exchange = exchange;
+ _exchange = exchangeFactory.apply(_statMap);
_collationKeys = collationKeys;
_collationDirections = collationDirections;
_isSortOnSender = isSortOnSender;
}
private static BlockExchange getBlockExchange(OpChainExecutionContext context, RelDistribution.Type distributionType,
- @Nullable List<Integer> distributionKeys, int receiverStageId) {
+ @Nullable List<Integer> distributionKeys, int receiverStageId, StatMap<StatKey> statMap) {
Preconditions.checkState(SUPPORTED_EXCHANGE_TYPES.contains(distributionType), "Unsupported distribution type: %s",
distributionType);
MailboxService mailboxService = context.getMailboxService();
@@ -97,12 +104,29 @@ public class MailboxSendOperator extends MultiStageOperator {
MailboxIdUtils.toRoutingInfos(requestId, context.getStageId(), context.getWorkerId(), receiverStageId,
mailboxInfos);
List<SendingMailbox> sendingMailboxes = routingInfos.stream()
- .map(v -> mailboxService.getSendingMailbox(v.getHostname(), v.getPort(), v.getMailboxId(), deadlineMs))
+ .map(v -> mailboxService.getSendingMailbox(v.getHostname(), v.getPort(), v.getMailboxId(), deadlineMs, statMap))
.collect(Collectors.toList());
+ statMap.merge(StatKey.FAN_OUT, sendingMailboxes.size());
return BlockExchange.getExchange(sendingMailboxes, distributionType, distributionKeys,
TransferableBlockUtils::splitBlock);
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.MAILBOX_SEND;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return Collections.singletonList(_sourceOperator);
@@ -119,12 +143,9 @@ public class MailboxSendOperator extends MultiStageOperator {
try {
TransferableBlock block = _sourceOperator.nextBlock();
if (block.isSuccessfulEndOfStreamBlock()) {
- // Stats need to be populated here because the block is being sent to the mailbox
- // and the receiving opChain will not be able to access the stats from the previous opChain
- TransferableBlock eosBlockWithStats = TransferableBlockUtils.getEndOfStreamTransferableBlock(
- OperatorUtils.getMetadataFromOperatorStats(_opChainStats.getOperatorStatsMap()));
+ updateEosBlock(block, _statMap);
// no need to check early terminate signal b/c the current block is already EOS
- sendTransferableBlock(eosBlockWithStats);
+ sendTransferableBlock(block);
} else {
if (sendTransferableBlock(block)) {
earlyTerminate();
@@ -133,7 +154,7 @@ public class MailboxSendOperator extends MultiStageOperator {
return block;
} catch (QueryCancelledException e) {
LOGGER.debug("Query was cancelled! for opChain: {}", _context.getId());
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return createLeafBlock();
} catch (TimeoutException e) {
LOGGER.warn("Timed out transferring data on opChain: {}", _context.getId(), e);
return TransferableBlockUtils.getErrorTransferableBlock(e);
@@ -149,6 +170,11 @@ public class MailboxSendOperator extends MultiStageOperator {
}
}
+ protected TransferableBlock createLeafBlock() {
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(
+ MultiStageQueryStats.createCancelledSend(_context.getStageId(), _statMap));
+ }
+
private boolean sendTransferableBlock(TransferableBlock block)
throws Exception {
boolean isEarlyTerminated = _exchange.send(block);
@@ -158,15 +184,6 @@ public class MailboxSendOperator extends MultiStageOperator {
return isEarlyTerminated;
}
- /**
- * This method is overridden to return true because this operator is last in the chain and needs to collect
- * execution time stats
- */
- @Override
- protected boolean shouldCollectStats() {
- return true;
- }
-
@Override
public void close() {
super.close();
@@ -178,4 +195,87 @@ public class MailboxSendOperator extends MultiStageOperator {
super.cancel(t);
_exchange.cancel(t);
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ STAGE(StatMap.Type.INT) {
+ @Override
+ public int merge(int value1, int value2) {
+ return StatMap.Key.eqNotZero(value1, value2);
+ }
+
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ /**
+ * Number of parallelism of the stage this operator is the root of.
+ * <p>
+ * The CPU times reported by this stage will be proportional to this number.
+ */
+ PARALLELISM(StatMap.Type.INT),
+ /**
+ * How many receive mailboxes are being written by this send operator.
+ */
+ FAN_OUT(StatMap.Type.INT) {
+ @Override
+ public int merge(int value1, int value2) {
+ return Math.max(value1, value2);
+ }
+ },
+ /**
+ * How many messages have been sent in heap format by this mailbox.
+ * <p>
+ * The lower the relation between RAW_MESSAGES and IN_MEMORY_MESSAGES, the more efficient the exchange is.
+ */
+ IN_MEMORY_MESSAGES(StatMap.Type.INT),
+ /**
+ * How many messages have been sent in raw format and therefore serialized by this mailbox.
+ * <p>
+ * The higher the relation between RAW_MESSAGES and IN_MEMORY_MESSAGES, the less efficient the exchange is.
+ */
+ RAW_MESSAGES(StatMap.Type.INT),
+ /**
+ * How many bytes have been serialized by this mailbox.
+ * <p>
+ * A high number here indicates that the mailbox is sending a lot of data to other servers.
+ */
+ SERIALIZED_BYTES(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ /**
+ * How long (in CPU time) it took to serialize the raw messages sent by this mailbox.
+ */
+ SERIALIZATION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },;
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java
index 3415bfb3fc..0085828e4d 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MinusOperator.java
@@ -23,12 +23,15 @@ import javax.annotation.Nullable;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.data.table.Record;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Minus/Except operator.
*/
public class MinusOperator extends SetOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(MinusOperator.class);
private static final String EXPLAIN_NAME = "MINUS";
public MinusOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
@@ -36,6 +39,16 @@ public class MinusOperator extends SetOperator {
super(opChainExecutionContext, upstreamOperators, dataSchema);
}
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.MINUS;
+ }
+
@Nullable
@Override
public String toExplainString() {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
index 581e2f8f10..2e1ba51f54 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
@@ -19,65 +19,75 @@
package org.apache.pinot.query.runtime.operator;
import com.google.common.base.Joiner;
+import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerOperator;
import org.apache.pinot.spi.exception.EarlyTerminationException;
import org.apache.pinot.spi.trace.InvocationScope;
import org.apache.pinot.spi.trace.Tracing;
import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-public abstract class MultiStageOperator implements Operator<TransferableBlock>, AutoCloseable {
- private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageOperator.class);
+public abstract class MultiStageOperator
+ implements Operator<TransferableBlock>, AutoCloseable {
protected final OpChainExecutionContext _context;
protected final String _operatorId;
- protected final OpChainStats _opChainStats;
protected boolean _isEarlyTerminated;
public MultiStageOperator(OpChainExecutionContext context) {
_context = context;
_operatorId = Joiner.on("_").join(getClass().getSimpleName(), _context.getStageId(), _context.getServer());
- _opChainStats = _context.getStats();
_isEarlyTerminated = false;
}
+ /**
+ * Returns the logger for the operator.
+ * <p>
+ * This method is used to generic multi-stage operator messages using the name of the specific operator.
+ * Implementations should not allocate new loggers for each call but instead reuse some (probably static and final)
+ * attribute.
+ */
+ protected abstract Logger logger();
+
+ public abstract Type getOperatorType();
+
+ public abstract void registerExecution(long time, int numRows);
+
@Override
public TransferableBlock nextBlock() {
if (Tracing.ThreadAccountantOps.isInterrupted()) {
throw new EarlyTerminationException("Interrupted while processing next block");
}
+ if (logger().isDebugEnabled()) {
+ logger().debug("Operator {}: Reading next block", _operatorId);
+ }
try (InvocationScope ignored = Tracing.getTracer().createScope(getClass())) {
TransferableBlock nextBlock;
- if (shouldCollectStats()) {
- OperatorStats operatorStats = _opChainStats.getOperatorStats(_context, _operatorId);
- operatorStats.startTimer();
- try {
- nextBlock = getNextBlock();
- } catch (Exception e) {
- nextBlock = TransferableBlockUtils.getErrorTransferableBlock(e);
- }
- operatorStats.recordRow(1, nextBlock.getNumRows());
- operatorStats.endTimer(nextBlock);
- } else {
- try {
- nextBlock = getNextBlock();
- } catch (Exception e) {
- nextBlock = TransferableBlockUtils.getErrorTransferableBlock(e);
- }
+ Stopwatch executeStopwatch = Stopwatch.createStarted();
+ try {
+ nextBlock = getNextBlock();
+ } catch (Exception e) {
+ nextBlock = TransferableBlockUtils.getErrorTransferableBlock(e);
+ }
+ registerExecution(executeStopwatch.elapsed(TimeUnit.MILLISECONDS), nextBlock.getNumRows());
+
+ if (logger().isDebugEnabled()) {
+ logger().debug("Operator {}. Block of type {} ready to send", _operatorId, nextBlock.getType());
}
return nextBlock;
}
}
- public String getOperatorId() {
- return _operatorId;
- }
-
// Make it protected because we should always call nextBlock()
protected abstract TransferableBlock getNextBlock()
throws Exception;
@@ -89,14 +99,24 @@ public abstract class MultiStageOperator implements Operator<TransferableBlock>,
}
}
- protected boolean shouldCollectStats() {
- return _context.isTraceEnabled();
+ /**
+ * Adds the current operator stats as the last operator in the open stats of the given holder.
+ *
+ * It is assumed that:
+ * <ol>
+ * <li>The current stage of the holder is equal to the stage id of this operator.</li>
+ * <li>The holder already contains the stats of the previous operators of the same stage in inorder</li>
+ * </ol>
+ */
+ protected void addStats(MultiStageQueryStats holder, StatMap<?> statMap) {
+ Preconditions.checkArgument(holder.getCurrentStageId() == _context.getStageId(),
+ "The holder's stage id should be the same as the current operator's stage id. Expected %s, got %s",
+ _context.getStageId(), holder.getCurrentStageId());
+ holder.getCurrentStats().addLastOperator(getOperatorType(), statMap);
}
@Override
- public List<MultiStageOperator> getChildOperators() {
- throw new UnsupportedOperationException();
- }
+ public abstract List<MultiStageOperator> getChildOperators();
// TODO: Ideally close() call should finish within request deadline.
// TODO: Consider passing deadline as part of the API.
@@ -106,7 +126,7 @@ public abstract class MultiStageOperator implements Operator<TransferableBlock>,
try {
op.close();
} catch (Exception e) {
- LOGGER.error("Failed to close operator: " + op + " with exception:" + e);
+ logger().error("Failed to close operator: " + op + " with exception:" + e);
// Continue processing because even one operator failed to be close, we should still close the rest.
}
}
@@ -117,9 +137,177 @@ public abstract class MultiStageOperator implements Operator<TransferableBlock>,
try {
op.cancel(e);
} catch (Exception e2) {
- LOGGER.error("Failed to cancel operator:" + op + "with error:" + e + " with exception:" + e2);
+ logger().error("Failed to cancel operator:" + op + "with error:" + e + " with exception:" + e2);
// Continue processing because even one operator failed to be cancelled, we should still cancel the rest.
}
}
}
+
+ /**
+ * Receives the EOS block from upstream operator and updates the stats.
+ * <p>
+ * The fact that the EOS belongs to the upstream operator is not an actual requirement. Actual requirements are listed
+ * in {@link #addStats(MultiStageQueryStats, StatMap)}
+ * @param upstreamEos
+ * @return
+ */
+ protected TransferableBlock updateEosBlock(TransferableBlock upstreamEos, StatMap<?> statMap) {
+ assert upstreamEos.isSuccessfulEndOfStreamBlock();
+ MultiStageQueryStats queryStats = upstreamEos.getQueryStats();
+ assert queryStats != null;
+ addStats(queryStats, statMap);
+ return upstreamEos;
+ }
+
+ /**
+ * This enum is used to identify the operation type.
+ * <p>
+ * This is mostly used in the context of stats collection, where we use this enum in the serialization form in order
+ * to identify the type of the stats in an efficient way.
+ */
+ public enum Type {
+ AGGREGATE(AggregateOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<AggregateOperator.StatKey> stats = (StatMap<AggregateOperator.StatKey>) map;
+ response.mergeNumGroupsLimitReached(stats.getBoolean(AggregateOperator.StatKey.NUM_GROUPS_LIMIT_REACHED));
+ response.mergeMaxRowsInOperator(stats.getLong(AggregateOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ FILTER(FilterOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<FilterOperator.StatKey> stats = (StatMap<FilterOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(FilterOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ HASH_JOIN(HashJoinOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<HashJoinOperator.StatKey> stats = (StatMap<HashJoinOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(HashJoinOperator.StatKey.EMITTED_ROWS));
+ response.mergeMaxRowsInJoinReached(stats.getBoolean(HashJoinOperator.StatKey.MAX_ROWS_IN_JOIN_REACHED));
+ }
+ },
+ INTERSECT(SetOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<SetOperator.StatKey> stats = (StatMap<SetOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(SetOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ LEAF(LeafStageTransferableBlockOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<LeafStageTransferableBlockOperator.StatKey> stats =
+ (StatMap<LeafStageTransferableBlockOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(LeafStageTransferableBlockOperator.StatKey.EMITTED_ROWS));
+
+ StatMap<BrokerResponseNativeV2.StatKey> brokerStats = new StatMap<>(BrokerResponseNativeV2.StatKey.class);
+ for (LeafStageTransferableBlockOperator.StatKey statKey : stats.keySet()) {
+ statKey.updateBrokerMetadata(brokerStats, stats);
+ }
+ response.addServerStats(brokerStats);
+ }
+ },
+ LITERAL(LiteralValueOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ // Do nothing
+ }
+ },
+ MAILBOX_RECEIVE(BaseMailboxReceiveOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<BaseMailboxReceiveOperator.StatKey> stats = (StatMap<BaseMailboxReceiveOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(BaseMailboxReceiveOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ MAILBOX_SEND(MailboxSendOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<MailboxSendOperator.StatKey> stats = (StatMap<MailboxSendOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(MailboxSendOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ MINUS(SetOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<SetOperator.StatKey> stats = (StatMap<SetOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(SetOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ PIPELINE_BREAKER(PipelineBreakerOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<PipelineBreakerOperator.StatKey> stats = (StatMap<PipelineBreakerOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(PipelineBreakerOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ SORT_OR_LIMIT(SortOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<SortOperator.StatKey> stats = (StatMap<SortOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(SortOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ TRANSFORM(TransformOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<TransformOperator.StatKey> stats = (StatMap<TransformOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(TransformOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ UNION(SetOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<SetOperator.StatKey> stats = (StatMap<SetOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(SetOperator.StatKey.EMITTED_ROWS));
+ }
+ },
+ WINDOW(WindowAggregateOperator.StatKey.class) {
+ @Override
+ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
+ @SuppressWarnings("unchecked")
+ StatMap<WindowAggregateOperator.StatKey> stats = (StatMap<WindowAggregateOperator.StatKey>) map;
+ response.mergeMaxRowsInOperator(stats.getLong(WindowAggregateOperator.StatKey.EMITTED_ROWS));
+ }
+ },;
+
+ private final Class _statKeyClass;
+
+ Type(Class<? extends StatMap.Key> statKeyClass) {
+ _statKeyClass = statKeyClass;
+ }
+
+ /**
+ * Gets the class of the stat key for this operator type.
+ * <p>
+ * Notice that this is not including the generic type parameter, because Java generic types are not expressive
+ * enough indicate what we want to say, so generics here are more problematic than useful.
+ */
+ public Class getStatKeyClass() {
+ return _statKeyClass;
+ }
+
+ /**
+ * Merges the stats from the given map into the given broker response.
+ * <p>
+ * Each literal has its own implementation of this method, which assumes the given map is of the correct type
+ * (compatible with {@link #getStatKeyClass()}). This is a way to avoid casting in the caller.
+ */
+ public abstract void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map);
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChain.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChain.java
index 360bef6324..5d989f169f 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChain.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChain.java
@@ -34,7 +34,6 @@ public class OpChain implements AutoCloseable {
private static final Logger LOGGER = LoggerFactory.getLogger(OpChain.class);
private final OpChainId _id;
- private final OpChainStats _stats;
private final MultiStageOperator _root;
private final Consumer<OpChainId> _finishCallback;
@@ -45,7 +44,6 @@ public class OpChain implements AutoCloseable {
public OpChain(OpChainExecutionContext context, MultiStageOperator root, Consumer<OpChainId> finishCallback) {
_id = context.getId();
- _stats = context.getStats();
_root = root;
_finishCallback = finishCallback;
}
@@ -54,11 +52,6 @@ public class OpChain implements AutoCloseable {
return _id;
}
- // TODO: Move OperatorStats here.
- public OpChainStats getStats() {
- return _stats;
- }
-
public Operator<TransferableBlock> getRoot() {
return _root;
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChainStats.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChainStats.java
deleted file mode 100644
index 6c259eaff2..0000000000
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OpChainStats.java
+++ /dev/null
@@ -1,109 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.apache.pinot.query.runtime.operator;
-
-import com.google.common.base.Stopwatch;
-import com.google.common.base.Suppliers;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.Supplier;
-import javax.annotation.concurrent.NotThreadSafe;
-import org.apache.pinot.common.datatable.DataTable;
-import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
-import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
-
-
-/**
- * {@code OpChainStats} tracks execution statistics for {@link OpChain}s.
- */
-@NotThreadSafe
-public class OpChainStats {
-
- // use memoized supplier so that the timing doesn't start until the
- // first time we get the timer
- private final Supplier<ThreadResourceUsageProvider> _exTimer =
- Suppliers.memoize(ThreadResourceUsageProvider::new)::get;
-
- // this is used to make sure that toString() doesn't have side
- // effects (accidentally starting the timer)
- private volatile boolean _exTimerStarted = false;
-
- private final Stopwatch _executeStopwatch = Stopwatch.createUnstarted();
- private final Stopwatch _queuedStopwatch = Stopwatch.createUnstarted();
- private final AtomicLong _queuedCount = new AtomicLong();
-
- private final String _id;
- private final ConcurrentHashMap<String, OperatorStats> _operatorStatsMap = new ConcurrentHashMap<>();
-
- public OpChainStats(String id) {
- _id = id;
- }
-
- public void executing() {
- startExecutionTimer();
- if (_queuedStopwatch.isRunning()) {
- _queuedStopwatch.stop();
- }
- }
-
- public void queued() {
- _queuedCount.incrementAndGet();
- if (!_queuedStopwatch.isRunning()) {
- _queuedStopwatch.start();
- }
- if (_executeStopwatch.isRunning()) {
- _executeStopwatch.stop();
- }
- }
-
- public ConcurrentHashMap<String, OperatorStats> getOperatorStatsMap() {
- return _operatorStatsMap;
- }
-
- public OperatorStats getOperatorStats(OpChainExecutionContext context, String operatorId) {
- return _operatorStatsMap.computeIfAbsent(operatorId, (id) -> {
- OperatorStats operatorStats = new OperatorStats(context);
- if (context.isTraceEnabled()) {
- operatorStats.recordSingleStat(DataTable.MetadataKey.OPERATOR_ID.getName(), operatorId);
- }
- return operatorStats;
- });
- }
-
- private void startExecutionTimer() {
- _exTimerStarted = true;
- _exTimer.get();
- if (!_executeStopwatch.isRunning()) {
- _executeStopwatch.start();
- }
- }
-
- public long getExecutionTime() {
- return _executeStopwatch.elapsed(TimeUnit.MILLISECONDS);
- }
-
- @Override
- public String toString() {
- return String.format("(%s) Queued Count: %s, Executing Time: %sms, Queued Time: %sms", _id, _queuedCount.get(),
- _exTimerStarted ? _executeStopwatch.elapsed(TimeUnit.MILLISECONDS) : 0,
- _queuedStopwatch.elapsed(TimeUnit.MILLISECONDS));
- }
-}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OperatorStats.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OperatorStats.java
deleted file mode 100644
index 32fa9f140f..0000000000
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/OperatorStats.java
+++ /dev/null
@@ -1,119 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-package org.apache.pinot.query.runtime.operator;
-
-import com.google.common.base.Stopwatch;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
-import org.apache.pinot.common.datatable.DataTable;
-import org.apache.pinot.query.routing.VirtualServerAddress;
-import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
-import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
-
-
-public class OperatorStats {
- private final Stopwatch _executeStopwatch = Stopwatch.createUnstarted();
-
- // TODO: add a operatorId for better tracking purpose.
- private final int _stageId;
- private final long _requestId;
-
- private final VirtualServerAddress _serverAddress;
-
- private int _numBlock = 0;
- private int _numRows = 0;
- private long _startTimeMs = -1;
- private long _endTimeMs = -1;
- private final Map<String, String> _executionStats;
-
- public OperatorStats(OpChainExecutionContext context) {
- this(context.getRequestId(), context.getStageId(), context.getServer());
- }
-
- //TODO: remove this constructor after the context constructor can be used in serialization and deserialization
- public OperatorStats(long requestId, int stageId, VirtualServerAddress serverAddress) {
- _stageId = stageId;
- _requestId = requestId;
- _serverAddress = serverAddress;
- _executionStats = new HashMap<>();
- }
-
- public void startTimer() {
- _startTimeMs = _startTimeMs == -1 ? System.currentTimeMillis() : _startTimeMs;
- if (!_executeStopwatch.isRunning()) {
- _executeStopwatch.start();
- }
- }
-
- public void endTimer(TransferableBlock block) {
- if (_executeStopwatch.isRunning()) {
- _executeStopwatch.stop();
- _endTimeMs = System.currentTimeMillis();
- }
- }
-
- public void recordRow(int numBlock, int numRows) {
- _numBlock += numBlock;
- _numRows += numRows;
- }
-
- public void recordSingleStat(String key, String stat) {
- _executionStats.put(key, stat);
- }
-
- public void recordExecutionStats(Map<String, String> executionStats) {
- _executionStats.putAll(executionStats);
- }
-
- public Map<String, String> getExecutionStats() {
- _executionStats.putIfAbsent(DataTable.MetadataKey.NUM_BLOCKS.getName(), String.valueOf(_numBlock));
- _executionStats.putIfAbsent(DataTable.MetadataKey.NUM_ROWS.getName(), String.valueOf(_numRows));
- _executionStats.putIfAbsent(DataTable.MetadataKey.OPERATOR_EXECUTION_TIME_MS.getName(),
- String.valueOf(_executeStopwatch.elapsed(TimeUnit.MILLISECONDS)));
- // wall time are recorded slightly longer than actual execution but it is OK.
-
- if (_startTimeMs != -1) {
- _executionStats.putIfAbsent(DataTable.MetadataKey.OPERATOR_EXEC_START_TIME_MS.getName(),
- String.valueOf(_startTimeMs));
- long endTimeMs = _endTimeMs == -1 ? System.currentTimeMillis() : _endTimeMs;
- _executionStats.putIfAbsent(DataTable.MetadataKey.OPERATOR_EXEC_END_TIME_MS.getName(),
- String.valueOf(endTimeMs));
- }
- return _executionStats;
- }
-
- public int getStageId() {
- return _stageId;
- }
-
- public long getRequestId() {
- return _requestId;
- }
-
- public VirtualServerAddress getServerAddress() {
- return _serverAddress;
- }
-
- @Override
- public String toString() {
- return OperatorUtils.operatorStatsToJson(this);
- }
-}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SetOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SetOperator.java
index edf8416f02..721c81d77f 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SetOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SetOperator.java
@@ -22,13 +22,16 @@ import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
+import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.ExplainPlanRows;
import org.apache.pinot.core.data.table.Record;
import org.apache.pinot.core.operator.ExecutionStatistics;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.segment.spi.IndexSegment;
@@ -50,8 +53,10 @@ public abstract class SetOperator extends MultiStageOperator {
private final DataSchema _dataSchema;
private boolean _isRightSetBuilt;
- private boolean _isTerminated;
- private TransferableBlock _upstreamErrorBlock;
+ protected TransferableBlock _upstreamErrorBlock;
+ @Nullable
+ private MultiStageQueryStats _rightQueryStats = null;
+ protected final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public SetOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
DataSchema dataSchema) {
@@ -62,7 +67,12 @@ public abstract class SetOperator extends MultiStageOperator {
_rightChildOperator = getChildOperators().get(1);
_rightRowSet = new HashSet<>();
_isRightSetBuilt = false;
- _isTerminated = false;
+ }
+
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
}
@Override
@@ -92,9 +102,6 @@ public abstract class SetOperator extends MultiStageOperator {
@Override
protected TransferableBlock getNextBlock() {
- if (_isTerminated) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
- }
if (!_isRightSetBuilt) {
// construct a SET with all the right side rows.
constructRightBlockSet();
@@ -121,6 +128,8 @@ public abstract class SetOperator extends MultiStageOperator {
_upstreamErrorBlock = block;
} else {
_isRightSetBuilt = true;
+ _rightQueryStats = block.getQueryStats();
+ assert _rightQueryStats != null;
}
}
@@ -133,7 +142,12 @@ public abstract class SetOperator extends MultiStageOperator {
return _upstreamErrorBlock;
}
if (leftBlock.isSuccessfulEndOfStreamBlock()) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _rightQueryStats != null;
+ MultiStageQueryStats leftQueryStats = leftBlock.getQueryStats();
+ assert leftQueryStats != null;
+ _rightQueryStats.mergeInOrder(leftQueryStats, getOperatorType(), _statMap);
+ _rightQueryStats.getCurrentStats().concat(leftQueryStats.getCurrentStats());
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_rightQueryStats);
}
for (Object[] row : leftBlock.getContainer()) {
if (handleRowMatched(row)) {
@@ -150,4 +164,29 @@ public abstract class SetOperator extends MultiStageOperator {
* @return true if the row is matched.
*/
protected abstract boolean handleRowMatched(Object[] row);
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ };
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
index b0a1923c80..0c7ba6cf82 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java
@@ -19,6 +19,7 @@
package org.apache.pinot.query.runtime.operator;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import java.util.ArrayList;
import java.util.Arrays;
@@ -27,11 +28,11 @@ import java.util.PriorityQueue;
import javax.annotation.Nullable;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.utils.SortUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -50,8 +51,11 @@ public class SortOperator extends MultiStageOperator {
private final PriorityQueue<Object[]> _priorityQueue;
private final ArrayList<Object[]> _rows;
private final int _numRowsToKeep;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
private boolean _hasConstructedSortedBlock;
+ @Nullable
+ private TransferableBlock _eosBlock = null;
public SortOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator,
List<RexExpression> collationKeys, List<RelFieldCollation.Direction> collationDirections,
@@ -63,8 +67,8 @@ public class SortOperator extends MultiStageOperator {
}
@VisibleForTesting
- SortOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, List<RexExpression> collationKeys,
- List<RelFieldCollation.Direction> collationDirections,
+ SortOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator,
+ List<RexExpression> collationKeys, List<RelFieldCollation.Direction> collationDirections,
List<RelFieldCollation.NullDirection> collationNullDirections, int fetch, int offset, DataSchema dataSchema,
boolean isInputSorted, int defaultHolderCapacity, int defaultResponseLimit) {
super(context);
@@ -91,6 +95,22 @@ public class SortOperator extends MultiStageOperator {
}
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.SORT_OR_LIMIT;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_upstreamOperator);
@@ -109,13 +129,16 @@ public class SortOperator extends MultiStageOperator {
@Override
protected TransferableBlock getNextBlock() {
if (_hasConstructedSortedBlock) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _eosBlock != null;
+ return _eosBlock;
}
TransferableBlock finalBlock = consumeInputBlocks();
// returning upstream error block if finalBlock contains error.
if (finalBlock.isErrorBlock()) {
return finalBlock;
}
+ _statMap.merge(StatKey.REQUIRE_SORT, _priorityQueue != null);
+ _eosBlock = updateEosBlock(finalBlock, _statMap);
return produceSortedBlock();
}
@@ -126,12 +149,12 @@ public class SortOperator extends MultiStageOperator {
List<Object[]> row = _rows.subList(_offset, _rows.size());
return new TransferableBlock(row, _dataSchema, DataBlock.Type.ROW);
} else {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return _eosBlock;
}
} else {
int resultSize = _priorityQueue.size() - _offset;
if (resultSize <= 0) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return _eosBlock;
}
Object[][] rowsArr = new Object[resultSize][];
for (int i = resultSize - 1; i >= 0; i--) {
@@ -154,8 +177,13 @@ public class SortOperator extends MultiStageOperator {
_rows.addAll(container);
} else {
_rows.addAll(container.subList(0, _numRowsToKeep - numRows));
- LOGGER.debug("Early terminate at SortOperator - operatorId={}, opChainId={}", _operatorId,
- _context.getId());
+ if (LOGGER.isDebugEnabled()) {
+ // this operatorId is an old name. It is being kept to avoid breaking changes on the log message.
+ String operatorId = Joiner.on("_")
+ .join(getClass().getSimpleName(), _context.getStageId(), _context.getServer());
+ LOGGER.debug("Early terminate at SortOperator - operatorId={}, opChainId={}", operatorId,
+ _context.getId());
+ }
// setting operator to be early terminated and awaits EOS block next.
earlyTerminate();
}
@@ -169,4 +197,30 @@ public class SortOperator extends MultiStageOperator {
}
return block;
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG),
+ REQUIRE_SORT(StatMap.Type.BOOLEAN) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ };
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
index 8949ad569a..76eb8dcc92 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
@@ -33,6 +33,8 @@ import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.operator.utils.SortUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -43,6 +45,8 @@ import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
* resorting via the PriorityQueue.
*/
public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxReceiveOperator.class);
+
private static final String EXPLAIN_NAME = "SORTED_MAILBOX_RECEIVE";
private final DataSchema _dataSchema;
@@ -66,6 +70,11 @@ public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator {
_isSortOnSender = isSortOnSender;
}
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Nullable
@Override
public String toExplainString() {
@@ -86,8 +95,10 @@ public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator {
return block;
} else {
assert block.isSuccessfulEndOfStreamBlock();
+ // the multiConsumer has already merged stages from upstream, but doesn't know about this operator
+ // specific stats.
+ _eosBlock = updateEosBlock(block, _statMap);
if (!_rows.isEmpty()) {
- _eosBlock = block;
// TODO: This might not be efficient because we are sorting all the received rows. We should use a k-way merge
// when sender side is sorted.
_rows.sort(
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/TransformOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/TransformOperator.java
index 234f32fbcd..143dcb05e2 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/TransformOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/TransformOperator.java
@@ -24,12 +24,15 @@ import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
@@ -43,6 +46,7 @@ import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
* and canonicalized function name matching (lower case).
*/
public class TransformOperator extends MultiStageOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TransformOperator.class);
private static final String EXPLAIN_NAME = "TRANSFORM";
private final MultiStageOperator _upstreamOperator;
@@ -50,6 +54,7 @@ public class TransformOperator extends MultiStageOperator {
private final int _resultColumnSize;
// TODO: Check type matching between resultSchema and the actual result.
private final DataSchema _resultSchema;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public TransformOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator,
DataSchema resultSchema, List<RexExpression> transforms, DataSchema upstreamDataSchema) {
@@ -66,11 +71,27 @@ public class TransformOperator extends MultiStageOperator {
_resultSchema = resultSchema;
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_upstreamOperator);
}
+ @Override
+ public Type getOperatorType() {
+ return Type.TRANSFORM;
+ }
+
@Nullable
@Override
public String toExplainString() {
@@ -81,7 +102,11 @@ public class TransformOperator extends MultiStageOperator {
protected TransferableBlock getNextBlock() {
TransferableBlock block = _upstreamOperator.nextBlock();
if (block.isEndOfStreamBlock()) {
- return block;
+ if (block.isSuccessfulEndOfStreamBlock()) {
+ return updateEosBlock(block, _statMap);
+ } else {
+ return block;
+ }
}
List<Object[]> container = block.getContainer();
List<Object[]> resultRows = new ArrayList<>(container.size());
@@ -94,4 +119,24 @@ public class TransformOperator extends MultiStageOperator {
}
return new TransferableBlock(resultRows, _resultSchema, DataBlock.Type.ROW);
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG);
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnionOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnionOperator.java
index 4f69b575b1..0b7106780e 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnionOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnionOperator.java
@@ -18,25 +18,43 @@
*/
package org.apache.pinot.query.runtime.operator;
+import com.google.common.base.Preconditions;
import java.util.List;
import javax.annotation.Nullable;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Union operator for UNION ALL queries.
*/
public class UnionOperator extends SetOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(UnionOperator.class);
private static final String EXPLAIN_NAME = "UNION";
+ @Nullable
+ private MultiStageQueryStats _queryStats = null;
+ private int _finishedChildren = 0;
public UnionOperator(OpChainExecutionContext opChainExecutionContext, List<MultiStageOperator> upstreamOperators,
DataSchema dataSchema) {
super(opChainExecutionContext, upstreamOperators, dataSchema);
}
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.UNION;
+ }
+
@Nullable
@Override
public String toExplainString() {
@@ -45,13 +63,41 @@ public class UnionOperator extends SetOperator {
@Override
protected TransferableBlock getNextBlock() {
- for (MultiStageOperator upstreamOperator : getChildOperators()) {
+ if (_upstreamErrorBlock != null) {
+ return _upstreamErrorBlock;
+ }
+ List<MultiStageOperator> childOperators = getChildOperators();
+ for (int i = _finishedChildren; i < childOperators.size(); i++) {
+ MultiStageOperator upstreamOperator = childOperators.get(i);
TransferableBlock block = upstreamOperator.nextBlock();
- if (!block.isEndOfStreamBlock()) {
+ if (block.isDataBlock()) {
+ return block;
+ } else if (block.isSuccessfulEndOfStreamBlock()) {
+ _finishedChildren++;
+ consumeEos(block);
+ } else {
+ assert block.isErrorBlock();
+ _upstreamErrorBlock = block;
return block;
}
}
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _queryStats != null : "Should have at least one EOS block from the upstream operators";
+ addStats(_queryStats, _statMap);
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_queryStats);
+ }
+
+ private void consumeEos(TransferableBlock block) {
+ MultiStageQueryStats queryStats = block.getQueryStats();
+ assert queryStats != null;
+ if (_queryStats == null) {
+ Preconditions.checkArgument(queryStats.getCurrentStageId() == _context.getStageId(),
+ "The current stage id of the stats holder: %s does not match the current stage id: %s",
+ queryStats.getCurrentStageId(), _context.getStageId());
+ _queryStats = queryStats;
+ } else {
+ _queryStats.mergeUpstream(queryStats);
+ _queryStats.getCurrentStats().concat(queryStats.getCurrentStats());
+ }
}
@Override
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java
index d2e37598a0..c797607660 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java
@@ -35,6 +35,7 @@ import javax.annotation.Nullable;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.commons.collections.CollectionUtils;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.data.table.Key;
@@ -101,6 +102,9 @@ public class WindowAggregateOperator extends MultiStageOperator {
private int _numRows;
private boolean _hasReturnedWindowAggregateBlock;
+ @Nullable
+ private TransferableBlock _eosBlock = null;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator,
List<RexExpression> groupSet, List<RexExpression> orderSet, List<RelFieldCollation.Direction> orderSetDirection,
@@ -151,11 +155,27 @@ public class WindowAggregateOperator extends MultiStageOperator {
_hasReturnedWindowAggregateBlock = false;
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_inputOperator);
}
+ @Override
+ public Type getOperatorType() {
+ return Type.WINDOW;
+ }
+
@Nullable
@Override
public String toExplainString() {
@@ -165,12 +185,13 @@ public class WindowAggregateOperator extends MultiStageOperator {
@Override
protected TransferableBlock getNextBlock() {
if (_hasReturnedWindowAggregateBlock) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return _eosBlock;
}
TransferableBlock finalBlock = consumeInputBlocks();
if (finalBlock.isErrorBlock()) {
return finalBlock;
}
+ _eosBlock = updateEosBlock(finalBlock, _statMap);
return produceWindowAggregatedBlock();
}
@@ -270,8 +291,8 @@ public class WindowAggregateOperator extends MultiStageOperator {
}
}
_hasReturnedWindowAggregateBlock = true;
- if (rows.size() == 0) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ if (rows.isEmpty()) {
+ return _eosBlock;
} else {
return new TransferableBlock(rows, _resultSchema, DataBlock.Type.ROW);
}
@@ -575,4 +596,29 @@ public class WindowAggregateOperator extends MultiStageOperator {
}
}
}
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ },
+ EMITTED_ROWS(StatMap.Type.LONG) {
+ @Override
+ public boolean includeDefaultInJson() {
+ return true;
+ }
+ };
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java
index f8d49b6328..3c470b1515 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java
@@ -89,6 +89,7 @@ public abstract class BlockExchange {
// Send metadata to only one randomly picked mailbox, and empty EOS block to other mailboxes
int numMailboxes = _sendingMailboxes.size();
int mailboxIdToSendMetadata = ThreadLocalRandom.current().nextInt(numMailboxes);
+ assert block.getQueryStats() != null;
for (int i = 0; i < numMailboxes; i++) {
SendingMailbox sendingMailbox = _sendingMailboxes.get(i);
TransferableBlock blockToSend =
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java
index 387a44e754..145028fc74 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java
@@ -25,6 +25,8 @@ import javax.annotation.Nullable;
import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -56,6 +58,12 @@ public abstract class BlockingMultiStreamConsumer<E> implements AutoCloseable {
protected abstract boolean isEos(E element);
+ /**
+ * This method is called whenever one of the consumer sends a EOS. It is guaranteed that the received element is
+ * an EOS as defined by {@link #isEos(Object)}
+ */
+ protected abstract void onConsumerFinish(E element);
+
protected abstract E onTimeout();
protected abstract E onException(Exception e);
@@ -175,6 +183,7 @@ public abstract class BlockingMultiStreamConsumer<E> implements AutoCloseable {
LOGGER.debug("==[RECEIVE]== EOS received : " + _id + " in mailbox: " + removed.getId()
+ " (" + _mailboxes.size() + " mailboxes alive)");
}
+ onConsumerFinish(block);
block = readBlockOrNull();
}
@@ -225,9 +234,13 @@ public abstract class BlockingMultiStreamConsumer<E> implements AutoCloseable {
}
public static class OfTransferableBlock extends BlockingMultiStreamConsumer<TransferableBlock> {
- public OfTransferableBlock(Object id, long deadlineMs,
+
+ private final MultiStageQueryStats _stats;
+
+ public OfTransferableBlock(OpChainExecutionContext context,
List<? extends AsyncStream<TransferableBlock>> asyncProducers) {
- super(id, deadlineMs, asyncProducers);
+ super(context.getId(), context.getDeadlineMs(), asyncProducers);
+ _stats = MultiStageQueryStats.emptyStats(context.getStageId());
}
@Override
@@ -240,6 +253,15 @@ public abstract class BlockingMultiStreamConsumer<E> implements AutoCloseable {
return element.isSuccessfulEndOfStreamBlock();
}
+ @Override
+ protected void onConsumerFinish(TransferableBlock element) {
+ if (element.getQueryStats() != null) {
+ _stats.mergeUpstream(element.getQueryStats());
+ } else {
+ _stats.mergeUpstream(element.getSerializedStatsByStage());
+ }
+ }
+
@Override
protected TransferableBlock onTimeout() {
return TransferableBlockUtils.getErrorTransferableBlock(QueryException.EXECUTION_TIMEOUT_ERROR);
@@ -252,7 +274,7 @@ public abstract class BlockingMultiStreamConsumer<E> implements AutoCloseable {
@Override
protected TransferableBlock onEos() {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_stats);
}
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/OperatorUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/OperatorUtils.java
index 7998185f8a..a10cde39dc 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/OperatorUtils.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/OperatorUtils.java
@@ -18,23 +18,12 @@
*/
package org.apache.pinot.query.runtime.operator.utils;
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.JsonNode;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.lang.StringUtils;
-import org.apache.pinot.common.datablock.MetadataBlock;
-import org.apache.pinot.common.datatable.DataTable;
-import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
-import org.apache.pinot.query.routing.VirtualServerAddress;
-import org.apache.pinot.query.runtime.operator.OperatorStats;
-import org.apache.pinot.spi.utils.JsonUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
public class OperatorUtils {
- private static final Logger LOGGER = LoggerFactory.getLogger(OperatorUtils.class);
private static final Map<String, String> OPERATOR_TOKEN_MAPPING = new HashMap<>();
static {
@@ -66,70 +55,4 @@ public class OperatorUtils {
functionName = OPERATOR_TOKEN_MAPPING.getOrDefault(functionName, functionName);
return functionName;
}
-
- public static void recordTableName(OperatorStats operatorStats, DispatchablePlanFragment dispatchablePlanFragment) {
- String tableName = dispatchablePlanFragment.getTableName();
- if (tableName != null) {
- operatorStats.recordSingleStat(DataTable.MetadataKey.TABLE.getName(), tableName);
- }
- }
-
- public static String operatorStatsToJson(OperatorStats operatorStats) {
- try {
- Map<String, Object> jsonOut = new HashMap<>();
- jsonOut.put("requestId", operatorStats.getRequestId());
- jsonOut.put("stageId", operatorStats.getStageId());
- jsonOut.put("serverAddress", operatorStats.getServerAddress().toString());
- jsonOut.put("executionStats", operatorStats.getExecutionStats());
- return JsonUtils.objectToString(jsonOut);
- } catch (Exception e) {
- LOGGER.warn("Error occurred while serializing operatorStats: {}", operatorStats, e);
- }
- return null;
- }
-
- public static OperatorStats operatorStatsFromJson(String json) {
- try {
- JsonNode operatorStatsNode = JsonUtils.stringToJsonNode(json);
- long requestId = operatorStatsNode.get("requestId").asLong();
- int stageId = operatorStatsNode.get("stageId").asInt();
- String serverAddressStr = operatorStatsNode.get("serverAddress").asText();
- VirtualServerAddress serverAddress = VirtualServerAddress.parse(serverAddressStr);
-
- OperatorStats operatorStats =
- new OperatorStats(requestId, stageId, serverAddress);
- operatorStats.recordExecutionStats(
- JsonUtils.jsonNodeToObject(operatorStatsNode.get("executionStats"), new TypeReference<Map<String, String>>() {
- }));
-
- return operatorStats;
- } catch (Exception e) {
- LOGGER.warn("Error occurred while deserializing operatorStats: {}", json, e);
- }
- return null;
- }
-
- public static Map<String, OperatorStats> getOperatorStatsFromMetadata(MetadataBlock metadataBlock) {
- Map<String, OperatorStats> operatorStatsMap = new HashMap<>();
- for (Map.Entry<String, String> entry : metadataBlock.getStats().entrySet()) {
- try {
- operatorStatsMap.put(entry.getKey(), operatorStatsFromJson(entry.getValue()));
- } catch (Exception e) {
- LOGGER.warn("Error occurred while fetching operator stats from metadata", e);
- }
- }
- return operatorStatsMap;
- }
-
- public static Map<String, String> getMetadataFromOperatorStats(Map<String, OperatorStats> operatorStatsMap) {
- Map<String, String> metadataStats = new HashMap<>();
- for (Map.Entry<String, OperatorStats> entry : operatorStatsMap.entrySet()) {
- try {
- metadataStats.put(entry.getKey(), operatorStatsToJson(entry.getValue()));
- } catch (Exception e) {
- LOGGER.warn("Error occurred while fetching metadata from operator stats", e);
- }
- }
- return metadataStats;
- }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/MultiStageQueryStats.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/MultiStageQueryStats.java
new file mode 100644
index 0000000000..5a5af37060
--- /dev/null
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/MultiStageQueryStats.java
@@ -0,0 +1,651 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime.plan;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.base.Preconditions;
+import java.io.DataInput;
+import java.io.DataInputStream;
+import java.io.DataOutput;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Objects;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+import org.apache.avro.util.ByteBufferInputStream;
+import org.apache.commons.io.output.UnsynchronizedByteArrayOutputStream;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.query.runtime.operator.BaseMailboxReceiveOperator;
+import org.apache.pinot.query.runtime.operator.LeafStageTransferableBlockOperator;
+import org.apache.pinot.query.runtime.operator.LiteralValueOperator;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * The stats of a given query.
+ * <p>
+ * For the same query, multiple instances of this class may exist. Each of them will have a partial view of the stats.
+ * Specifically, while the query is being executed, each operator will return its own partial view of the stats when
+ * EOS block is sent.
+ * <p>
+ * Simple operations with a single upstream, like filters or transforms, would just add their own information to the
+ * stats. More complex operations, like joins or receiving mailboxes, will merge the stats from all their upstreams and
+ * add their own stats.
+ * <p>
+ * The complete stats for the query are obtained in the execution root (usually the broker) by merging the partial
+ * views.
+ * <p>
+ * In order to reduce allocation, this class is mutable. Some operators may create their own stats, but most of them
+ * will receive a stats object from the upstream operator and modify it by adding their own stats and sometimes merging
+ * them with other upstream stats.
+ */
+public class MultiStageQueryStats {
+ private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryStats.class);
+ private final int _currentStageId;
+ private final StageStats.Open _currentStats;
+ /**
+ * Known stats for stages whose id is higher than the current one.
+ * <p>
+ * A stage may not know all the stats whose id is higher than the current one, so this list may contain null values.
+ * It may also grow in size when different merge methods are called.
+ * <p>
+ * For example the stats of the left hand side of a join may know stats of stages 3 and 4 and the right side may know
+ * stats of stages 5. When merging the stats of the join, the stats of stages 5 will be added to this list.
+ *
+ * @see #mergeUpstream(List)
+ * @see #mergeUpstream(MultiStageQueryStats)
+ * @see #mergeInOrder(MultiStageQueryStats, MultiStageOperator.Type, StatMap)
+ */
+ private final ArrayList<StageStats.Closed> _closedStats;
+ private static final MultiStageOperator.Type[] ALL_TYPES = MultiStageOperator.Type.values();
+
+ private MultiStageQueryStats(int stageId) {
+ _currentStageId = stageId;
+ _currentStats = new StageStats.Open();
+ _closedStats = new ArrayList<>();
+ }
+
+ private static MultiStageQueryStats create(int stageId, MultiStageOperator.Type type, @Nullable StatMap<?> opStats) {
+ MultiStageQueryStats multiStageQueryStats = new MultiStageQueryStats(stageId);
+ multiStageQueryStats.getCurrentStats().addLastOperator(type, opStats);
+ return multiStageQueryStats;
+ }
+
+ public static MultiStageQueryStats emptyStats(int stageId) {
+ return new MultiStageQueryStats(stageId);
+ }
+
+ public static MultiStageQueryStats createLeaf(int stageId,
+ StatMap<LeafStageTransferableBlockOperator.StatKey> opStats) {
+ return create(stageId, MultiStageOperator.Type.LEAF, opStats);
+ }
+
+ public static MultiStageQueryStats createLiteral(int stageId, StatMap<LiteralValueOperator.StatKey> statMap) {
+ return create(stageId, MultiStageOperator.Type.LITERAL, statMap);
+ }
+
+ public static MultiStageQueryStats createCancelledSend(int stageId,
+ StatMap<MailboxSendOperator.StatKey> statMap) {
+ return create(stageId, MultiStageOperator.Type.MAILBOX_SEND, statMap);
+ }
+
+ public static MultiStageQueryStats createReceive(int stageId, StatMap<BaseMailboxReceiveOperator.StatKey> stats) {
+ return create(stageId, MultiStageOperator.Type.MAILBOX_RECEIVE, stats);
+ }
+
+ public int getCurrentStageId() {
+ return _currentStageId;
+ }
+
+ /**
+ * Serialize the current stats in a way it is compatible with {@link #mergeUpstream(List)}.
+ * <p>
+ * The serialized stats are returned in a list where the index is the stage id. Stages downstream or not related to
+ * the current one will be null.
+ */
+ public List<ByteBuffer> serialize()
+ throws IOException {
+
+ ArrayList<ByteBuffer> serializedStats = new ArrayList<>(getMaxStageId());
+ for (int i = 0; i < _currentStageId; i++) {
+ serializedStats.add(null);
+ }
+
+ try (UnsynchronizedByteArrayOutputStream baos = new UnsynchronizedByteArrayOutputStream.Builder().get();
+ DataOutputStream output = new DataOutputStream(baos)) {
+
+ _currentStats.serialize(output);
+ ByteBuffer currentBuf = ByteBuffer.wrap(baos.toByteArray());
+
+ serializedStats.add(currentBuf);
+
+ for (StageStats.Closed closedStats : _closedStats) {
+ if (closedStats == null) {
+ serializedStats.add(null);
+ continue;
+ }
+ baos.reset();
+ closedStats.serialize(output);
+ ByteBuffer buf = ByteBuffer.wrap(baos.toByteArray());
+ serializedStats.add(buf);
+ }
+ }
+ Preconditions.checkState(serializedStats.size() == getMaxStageId() + 1,
+ "Serialized stats size is different from expected size. Expected %s, got %s",
+ getMaxStageId() + 1, serializedStats.size());
+ return serializedStats;
+ }
+
+ public StageStats.Open getCurrentStats() {
+ return _currentStats;
+ }
+
+ /**
+ * Returns the higher stage id known by this object.
+ */
+ public int getMaxStageId() {
+ return _currentStageId + _closedStats.size();
+ }
+
+ /**
+ * Get the stats of a stage whose id is higher than the current one.
+ * <p>
+ * This method returns null in case the stage id is unknown by this stage or no stats are stored for it.
+ */
+ @Nullable
+ public StageStats.Closed getUpstreamStageStats(int stageId) {
+ if (stageId <= _currentStageId) {
+ throw new IllegalArgumentException("Stage " + stageId + " cannot be upstream of current stage "
+ + _currentStageId);
+ }
+
+ int index = stageId - _currentStageId - 1;
+ if (index >= _closedStats.size()) {
+ return null;
+ }
+ return _closedStats.get(index);
+ }
+
+ public void mergeInOrder(MultiStageQueryStats otherStats, MultiStageOperator.Type type,
+ StatMap<?> statMap) {
+ Preconditions.checkArgument(_currentStageId == otherStats._currentStageId,
+ "Cannot merge stats from different stages (%s and %s)", _currentStageId, otherStats._currentStageId);
+ mergeUpstream(otherStats);
+ StageStats.Open currentStats = getCurrentStats();
+ currentStats.concat(otherStats.getCurrentStats());
+ currentStats.addLastOperator(type, statMap);
+ }
+
+ private void growUpToStage(int stageId) {
+ _closedStats.ensureCapacity(stageId - _currentStageId);
+ while (getMaxStageId() < stageId) {
+ _closedStats.add(null);
+ }
+ }
+
+ /**
+ * Merge upstream stats from another MultiStageQueryStats object into this one.
+ * <p>
+ * Only the stages whose id is higher than the current one are merged. The reason to do so is that upstream stats
+ * should be already closed while current stage may need some extra tuning.
+ * <p>
+ * For example set operations may need to merge the stats from all its upstreams before concatenating stats of the
+ * current stage.
+ */
+ public void mergeUpstream(MultiStageQueryStats otherStats) {
+ Preconditions.checkArgument(_currentStageId <= otherStats._currentStageId,
+ "Cannot merge stats from early stage %s into stats of later stage %s",
+ otherStats._currentStageId, _currentStageId);
+
+ growUpToStage(otherStats.getMaxStageId());
+
+ int currentDiff = otherStats._currentStageId - _currentStageId;
+ if (currentDiff > 0) {
+ StageStats.Closed close = otherStats._currentStats.close();
+ int selfIdx = currentDiff - 1;
+ StageStats.Closed myStats = _closedStats.get(selfIdx);
+ if (myStats == null) {
+ _closedStats.set(selfIdx, close);
+ } else {
+ myStats.merge(close);
+ }
+ }
+
+ for (int i = 0; i < otherStats._closedStats.size(); i++) {
+ StageStats.Closed otherStatsForStage = otherStats._closedStats.get(i);
+ if (otherStatsForStage == null) {
+ continue;
+ }
+ int selfIdx = i + currentDiff;
+ StageStats.Closed myStats = _closedStats.get(selfIdx);
+ try {
+ if (myStats == null) {
+ _closedStats.set(selfIdx, otherStatsForStage);
+ assert getUpstreamStageStats(i + otherStats._currentStageId + 1) == otherStatsForStage;
+ } else {
+ myStats.merge(otherStatsForStage);
+ }
+ } catch (IllegalArgumentException | IllegalStateException ex) {
+ LOGGER.warn("Error merging stats on stage " + i + ". Ignoring the new stats", ex);
+ }
+ }
+ }
+
+ public void mergeUpstream(List<ByteBuffer> otherStats) {
+ for (int i = 0; i <= _currentStageId && i < otherStats.size(); i++) {
+ if (otherStats.get(i) != null) {
+ throw new IllegalArgumentException("Cannot merge stats from early stage " + i + " into stats of "
+ + "later stage " + _currentStageId);
+ }
+ }
+ growUpToStage(otherStats.size() - 1);
+
+ for (int i = _currentStageId + 1; i < otherStats.size(); i++) {
+ ByteBuffer otherBuf = otherStats.get(i);
+ if (otherBuf != null) {
+ StageStats.Closed myStats = getUpstreamStageStats(i);
+ try (InputStream is = new ByteBufferInputStream(Collections.singletonList(otherBuf));
+ DataInputStream dis = new DataInputStream(is)) {
+ if (myStats == null) {
+ StageStats.Closed deserialized = StageStats.Closed.deserialize(dis);
+ _closedStats.set(i - _currentStageId - 1, deserialized);
+ assert getUpstreamStageStats(i) == deserialized;
+ } else {
+ myStats.merge(dis);
+ }
+ } catch (IOException ex) {
+ LOGGER.warn("Error deserializing stats on stage " + i + ". Considering the new stats empty", ex);
+ } catch (IllegalArgumentException | IllegalStateException ex) {
+ LOGGER.warn("Error merging stats on stage " + i + ". Ignoring the new stats", ex);
+ }
+ }
+ }
+ }
+
+ public JsonNode asJson() {
+ ObjectNode node = JsonUtils.newObjectNode();
+ node.put("stage", _currentStageId);
+ node.set("open", _currentStats.asJson());
+
+ ArrayNode closedStats = JsonUtils.newArrayNode();
+ for (StageStats.Closed closed : _closedStats) {
+ if (closed == null) {
+ closedStats.addNull();
+ } else {
+ closedStats.add(closed.asJson());
+ }
+ }
+ node.set("closed", closedStats);
+ return node;
+ }
+
+ @Override
+ public String toString() {
+ return asJson().toString();
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ MultiStageQueryStats that = (MultiStageQueryStats) o;
+ return _currentStageId == that._currentStageId && Objects.equals(_currentStats, that._currentStats)
+ && Objects.equals(_closedStats, that._closedStats);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(_currentStageId, _currentStats, _closedStats);
+ }
+
+ /**
+ * {@code StageStats} tracks execution statistics for a single stage.
+ * <p>
+ * Instances of this class may not have the complete stat information for the stage. Specifically, while the query
+ * is being executed, each OpChain will contain its own partial view of the stats.
+ * The final stats for the stage are obtained in the execution root (usually the broker) by
+ * {@link Closed#merge(StageStats) merging} the partial views from all OpChains.
+ */
+ public abstract static class StageStats {
+
+ /**
+ * The types of the operators in the stage.
+ * <p>
+ * The operator index used here is the index of the operator in the operator tree, in the order of the inorder
+ * traversal. That means that the first value is the leftmost leaf, and the last value is the root.
+ * <p>
+ * This list contains no null values.
+ */
+ protected final List<MultiStageOperator.Type> _operatorTypes;
+ /**
+ * The stats associated with the given operator index.
+ * <p>
+ * The operator index used here is the index of the operator in the operator tree, in the order of the inorder
+ * traversal. That means that the first value is the leftmost leaf, and the last value is the root.
+ * <p>
+ * This list contains no null values.
+ */
+ protected final List<StatMap<?>> _operatorStats;
+
+ private StageStats() {
+ this(new ArrayList<>(), new ArrayList<>());
+ }
+
+ private StageStats(List<MultiStageOperator.Type> operatorTypes, List<StatMap<?>> operatorStats) {
+ Preconditions.checkArgument(operatorTypes.size() == operatorStats.size(),
+ "Operator types and stats must have the same size (%s != %s)",
+ operatorTypes.size(), operatorStats.size());
+ for (int i = 0; i < operatorTypes.size(); i++) {
+ if (operatorTypes.get(i) == null) {
+ throw new IllegalArgumentException("Unexpected null operator type at index " + i);
+ }
+ }
+ for (int i = 0; i < operatorStats.size(); i++) {
+ if (operatorStats.get(i) == null) {
+ throw new IllegalArgumentException("Unexpected null operator stats of type " + operatorTypes.get(i)
+ + " at index " + i);
+ }
+ }
+ _operatorTypes = operatorTypes;
+ _operatorStats = operatorStats;
+ }
+
+ /**
+ * Return the stats associated with the given operator index.
+ * <p>
+ * The operator index used here is the index of the operator in the operator tree, in the order of the inorder
+ * traversal.
+ * That means that the first value is the leftmost leaf, and the last value is the root.
+ * <p>
+ * It is the operator responsibility to store here its own stats and that must be done just before the end of stream
+ * block is sent. This means that calling this method before the stats is added will throw an index out of bounds.
+ *
+ * @param operatorIdx The operator index in inorder traversal of the operator tree.
+ * @return The value of the stat or null if no stat is registered.
+ * @throws IndexOutOfBoundsException if there is no stats for the given operator index.
+ */
+ public StatMap<?> getOperatorStats(int operatorIdx) {
+ return _operatorStats.get(operatorIdx);
+ }
+
+ public MultiStageOperator.Type getOperatorType(int index) {
+ return _operatorTypes.get(index);
+ }
+
+ public MultiStageOperator.Type getLastType() {
+ return _operatorTypes.get(_operatorTypes.size() - 1);
+ }
+
+ public StatMap<?> getLastOperatorStats() {
+ return _operatorStats.get(_operatorStats.size() - 1);
+ }
+
+ public int getLastOperatorIndex() {
+ return _operatorStats.size() - 1;
+ }
+
+ public JsonNode asJson() {
+ ArrayNode json = JsonUtils.newArrayNode();
+
+ for (int i = 0; i < _operatorStats.size(); i++) {
+ ObjectNode statNode = JsonUtils.newObjectNode();
+ statNode.put("type", _operatorTypes.get(i).name());
+ StatMap<?> stats = _operatorStats.get(i);
+ if (!stats.isEmpty()) {
+ statNode.set("stats", stats.asJson());
+ }
+ json.add(statNode);
+ }
+
+ return json;
+ }
+
+ /**
+ * Serialize the stats to the given output.
+ * <p>
+ * Stats can be then deserialized as {@link Closed} with {@link Closed#deserialize(DataInput)}.
+ */
+ public void serialize(DataOutput output)
+ throws IOException {
+ // TODO: we can serialize with short or variable size
+ output.writeInt(_operatorTypes.size());
+ assert MultiStageOperator.Type.values().length < Byte.MAX_VALUE : "Too many operator types. "
+ + "Need to increase the number of bytes size per operator type";
+ for (int i = 0; i < _operatorTypes.size(); i++) {
+ output.writeByte(_operatorTypes.get(i).ordinal());
+ StatMap<?> statMap = _operatorStats.get(i);
+ statMap.serialize(output);
+ }
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ StageStats that = (StageStats) o;
+ return Objects.equals(_operatorTypes, that._operatorTypes) && Objects.equals(_operatorStats, that._operatorStats);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(_operatorTypes, _operatorStats);
+ }
+
+ @Override
+ public String toString() {
+ return asJson().toString();
+ }
+
+ /**
+ * Closed stats represent the stats of upstream stages.
+ * <p>
+ * These stats can be serialized or deserialized and merged with other stats of the same stage, but operators must
+ * never add entries for itself in upstream stats.
+ */
+ public static class Closed extends StageStats {
+ public Closed(List<MultiStageOperator.Type> operatorTypes, List<StatMap<?>> operatorStats) {
+ super(operatorTypes, operatorStats);
+ }
+
+ /**
+ * Merges the stats from another StageStats object into this one.
+ * <p>
+ * This object is modified in place while the other object is not modified.
+ * Both stats must belong to the same stage.
+ */
+ // We need to deal with unchecked because Java type system is not expressive enough to handle this at static time
+ // But we know we are merging the same stat types because they were created for the same operation.
+ // There is also a dynamic check in the StatMap.merge() method.
+ @SuppressWarnings("unchecked")
+ public void merge(StageStats other) {
+ Preconditions.checkState(_operatorTypes.equals(other._operatorTypes), "Cannot merge stats from "
+ + "different stages. Found types %s and %s", _operatorTypes, other._operatorTypes);
+ for (int i = 0; i < _operatorStats.size(); i++) {
+ StatMap otherStats = other._operatorStats.get(i);
+ StatMap myStats = _operatorStats.get(i);
+ myStats.merge(otherStats);
+ }
+ }
+
+ public void merge(DataInputStream input)
+ throws IOException {
+ int numOperators = input.readInt();
+ if (numOperators != _operatorTypes.size()) {
+ try {
+ Closed deserialized = deserialize(input, numOperators);
+ throw new RuntimeException("Cannot merge stats from stages with different operators. Expected "
+ + _operatorTypes + " operators, got " + numOperators + ". Deserialized stats: " + deserialized);
+ } catch (IOException e) {
+ throw new IOException("Cannot merge stats from stages with different operators. Expected "
+ + _operatorTypes + " operators, got " + numOperators, e);
+ } catch (RuntimeException e) {
+ throw new RuntimeException("Cannot merge stats from stages with different operators. Expected "
+ + _operatorTypes + " operators, got " + numOperators, e);
+ }
+ }
+ for (int i = 0; i < numOperators; i++) {
+ byte ordinal = input.readByte();
+ if (ordinal != _operatorTypes.get(i).ordinal()) {
+ throw new IllegalStateException("Cannot merge stats from stages with different operators. Expected "
+ + " operator " + _operatorTypes.get(i) + "at index " + i + ", got " + ordinal);
+ }
+ _operatorStats.get(i).merge(input);
+ }
+ }
+
+ /**
+ * Same as {@link #merge(StageStats)} but reads the stats from a DataInput, so it should be slightly faster given
+ * it doesn't need to create new objects.
+ */
+ public static Closed deserialize(DataInput input)
+ throws IOException {
+ return deserialize(input, input.readInt());
+ }
+
+ public void forEach(BiConsumer<MultiStageOperator.Type, StatMap<?>> consumer) {
+ Iterator<MultiStageOperator.Type> typeIterator = _operatorTypes.iterator();
+ Iterator<StatMap<?>> statIterator = _operatorStats.iterator();
+ while (typeIterator.hasNext()) {
+ consumer.accept(typeIterator.next(), statIterator.next());
+ }
+ }
+ }
+
+ public static Closed deserialize(DataInput input, int numOperators)
+ throws IOException {
+ List<MultiStageOperator.Type> operatorTypes = new ArrayList<>(numOperators);
+ List<StatMap<?>> operatorStats = new ArrayList<>(numOperators);
+
+ MultiStageOperator.Type[] allTypes = ALL_TYPES;
+ try {
+ for (int i = 0; i < numOperators; i++) {
+ byte ordinal = input.readByte();
+ if (ordinal < 0 || ordinal >= allTypes.length) {
+ throw new IllegalStateException(
+ "Invalid operator type ordinal " + ordinal + " at index " + i + ". " + "Deserialized so far: "
+ + new Closed(operatorTypes, operatorStats));
+ }
+ MultiStageOperator.Type type = allTypes[ordinal];
+ operatorTypes.add(type);
+
+ @SuppressWarnings("unchecked")
+ StatMap<?> opStatMap = StatMap.deserialize(input, type.getStatKeyClass());
+ operatorStats.add(opStatMap);
+ }
+ return new Closed(operatorTypes, operatorStats);
+ } catch (IOException e) {
+ throw new IOException("Error deserializing stats. Deserialized so far: "
+ + new Closed(operatorTypes.subList(0, operatorStats.size()), operatorStats), e);
+ } catch (RuntimeException e) {
+ throw new RuntimeException("Error deserializing stats. Deserialized so far: "
+ + new Closed(operatorTypes.subList(0, operatorStats.size()), operatorStats), e);
+ }
+ }
+
+ /**
+ * Open stats represent the stats of the current stage.
+ * <p>
+ * These stats can be modified by the operator that is currently executing. Specifically they can add stats for
+ * the current operator or merge with other open stats from the same stage.
+ */
+ public static class Open extends StageStats {
+ private Open() {
+ super();
+ }
+
+ public Open addLastOperator(MultiStageOperator.Type type, StatMap<?> statMap) {
+ Preconditions.checkArgument(statMap.getKeyClass().equals(type.getStatKeyClass()),
+ "Expected stats of class %s for type %s but found class %s",
+ type.getStatKeyClass(), type, statMap.getKeyClass());
+ if (!_operatorStats.isEmpty() && _operatorStats.get(_operatorStats.size() - 1) == statMap) {
+ // This is mostly useful to detect errors in the code.
+ // In the future we may choose to evaluate it only if asserts are enabled
+ throw new IllegalArgumentException("Cannot add the same stat map twice.");
+ }
+ Preconditions.checkNotNull(type, "Cannot add null operator type");
+ Preconditions.checkNotNull(statMap, "Cannot add null stats");
+ _operatorTypes.add(type);
+ _operatorStats.add(statMap);
+ return this;
+ }
+
+ /**
+ * Adds the given stats at the end of this object.
+ */
+ public void concat(StageStats.Open other) {
+ _operatorTypes.addAll(other._operatorTypes);
+ _operatorStats.addAll(other._operatorStats);
+ }
+
+ public Closed close() {
+ return new Closed(_operatorTypes, _operatorStats);
+ }
+ }
+ }
+
+ public static class Builder {
+ private final MultiStageQueryStats _stats;
+
+ public Builder(int stageId) {
+ _stats = new MultiStageQueryStats(stageId);
+ }
+
+ public Builder customizeOpen(Consumer<StageStats.Open> customizer) {
+ customizer.accept(_stats._currentStats);
+ return this;
+ }
+
+ /**
+ * Adds a new operator to the stats.
+ * @param consumer a function that will be called with a new and empty open stats object and returns the closed stat
+ * to be added. The received object can be freely modified and close.
+ */
+ public Builder addLast(Function<StageStats.Open, StageStats.Closed> consumer) {
+ StageStats.Open open = new StageStats.Open();
+ _stats._closedStats.add(consumer.apply(open));
+ return this;
+ }
+
+ public MultiStageQueryStats build() {
+ return _stats;
+ }
+ }
+}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java
index 51d61e5dbc..3290478de8 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java
@@ -20,12 +20,12 @@ package org.apache.pinot.query.runtime.plan;
import java.util.Collections;
import java.util.Map;
+import javax.annotation.Nullable;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.routing.StageMetadata;
import org.apache.pinot.query.routing.VirtualServerAddress;
import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.operator.OpChainId;
-import org.apache.pinot.query.runtime.operator.OpChainStats;
import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerResult;
import org.apache.pinot.query.runtime.plan.server.ServerPlanRequestContext;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -45,7 +45,7 @@ public class OpChainExecutionContext {
private final WorkerMetadata _workerMetadata;
private final VirtualServerAddress _server;
private final OpChainId _id;
- private final OpChainStats _stats;
+ @Nullable
private final PipelineBreakerResult _pipelineBreakerResult;
private final boolean _traceEnabled;
@@ -53,7 +53,7 @@ public class OpChainExecutionContext {
public OpChainExecutionContext(MailboxService mailboxService, long requestId, long deadlineMs,
Map<String, String> opChainMetadata, StageMetadata stageMetadata, WorkerMetadata workerMetadata,
- PipelineBreakerResult pipelineBreakerResult) {
+ @Nullable PipelineBreakerResult pipelineBreakerResult) {
_mailboxService = mailboxService;
_requestId = requestId;
_deadlineMs = deadlineMs;
@@ -63,11 +63,7 @@ public class OpChainExecutionContext {
_server =
new VirtualServerAddress(mailboxService.getHostname(), mailboxService.getPort(), workerMetadata.getWorkerId());
_id = new OpChainId(requestId, workerMetadata.getWorkerId(), stageMetadata.getStageId());
- _stats = new OpChainStats(_id.toString());
_pipelineBreakerResult = pipelineBreakerResult;
- if (pipelineBreakerResult != null && pipelineBreakerResult.getOpChainStats() != null) {
- _stats.getOperatorStatsMap().putAll(pipelineBreakerResult.getOpChainStats().getOperatorStatsMap());
- }
_traceEnabled = Boolean.parseBoolean(opChainMetadata.get(CommonConstants.Broker.Request.TRACE));
}
@@ -111,10 +107,7 @@ public class OpChainExecutionContext {
return _id;
}
- public OpChainStats getStats() {
- return _stats;
- }
-
+ @Nullable
public PipelineBreakerResult getPipelineBreakerResult() {
return _pipelineBreakerResult;
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutor.java
index a033df03d7..2e0cc7003d 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutor.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutor.java
@@ -121,7 +121,7 @@ public class PipelineBreakerExecutor {
long timeoutMs = opChainExecutionContext.getDeadlineMs() - System.currentTimeMillis();
if (latch.await(timeoutMs, TimeUnit.MILLISECONDS)) {
return new PipelineBreakerResult(pipelineBreakerContext.getNodeIdMap(), pipelineBreakerOperator.getResultMap(),
- pipelineBreakerOperator.getErrorBlock(), pipelineBreakerOpChain.getStats());
+ pipelineBreakerOperator.getErrorBlock(), pipelineBreakerOperator.getQueryStats());
} else {
throw new TimeoutException(
String.format("Timed out waiting for pipeline breaker results after: %dms", timeoutMs));
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
index 9fe2588827..ee3fd72101 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.query.runtime.plan.pipeline;
+import com.google.common.base.Preconditions;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
@@ -26,20 +27,28 @@ import java.util.List;
import java.util.Map;
import java.util.Queue;
import javax.annotation.Nullable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
-class PipelineBreakerOperator extends MultiStageOperator {
+public class PipelineBreakerOperator extends MultiStageOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(PipelineBreakerOperator.class);
private static final String EXPLAIN_NAME = "PIPELINE_BREAKER";
private final Map<Integer, Operator<TransferableBlock>> _workerMap;
private Map<Integer, List<TransferableBlock>> _resultMap;
private TransferableBlock _errorBlock;
+ @Nullable
+ private MultiStageQueryStats _queryStats = null;
+ private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
public PipelineBreakerOperator(OpChainExecutionContext context, Map<Integer, Operator<TransferableBlock>> workerMap) {
super(context);
@@ -50,6 +59,33 @@ class PipelineBreakerOperator extends MultiStageOperator {
}
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ public List<MultiStageOperator> getChildOperators() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.PIPELINE_BREAKER;
+ }
+
+ public MultiStageQueryStats getQueryStats() {
+ assert _queryStats != null || _errorBlock != null
+ : "This method should not be called before blocks have been processed";
+ return _queryStats;
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
public Map<Integer, List<TransferableBlock>> getResultMap() {
return _resultMap;
}
@@ -84,6 +120,7 @@ class PipelineBreakerOperator extends MultiStageOperator {
dataBlocks.add(block);
block = operator.nextBlock();
}
+ _queryStats = block.getQueryStats();
} else {
_resultMap = new HashMap<>();
for (int workerKey : _workerMap.keySet()) {
@@ -101,9 +138,37 @@ class PipelineBreakerOperator extends MultiStageOperator {
if (block.isDataBlock()) {
_resultMap.get(entry.getKey()).add(block);
entries.offer(entry);
+ } else if (block.isSuccessfulEndOfStreamBlock()) {
+ MultiStageQueryStats queryStats = block.getQueryStats();
+ assert queryStats != null;
+ if (_queryStats == null) {
+ Preconditions.checkArgument(queryStats.getCurrentStageId() == _context.getStageId(),
+ "The current stage id of the stats holder: %s does not match the current stage id: %s",
+ queryStats.getCurrentStageId(), _context.getStageId());
+ _queryStats = queryStats;
+ } else {
+ _queryStats.mergeUpstream(queryStats);
+ }
}
}
}
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ assert _queryStats != null;
+ addStats(_queryStats, _statMap);
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(_queryStats);
+ }
+
+ public enum StatKey implements StatMap.Key {
+ EXECUTION_TIME_MS(StatMap.Type.LONG),
+ EMITTED_ROWS(StatMap.Type.LONG);
+ private final StatMap.Type _type;
+
+ StatKey(StatMap.Type type) {
+ _type = type;
+ }
+
+ @Override
+ public StatMap.Type getType() {
+ return _type;
+ }
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerResult.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerResult.java
index 2e2b003e34..fef6df2cb1 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerResult.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerResult.java
@@ -23,7 +23,7 @@ import java.util.Map;
import javax.annotation.Nullable;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.operator.OpChainStats;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
/**
@@ -33,14 +33,14 @@ public class PipelineBreakerResult {
private final Map<PlanNode, Integer> _nodeIdMap;
private final Map<Integer, List<TransferableBlock>> _resultMap;
private final TransferableBlock _errorBlock;
- private final OpChainStats _opChainStats;
+ private final MultiStageQueryStats _multiStageQueryStats;
public PipelineBreakerResult(Map<PlanNode, Integer> nodeIdMap, Map<Integer, List<TransferableBlock>> resultMap,
- @Nullable TransferableBlock errorBlock, @Nullable OpChainStats opChainStats) {
+ @Nullable TransferableBlock errorBlock, @Nullable MultiStageQueryStats multiStageQueryStats) {
_nodeIdMap = nodeIdMap;
_resultMap = resultMap;
_errorBlock = errorBlock;
- _opChainStats = opChainStats;
+ _multiStageQueryStats = multiStageQueryStats;
}
public Map<PlanNode, Integer> getNodeIdMap() {
@@ -57,7 +57,7 @@ public class PipelineBreakerResult {
}
@Nullable
- public OpChainStats getOpChainStats() {
- return _opChainStats;
+ public MultiStageQueryStats getStageQueryStats() {
+ return _multiStageQueryStats;
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestContext.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestContext.java
index cc9d131962..ac91b45707 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestContext.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestContext.java
@@ -20,6 +20,7 @@ package org.apache.pinot.query.runtime.plan.server;
import java.util.List;
import java.util.concurrent.ExecutorService;
+import javax.annotation.Nullable;
import org.apache.pinot.common.request.PinotQuery;
import org.apache.pinot.core.query.executor.QueryExecutor;
import org.apache.pinot.core.query.request.ServerQueryRequest;
@@ -39,6 +40,7 @@ public class ServerPlanRequestContext {
private final StagePlan _stagePlan;
private final QueryExecutor _leafQueryExecutor;
private final ExecutorService _executorService;
+ @Nullable
private final PipelineBreakerResult _pipelineBreakerResult;
private final PinotQuery _pinotQuery;
@@ -46,7 +48,7 @@ public class ServerPlanRequestContext {
private List<ServerQueryRequest> _serverQueryRequests;
public ServerPlanRequestContext(StagePlan stagePlan, QueryExecutor leafQueryExecutor,
- ExecutorService executorService, PipelineBreakerResult pipelineBreakerResult) {
+ ExecutorService executorService, @Nullable PipelineBreakerResult pipelineBreakerResult) {
_stagePlan = stagePlan;
_leafQueryExecutor = leafQueryExecutor;
_executorService = executorService;
@@ -66,6 +68,7 @@ public class ServerPlanRequestContext {
return _executorService;
}
+ @Nullable
public PipelineBreakerResult getPipelineBreakerResult() {
return _pipelineBreakerResult;
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
index 99a94d0a24..47b563b1fb 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java
@@ -36,15 +36,12 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
-import javax.annotation.Nullable;
import org.apache.calcite.runtime.PairList;
-import org.apache.commons.collections.MapUtils;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.proto.Worker;
import org.apache.pinot.common.response.broker.ResultTable;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
-import org.apache.pinot.core.query.reduce.ExecutionStatsAggregator;
import org.apache.pinot.core.util.DataBlockExtractUtils;
import org.apache.pinot.core.util.trace.TracedThreadFactory;
import org.apache.pinot.query.mailbox.MailboxService;
@@ -62,9 +59,7 @@ import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.MailboxReceiveOperator;
-import org.apache.pinot.query.runtime.operator.OpChainStats;
-import org.apache.pinot.query.runtime.operator.OperatorStats;
-import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.trace.RequestContext;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -89,18 +84,17 @@ public class QueryDispatcher {
new TracedThreadFactory(Thread.NORM_PRIORITY, false, PINOT_BROKER_QUERY_DISPATCHER_FORMAT));
}
- public ResultTable submitAndReduce(RequestContext context, DispatchableSubPlan dispatchableSubPlan, long timeoutMs,
- Map<String, String> queryOptions, @Nullable Map<Integer, ExecutionStatsAggregator> executionStatsAggregator)
+ public QueryResult submitAndReduce(RequestContext context, DispatchableSubPlan dispatchableSubPlan, long timeoutMs,
+ Map<String, String> queryOptions)
throws Exception {
long requestId = context.getRequestId();
try {
submit(requestId, dispatchableSubPlan, timeoutMs, queryOptions);
long reduceStartTimeNs = System.nanoTime();
- ResultTable resultTable =
- runReducer(requestId, dispatchableSubPlan, timeoutMs, queryOptions, executionStatsAggregator,
- _mailboxService);
+ QueryResult queryResult =
+ runReducer(requestId, dispatchableSubPlan, timeoutMs, queryOptions, _mailboxService);
context.setReduceTimeNanos(System.nanoTime() - reduceStartTimeNs);
- return resultTable;
+ return queryResult;
} catch (Throwable e) {
// TODO: Consider always cancel when it returns (early terminate)
cancel(requestId, dispatchableSubPlan);
@@ -252,9 +246,8 @@ public class QueryDispatcher {
}
@VisibleForTesting
- public static ResultTable runReducer(long requestId, DispatchableSubPlan dispatchableSubPlan, long timeoutMs,
- Map<String, String> queryOptions, @Nullable Map<Integer, ExecutionStatsAggregator> statsAggregatorMap,
- MailboxService mailboxService) {
+ public static QueryResult runReducer(long requestId, DispatchableSubPlan dispatchableSubPlan, long timeoutMs,
+ Map<String, String> queryOptions, MailboxService mailboxService) {
// NOTE: Reduce stage is always stage 0
DispatchablePlanFragment dispatchableStagePlan = dispatchableSubPlan.getQueryStageList().get(0);
PlanFragment planFragment = dispatchableStagePlan.getPlanFragment();
@@ -272,29 +265,10 @@ public class QueryDispatcher {
MailboxReceiveOperator receiveOperator =
new MailboxReceiveOperator(opChainExecutionContext, receiveNode.getDistributionType(),
receiveNode.getSenderStageId());
- ResultTable resultTable =
- getResultTable(receiveOperator, receiveNode.getDataSchema(), dispatchableSubPlan.getQueryResultFields());
- collectStats(dispatchableSubPlan, opChainExecutionContext.getStats(), statsAggregatorMap);
- return resultTable;
+ return getQueryResult(receiveOperator, receiveNode.getDataSchema(), dispatchableSubPlan.getQueryResultFields());
}
- private static void collectStats(DispatchableSubPlan dispatchableSubPlan, OpChainStats opChainStats,
- @Nullable Map<Integer, ExecutionStatsAggregator> statsAggregatorMap) {
- if (MapUtils.isNotEmpty(statsAggregatorMap)) {
- for (OperatorStats operatorStats : opChainStats.getOperatorStatsMap().values()) {
- ExecutionStatsAggregator rootStatsAggregator = statsAggregatorMap.get(0);
- rootStatsAggregator.aggregate(null, operatorStats.getExecutionStats(), new HashMap<>());
- ExecutionStatsAggregator stageStatsAggregator = statsAggregatorMap.get(operatorStats.getStageId());
- if (stageStatsAggregator != null) {
- OperatorUtils.recordTableName(operatorStats,
- dispatchableSubPlan.getQueryStageList().get(operatorStats.getStageId()));
- stageStatsAggregator.aggregate(null, operatorStats.getExecutionStats(), new HashMap<>());
- }
- }
- }
- }
-
- private static ResultTable getResultTable(MailboxReceiveOperator receiveOperator, DataSchema sourceDataSchema,
+ private static QueryResult getQueryResult(MailboxReceiveOperator receiveOperator, DataSchema sourceDataSchema,
PairList<Integer, String> resultFields) {
int numColumns = resultFields.size();
String[] columnNames = new String[numColumns];
@@ -308,6 +282,7 @@ public class QueryDispatcher {
ArrayList<Object[]> resultRows = new ArrayList<>();
TransferableBlock block = receiveOperator.nextBlock();
+
while (!TransferableBlockUtils.isEndOfStream(block)) {
DataBlock dataBlock = block.getDataBlock();
int numRows = dataBlock.getNumberOfRows();
@@ -328,11 +303,16 @@ public class QueryDispatcher {
}
block = receiveOperator.nextBlock();
}
+ MultiStageQueryStats queryStats;
if (block.isErrorBlock()) {
throw new RuntimeException("Received error query execution result block: " + block.getExceptions());
+ } else {
+ assert block.isSuccessfulEndOfStreamBlock();
+ queryStats = block.getQueryStats();
+ assert queryStats != null;
}
- return new ResultTable(resultDataSchema, resultRows);
+ return new QueryResult(new ResultTable(resultDataSchema, resultRows), queryStats);
}
public void shutdown() {
@@ -341,4 +321,29 @@ public class QueryDispatcher {
}
_dispatchClientMap.clear();
}
+
+ public static class QueryResult {
+ private final ResultTable _resultTable;
+ private final List<MultiStageQueryStats.StageStats.Closed> _queryStats;
+
+ public QueryResult(ResultTable resultTable, MultiStageQueryStats queryStats) {
+ _resultTable = resultTable;
+
+ Preconditions.checkArgument(queryStats.getCurrentStageId() == 0,
+ "Expecting query stats for stage 0, got: %s", queryStats.getCurrentStageId());
+ _queryStats = new ArrayList<>(queryStats.getMaxStageId());
+ _queryStats.add(queryStats.getCurrentStats().close());
+ for (int i = 1; i <= queryStats.getMaxStageId(); i++) {
+ _queryStats.add(queryStats.getUpstreamStageStats(i));
+ }
+ }
+
+ public ResultTable getResultTable() {
+ return _resultTable;
+ }
+
+ public List<MultiStageQueryStats.StageStats.Closed> getQueryStats() {
+ return _queryStats;
+ }
+ }
}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/MailboxServiceTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/MailboxServiceTest.java
index 3f20d33956..4991448fc1 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/MailboxServiceTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/MailboxServiceTest.java
@@ -22,11 +22,13 @@ import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.physical.MailboxIdUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.query.runtime.operator.OperatorTestUtil;
import org.apache.pinot.query.testutils.QueryTestUtils;
import org.apache.pinot.spi.env.PinotConfiguration;
@@ -34,6 +36,7 @@ import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.util.TestUtils;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeTest;
import org.testng.annotations.Test;
import static org.testng.Assert.*;
@@ -49,6 +52,7 @@ public class MailboxServiceTest {
private MailboxService _mailboxService2;
private long _requestId = 0;
+ private StatMap<MailboxSendOperator.StatKey> _stats;
@BeforeClass
public void setUp() {
@@ -67,6 +71,11 @@ public class MailboxServiceTest {
_mailboxService2.shutdown();
}
+ @BeforeTest
+ public void setUpStats() {
+ _stats = new StatMap<>(MailboxSendOperator.StatKey.class);
+ }
+
@Test
public void testLocalHappyPathSendFirst()
throws Exception {
@@ -74,11 +83,11 @@ public class MailboxServiceTest {
// Sends are non-blocking as long as channel capacity is not breached
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
for (int i = 0; i < ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS - 1; i++) {
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{i}));
}
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
sendingMailbox.complete();
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
@@ -113,11 +122,11 @@ public class MailboxServiceTest {
// Sends are non-blocking as long as channel capacity is not breached
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
for (int i = 0; i < ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS - 1; i++) {
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{i}));
}
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
sendingMailbox.complete();
assertEquals(numCallbacks.get(), ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS);
@@ -143,7 +152,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
receivingMailbox.registeredReader(numCallbacks::getAndIncrement);
@@ -169,8 +178,9 @@ public class MailboxServiceTest {
@Test
public void testLocalCancelledBySenderBeforeSend() {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
+ StatMap<MailboxSendOperator.StatKey> stats = new StatMap<>(MailboxSendOperator.StatKey.class);
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
receivingMailbox.registeredReader(numCallbacks::getAndIncrement);
@@ -197,7 +207,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
receivingMailbox.registeredReader(numCallbacks::getAndIncrement);
@@ -226,7 +236,7 @@ public class MailboxServiceTest {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
long deadlineMs = System.currentTimeMillis() + 1000;
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, deadlineMs);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, deadlineMs, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
receivingMailbox.registeredReader(numCallbacks::getAndIncrement);
@@ -261,7 +271,7 @@ public class MailboxServiceTest {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
_mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId,
- System.currentTimeMillis() + 1000);
+ System.currentTimeMillis() + 1000, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
receivingMailbox.registeredReader(numCallbacks::getAndIncrement);
@@ -273,7 +283,7 @@ public class MailboxServiceTest {
// Next send will throw exception because buffer is full
try {
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
fail("Except exception when sending data after buffer is full");
} catch (Exception e) {
// Expected
@@ -298,7 +308,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService1.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
receivingMailbox.registeredReader(() -> {
});
@@ -313,7 +323,7 @@ public class MailboxServiceTest {
// send another block b/c it doesn't guarantee the next block must be EOS
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{0}));
// send a metadata block
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
// sending side should early terminate
assertTrue(sendingMailbox.isEarlyTerminated());
@@ -326,11 +336,11 @@ public class MailboxServiceTest {
// Sends are non-blocking as long as channel capacity is not breached
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
for (int i = 0; i < ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS - 1; i++) {
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{i}));
}
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
sendingMailbox.complete();
// Wait until all the mails are delivered
@@ -374,11 +384,11 @@ public class MailboxServiceTest {
// Sends are non-blocking as long as channel capacity is not breached
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
for (int i = 0; i < ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS - 1; i++) {
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{i}));
}
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
sendingMailbox.complete();
// Wait until all the mails are delivered
@@ -406,7 +416,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
CountDownLatch receiveMailLatch = new CountDownLatch(2);
@@ -441,7 +451,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
CountDownLatch receiveMailLatch = new CountDownLatch(1);
@@ -475,7 +485,7 @@ public class MailboxServiceTest {
throws Exception {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
CountDownLatch receiveMailLatch = new CountDownLatch(1);
@@ -509,7 +519,7 @@ public class MailboxServiceTest {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
long deadlineMs = System.currentTimeMillis() + 1000;
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, deadlineMs);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, deadlineMs, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
CountDownLatch receiveMailLatch = new CountDownLatch(2);
@@ -552,7 +562,7 @@ public class MailboxServiceTest {
String mailboxId = MailboxIdUtils.toMailboxId(_requestId++, SENDER_STAGE_ID, 0, RECEIVER_STAGE_ID, 0);
SendingMailbox sendingMailbox =
_mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId,
- System.currentTimeMillis() + 1000);
+ System.currentTimeMillis() + 1000, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
AtomicInteger numCallbacks = new AtomicInteger();
CountDownLatch receiveMailLatch = new CountDownLatch(ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS + 1);
@@ -567,7 +577,7 @@ public class MailboxServiceTest {
}
// Next send will be blocked on the receiver side and cause exception after timeout
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
receiveMailLatch.await();
assertEquals(numCallbacks.get(), ReceivingMailbox.DEFAULT_MAX_PENDING_BLOCKS + 1);
@@ -591,7 +601,7 @@ public class MailboxServiceTest {
// Sends are non-blocking as long as channel capacity is not breached
SendingMailbox sendingMailbox =
- _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE);
+ _mailboxService2.getSendingMailbox("localhost", _mailboxService1.getPort(), mailboxId, Long.MAX_VALUE, _stats);
ReceivingMailbox receivingMailbox = _mailboxService1.getReceivingMailbox(mailboxId);
receivingMailbox.registeredReader(() -> {
});
@@ -608,7 +618,7 @@ public class MailboxServiceTest {
// send another block b/c it doesn't guarantee the next block must be EOS
sendingMailbox.send(OperatorTestUtil.block(DATA_SCHEMA, new Object[]{0}));
// send a metadata block
- sendingMailbox.send(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ sendingMailbox.send(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(SENDER_STAGE_ID));
sendingMailbox.complete();
// sending side should early terminate
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTest.java
new file mode 100644
index 0000000000..8ede37c65b
--- /dev/null
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTest.java
@@ -0,0 +1,44 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime.blocks;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStatsTest;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+public class TransferableBlockTest {
+
+ @Test(dataProviderClass = MultiStageQueryStatsTest.class, dataProvider = "stats")
+ public void serializeDeserialize(MultiStageQueryStats queryStats) {
+ TransferableBlock transferableBlock = new TransferableBlock(queryStats);
+ List<ByteBuffer> fromStatsBytes = transferableBlock.getSerializedStatsByStage();
+
+ DataBlock dataBlock = transferableBlock.getDataBlock();
+
+ TransferableBlock fromBlock = TransferableBlockUtils.wrap(dataBlock);
+ List<ByteBuffer> fromBlockBytes = fromBlock.getSerializedStatsByStage();
+
+ Assert.assertEquals(fromStatsBytes, fromBlockBytes, "Serialized bytes from stats and block should be equal");
+ }
+}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTestUtils.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTestUtils.java
new file mode 100644
index 0000000000..a300326fb2
--- /dev/null
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockTestUtils.java
@@ -0,0 +1,45 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime.blocks;
+
+import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.testng.Assert;
+
+
+public class TransferableBlockTestUtils {
+ private TransferableBlockTestUtils() {
+ // do not instantiate.
+ }
+
+ public static TransferableBlock getEndOfStreamTransferableBlock(int stageId) {
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(MultiStageQueryStats.emptyStats(stageId));
+ }
+
+ public static void assertSuccessEos(TransferableBlock block) {
+ Assert.assertEquals(block.getType(), DataBlock.Type.METADATA, "Block type should be metadata");
+ Assert.assertTrue(block.isSuccessfulEndOfStreamBlock(), "Block should be successful EOS");
+ }
+
+ public static void assertDataBlock(TransferableBlock block) {
+ if (block.getType() != DataBlock.Type.ROW && block.getType() != DataBlock.Type.COLUMNAR) {
+ Assert.fail("Block type should be row or columnar but found " + block.getType());
+ }
+ }
+}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtilsTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtilsTest.java
index 94900dd334..d136f91795 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtilsTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/blocks/TransferableBlockUtilsTest.java
@@ -89,7 +89,7 @@ public class TransferableBlockUtilsTest {
validateNonSplittableBlock(columnarBlock);
// METADATA
- MetadataBlock metadataBlock = new MetadataBlock(MetadataBlock.MetadataBlockType.EOS);
+ MetadataBlock metadataBlock = MetadataBlock.newEos();
validateNonSplittableBlock(metadataBlock);
}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
index 7a35202529..aff6a68853 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
@@ -28,6 +28,7 @@ import org.apache.pinot.common.utils.NamedThreadFactory;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.routing.StageMetadata;
import org.apache.pinot.query.routing.WorkerMetadata;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.MultiStageOperator;
import org.apache.pinot.query.runtime.operator.OpChain;
@@ -91,7 +92,7 @@ public class OpChainSchedulerServiceTest {
CountDownLatch latch = new CountDownLatch(1);
Mockito.when(_operatorA.nextBlock()).thenAnswer(inv -> {
latch.countDown();
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0);
});
schedulerService.register(opChain);
@@ -108,7 +109,7 @@ public class OpChainSchedulerServiceTest {
CountDownLatch latch = new CountDownLatch(1);
Mockito.when(_operatorA.nextBlock()).thenAnswer(inv -> {
latch.countDown();
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0);
});
schedulerService.register(opChain);
@@ -123,7 +124,7 @@ public class OpChainSchedulerServiceTest {
OpChainSchedulerService schedulerService = new OpChainSchedulerService(_executor);
CountDownLatch latch = new CountDownLatch(1);
- Mockito.when(_operatorA.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ Mockito.when(_operatorA.nextBlock()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.doAnswer(inv -> {
latch.countDown();
return null;
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
index c207118aaa..790bd3a901 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
@@ -18,7 +18,6 @@
*/
package org.apache.pinot.query.runtime.operator;
-import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.Collections;
@@ -27,7 +26,7 @@ import java.util.Map;
import org.apache.calcite.rel.hint.RelHint;
import org.apache.calcite.sql.SqlKind;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
-import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.logical.RexExpression;
@@ -35,6 +34,7 @@ import org.apache.pinot.query.planner.plannode.AbstractPlanNode;
import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
import org.apache.pinot.query.routing.VirtualServerAddress;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.mockito.Mock;
@@ -90,7 +90,7 @@ public class AggregateOperatorTest {
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT, DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
Collections.singletonList(-1), null);
// When:
@@ -107,11 +107,11 @@ public class AggregateOperatorTest {
List<RexExpression> calls = ImmutableList.of(getSum(new RexExpression.InputRef(1)));
List<RexExpression> group = ImmutableList.of(new RexExpression.InputRef(0));
- Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT, DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
Collections.singletonList(-1), null);
// When:
@@ -130,11 +130,11 @@ public class AggregateOperatorTest {
DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, DOUBLE});
Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1.0}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT, DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
Collections.singletonList(-1), null);
// When:
@@ -158,11 +158,11 @@ public class AggregateOperatorTest {
Mockito.when(_input.nextBlock())
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1.0}, new Object[]{2, 2.0}))
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 3.0}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT, DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
Collections.singletonList(-1), null);
// When:
@@ -189,12 +189,12 @@ public class AggregateOperatorTest {
Mockito.when(_input.nextBlock())
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1.0, 0}, new Object[]{2, 2.0, 1}))
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 3.0, 1}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema =
new DataSchema(new String[]{"group", "sum", "sumWithFilter"}, new ColumnDataType[]{INT, DOUBLE, DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
filterArgIds, null);
// When:
@@ -215,7 +215,7 @@ public class AggregateOperatorTest {
RexExpression.FunctionCall agg = getSum(new RexExpression.InputRef(0));
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{STRING, DOUBLE});
AggregateOperator sum0GroupBy1 =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), upstreamOperator, outSchema,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), upstreamOperator, outSchema,
Collections.singletonList(agg), Collections.singletonList(new RexExpression.InputRef(1)), AggType.DIRECT,
Collections.singletonList(-1), null);
TransferableBlock result = sum0GroupBy1.getNextBlock();
@@ -239,7 +239,7 @@ public class AggregateOperatorTest {
DataSchema outSchema = new DataSchema(new String[]{"unknown"}, new ColumnDataType[]{DOUBLE});
// When:
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, AggType.DIRECT,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group, AggType.DIRECT,
Collections.singletonList(-1), null);
}
@@ -254,11 +254,11 @@ public class AggregateOperatorTest {
// TODO: it is necessary to produce two values here, the operator only throws on second
// (see the comment in Aggregate operator)
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"}, new Object[]{2, "foo"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE});
AggregateOperator operator =
- new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group,
+ new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, outSchema, calls, group,
AggType.INTERMEDIATE, Collections.singletonList(-1), null);
// When:
@@ -280,10 +280,10 @@ public class AggregateOperatorTest {
Mockito.when(_input.nextBlock())
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1.0}, new Object[]{3, 2.0}))
.thenReturn(OperatorTestUtil.block(inSchema, new Object[]{3, 3.0}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema outSchema = new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT, DOUBLE});
- OpChainExecutionContext context = OperatorTestUtil.getDefaultContext();
+ OpChainExecutionContext context = OperatorTestUtil.getTracingContext();
Map<String, String> hintsMap = ImmutableMap.of(PinotHintOptions.AggregateOptions.NUM_GROUPS_LIMIT, "1");
AggregateOperator operator =
new AggregateOperator(context, _input, outSchema, calls, group, AggType.DIRECT, Collections.singletonList(-1),
@@ -297,13 +297,12 @@ public class AggregateOperatorTest {
Mockito.verify(_input).earlyTerminate();
// Then:
- Assert.assertTrue(block1.getNumRows() == 1, "when group limit reach it should only return that many groups");
+ Assert.assertEquals(block1.getNumRows(), 1, "when group limit reach it should only return that many groups");
Assert.assertTrue(block2.isEndOfStreamBlock(), "Second block is EOS (done processing)");
- String operatorId =
- Joiner.on("_").join(AggregateOperator.class.getSimpleName(), context.getStageId(), context.getServer());
- OperatorStats operatorStats = context.getStats().getOperatorStats(context, operatorId);
- Assert.assertEquals(operatorStats.getExecutionStats().get(DataTable.MetadataKey.NUM_GROUPS_LIMIT_REACHED.getName()),
- "true");
+ StatMap<AggregateOperator.StatKey> aggrStats =
+ OperatorTestUtil.getStatMap(AggregateOperator.StatKey.class, block2);
+ Assert.assertTrue(aggrStats.getBoolean(AggregateOperator.StatKey.NUM_GROUPS_LIMIT_REACHED),
+ "num groups limit should be reached");
}
private static RexExpression.FunctionCall getSum(RexExpression arg) {
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/FilterOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/FilterOperatorTest.java
index f68fbdea99..27825f916a 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/FilterOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/FilterOperatorTest.java
@@ -26,6 +26,7 @@ import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.mockito.Mock;
import org.mockito.Mockito;
@@ -61,7 +62,7 @@ public class FilterOperatorTest {
ColumnDataType.BOOLEAN
});
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, booleanLiteral);
TransferableBlock errorBlock = op.getNextBlock();
Assert.assertTrue(errorBlock.isErrorBlock());
Assert.assertTrue(errorBlock.getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains("filterError"));
@@ -73,9 +74,10 @@ public class FilterOperatorTest {
DataSchema inputSchema = new DataSchema(new String[]{"intCol"}, new ColumnDataType[]{
ColumnDataType.INT
});
- Mockito.when(_upstreamOperator.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ Mockito.when(_upstreamOperator.nextBlock())
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, booleanLiteral);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertTrue(dataBlock.isEndOfStreamBlock());
}
@@ -88,9 +90,9 @@ public class FilterOperatorTest {
});
Mockito.when(_upstreamOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{0}, new Object[]{1}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, booleanLiteral);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -108,7 +110,7 @@ public class FilterOperatorTest {
Mockito.when(_upstreamOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2}));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, booleanLiteral);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -125,7 +127,7 @@ public class FilterOperatorTest {
Mockito.when(_upstreamOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2}));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, booleanLiteral);
op.getNextBlock();
}
@@ -138,7 +140,7 @@ public class FilterOperatorTest {
});
Mockito.when(_upstreamOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2}));
- FilterOperator op = new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, ref0);
+ FilterOperator op = new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, ref0);
op.getNextBlock();
}
@@ -150,7 +152,7 @@ public class FilterOperatorTest {
});
Mockito.when(_upstreamOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1, 1}, new Object[]{2, 0}));
- FilterOperator op = new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, ref1);
+ FilterOperator op = new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, ref1);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -170,7 +172,7 @@ public class FilterOperatorTest {
ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1)));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, andCall);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, andCall);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -190,7 +192,7 @@ public class FilterOperatorTest {
ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1)));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, orCall);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, orCall);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -212,7 +214,7 @@ public class FilterOperatorTest {
ImmutableList.of(new RexExpression.InputRef(0)));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, notCall);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, notCall);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -232,7 +234,7 @@ public class FilterOperatorTest {
new RexExpression.FunctionCall(SqlKind.GREATER_THAN, ColumnDataType.BOOLEAN, "greaterThan",
ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1)));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, greaterThan);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, greaterThan);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -252,7 +254,7 @@ public class FilterOperatorTest {
new RexExpression.FunctionCall(SqlKind.OTHER, ColumnDataType.BOOLEAN, "startsWith",
ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(ColumnDataType.STRING, "star")));
FilterOperator op =
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, startsWith);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, startsWith);
TransferableBlock dataBlock = op.getNextBlock();
Assert.assertFalse(dataBlock.isErrorBlock());
List<Object[]> result = dataBlock.getContainer();
@@ -272,6 +274,6 @@ public class FilterOperatorTest {
RexExpression.FunctionCall startsWith =
new RexExpression.FunctionCall(SqlKind.OTHER, ColumnDataType.BOOLEAN, "startsWithError",
ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(ColumnDataType.STRING, "star")));
- new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, startsWith);
+ new FilterOperator(OperatorTestUtil.getTracingContext(), _upstreamOperator, inputSchema, startsWith);
}
}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java
index 48dbe8ef3e..45afa6dbc6 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java
@@ -18,7 +18,6 @@
*/
package org.apache.pinot.query.runtime.operator;
-import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.ArrayList;
@@ -30,7 +29,7 @@ import org.apache.calcite.rel.core.JoinRelType;
import org.apache.calcite.rel.hint.RelHint;
import org.apache.calcite.sql.SqlKind;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
-import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
@@ -38,6 +37,7 @@ import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.plannode.JoinNode;
import org.apache.pinot.query.routing.VirtualServerAddress;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.mockito.Mock;
@@ -94,10 +94,10 @@ public class HashJoinOperatorTest {
List<RexExpression> joinClauses = new ArrayList<>();
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema resultSchema =
new DataSchema(new String[]{"int_col1", "string_col1", "int_col2", "string_col2"}, new ColumnDataType[]{
ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.STRING
@@ -105,7 +105,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses, Collections.emptyList());
HashJoinOperator joinOnString =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = joinOnString.nextBlock();
List<Object[]> resultRows = result.getContainer();
@@ -127,10 +127,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
new DataSchema(new String[]{"int_col1", "string_col1", "int_col2", "string_co2"}, new ColumnDataType[]{
@@ -139,7 +139,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator joinOnInt =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = joinOnInt.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows = Arrays.asList(new Object[]{2, "BB", 2, "Aa"}, new Object[]{2, "BB", 2, "BB"});
@@ -158,10 +158,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
new DataSchema(new String[]{"int_col1", "string_col1", "int_col2", "string_co2"}, new ColumnDataType[]{
@@ -170,7 +170,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses, Collections.emptyList());
HashJoinOperator joinOnInt =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = joinOnInt.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows =
@@ -195,10 +195,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
@@ -208,7 +208,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT,
getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
@@ -226,10 +226,10 @@ public class HashJoinOperatorTest {
DataSchema rightSchema = new DataSchema(new String[]{"int_col", "string_col"}, new ColumnDataType[]{
ColumnDataType.INT, ColumnDataType.STRING
});
- Mockito.when(_leftOperator.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ Mockito.when(_leftOperator.nextBlock()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{1, "BB"}, new Object[]{1, "CC"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
DataSchema resultSchema =
new DataSchema(new String[]{"int_col1", "string_col1", "int_co2", "string_col2"}, new ColumnDataType[]{
@@ -239,7 +239,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
Assert.assertTrue(result.isEndOfStreamBlock());
@@ -254,10 +254,10 @@ public class HashJoinOperatorTest {
ColumnDataType.INT, ColumnDataType.STRING
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{1, "BB"}, new Object[]{1, "CC"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
@@ -267,7 +267,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
@@ -287,8 +287,8 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{1, "BB"}, new Object[]{1, "CC"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
- Mockito.when(_rightOperator.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
+ Mockito.when(_rightOperator.nextBlock()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
@@ -299,7 +299,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
@@ -316,10 +316,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
List<RexExpression> functionOperands = new ArrayList<>();
@@ -333,7 +333,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows =
@@ -354,10 +354,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{1, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
List<RexExpression> functionOperands = new ArrayList<>();
@@ -371,7 +371,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows = Arrays.asList(new Object[]{1, "Aa", 2, "Aa"}, new Object[]{2, "BB", 1, "BB"});
@@ -391,10 +391,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema = new DataSchema(new String[]{"foo", "bar", "foo", "bar"}, new ColumnDataType[]{
@@ -403,7 +403,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.RIGHT,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator joinOnNum =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = joinOnNum.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows = Arrays.asList(new Object[]{2, "BB", 2, "Aa"}, new Object[]{2, "BB", 2, "BB"});
@@ -418,7 +418,7 @@ public class HashJoinOperatorTest {
Assert.assertEquals(resultRows.get(0), expectedRows.get(0));
// Third block is EOS block.
result = joinOnNum.nextBlock();
- Assert.assertTrue(result.isSuccessfulEndOfStreamBlock());
+ TransferableBlockTestUtils.assertSuccessEos(result);
}
@Test
@@ -431,10 +431,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}, new Object[]{4, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema = new DataSchema(new String[]{"foo", "bar", "foo", "bar"}, new ColumnDataType[]{
@@ -443,7 +443,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.SEMI,
getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows =
@@ -465,10 +465,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}, new Object[]{4, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema = new DataSchema(new String[]{"foo", "bar", "foo", "bar"}, new ColumnDataType[]{
ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.INT, ColumnDataType.STRING
@@ -476,7 +476,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.FULL,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows = ImmutableList.of(new Object[]{1, "Aa", null, null}, new Object[]{2, "BB", 2, "Aa"},
@@ -490,11 +490,12 @@ public class HashJoinOperatorTest {
result = join.nextBlock();
resultRows = result.getContainer();
expectedRows = ImmutableList.of(new Object[]{null, null, 3, "BB"});
+ TransferableBlockTestUtils.assertDataBlock(result);
Assert.assertEquals(resultRows.size(), expectedRows.size());
Assert.assertEquals(resultRows.get(0), expectedRows.get(0));
// Third block is EOS block.
result = join.nextBlock();
- Assert.assertTrue(result.isSuccessfulEndOfStreamBlock());
+ TransferableBlockTestUtils.assertSuccessEos(result);
}
@Test
@@ -507,10 +508,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}, new Object[]{4, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema = new DataSchema(new String[]{"foo", "bar", "foo", "bar"}, new ColumnDataType[]{
@@ -519,7 +520,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.ANTI,
getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
List<Object[]> resultRows = result.getContainer();
List<Object[]> expectedRows = ImmutableList.of(new Object[]{4, "CC", null, null});
@@ -539,7 +540,7 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{1, "BB"}, new Object[]{1, "CC"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock())
.thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("testInnerJoinRightError")));
@@ -551,7 +552,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
Assert.assertTrue(result.isErrorBlock());
@@ -569,7 +570,7 @@ public class HashJoinOperatorTest {
});
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{1, "BB"}, new Object[]{1, "CC"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_leftOperator.nextBlock())
.thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("testInnerJoinLeftError")));
@@ -581,7 +582,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, Collections.emptyList());
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
Assert.assertTrue(result.isErrorBlock());
@@ -598,10 +599,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
@@ -613,7 +614,7 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, getJoinHints(hintsMap));
HashJoinOperator join =
- new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node);
+ new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftOperator, _rightOperator, leftSchema, node);
TransferableBlock result = join.nextBlock();
Assert.assertTrue(result.isErrorBlock());
@@ -631,10 +632,10 @@ public class HashJoinOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(leftSchema, new Object[]{1, "Aa"}, new Object[]{2, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(rightSchema, new Object[]{2, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
List<RexExpression> joinClauses = new ArrayList<>();
DataSchema resultSchema =
@@ -646,19 +647,19 @@ public class HashJoinOperatorTest {
JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER,
getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses, getJoinHints(hintsMap));
- OpChainExecutionContext context = OperatorTestUtil.getDefaultContext();
+ OpChainExecutionContext context = OperatorTestUtil.getTracingContext();
HashJoinOperator join = new HashJoinOperator(context, _leftOperator, _rightOperator, leftSchema, node);
- TransferableBlock result = join.nextBlock();
+ TransferableBlock firstBlock = join.nextBlock();
Mockito.verify(_rightOperator).earlyTerminate();
- Assert.assertFalse(result.isErrorBlock());
- Assert.assertEquals(result.getNumRows(), 1);
-
- String operatorId =
- Joiner.on("_").join(HashJoinOperator.class.getSimpleName(), context.getStageId(), context.getServer());
- OperatorStats operatorStats = context.getStats().getOperatorStats(context, operatorId);
- Assert.assertEquals(
- operatorStats.getExecutionStats().get(DataTable.MetadataKey.MAX_ROWS_IN_JOIN_REACHED.getName()), "true");
+ Assert.assertTrue(firstBlock.isDataBlock(), "First block should be a data block but is " + firstBlock.getClass());
+ Assert.assertEquals(firstBlock.getNumRows(), 1);
+
+ TransferableBlock secondBlock = join.nextBlock();
+ StatMap<HashJoinOperator.StatKey> joinStats =
+ OperatorTestUtil.getStatMap(HashJoinOperator.StatKey.class, secondBlock);
+ Assert.assertTrue(joinStats.getBoolean(HashJoinOperator.StatKey.MAX_ROWS_IN_JOIN_REACHED),
+ "Max rows in join should be reached");
}
}
// TODO: Add more inequi join tests.
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/IntersectOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/IntersectOperatorTest.java
index 04c14caea1..51cc068f90 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/IntersectOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/IntersectOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.routing.VirtualServerAddress;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
@@ -66,13 +66,13 @@ public class IntersectOperatorTest {
});
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{4, "DD"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
IntersectOperator intersectOperator =
- new IntersectOperator(OperatorTestUtil.getDefaultContext(), ImmutableList.of(_leftOperator, _rightOperator),
+ new IntersectOperator(OperatorTestUtil.getTracingContext(), ImmutableList.of(_leftOperator, _rightOperator),
schema);
TransferableBlock result = intersectOperator.nextBlock();
@@ -95,14 +95,14 @@ public class IntersectOperatorTest {
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"},
new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{4, "DD"},
new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{4, "DD"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
IntersectOperator intersectOperator =
- new IntersectOperator(OperatorTestUtil.getDefaultContext(), ImmutableList.of(_leftOperator, _rightOperator),
+ new IntersectOperator(OperatorTestUtil.getTracingContext(), ImmutableList.of(_leftOperator, _rightOperator),
schema);
TransferableBlock result = intersectOperator.nextBlock();
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
index b4365b8f3d..2bca79dcf8 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafStageTransferableBlockOperatorTest.java
@@ -109,7 +109,7 @@ public class LeafStageTransferableBlockOperatorTest {
InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(1), schema,
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema,
queryExecutor, _executorService);
_operatorRef.set(operator);
@@ -140,7 +140,7 @@ public class LeafStageTransferableBlockOperatorTest {
InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(1),
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1),
desiredSchema, queryExecutor, _executorService);
_operatorRef.set(operator);
@@ -167,7 +167,7 @@ public class LeafStageTransferableBlockOperatorTest {
InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(1), schema,
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema,
queryExecutor, _executorService);
_operatorRef.set(operator);
@@ -198,7 +198,7 @@ public class LeafStageTransferableBlockOperatorTest {
InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(2), schema,
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(2), schema,
queryExecutor, _executorService);
_operatorRef.set(operator);
@@ -224,7 +224,7 @@ public class LeafStageTransferableBlockOperatorTest {
errorBlock.addException(QueryException.QUERY_EXECUTION_ERROR.getErrorCode(), "foobar");
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, errorBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(1), schema,
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema,
queryExecutor, _executorService);
_operatorRef.set(operator);
@@ -252,7 +252,7 @@ public class LeafStageTransferableBlockOperatorTest {
new InstanceResponseBlock(new SelectionResultsBlock(resultSchema, Collections.emptyList(), queryContext));
QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, emptySelectionResponseBlock);
LeafStageTransferableBlockOperator operator =
- new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), mockQueryRequests(1),
+ new LeafStageTransferableBlockOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1),
desiredSchema, queryExecutor, _executorService);
_operatorRef.set(operator);
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LiteralValueOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LiteralValueOperatorTest.java
index 959038c677..bce3ba4e47 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LiteralValueOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LiteralValueOperatorTest.java
@@ -63,7 +63,7 @@ public class LiteralValueOperatorTest {
new RexExpression.Literal(ColumnDataType.INT, 1)),
ImmutableList.of(new RexExpression.Literal(ColumnDataType.STRING, ""),
new RexExpression.Literal(ColumnDataType.INT, 2)));
- LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getDefaultContext(), schema, literals);
+ LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getTracingContext(), schema, literals);
// When:
TransferableBlock transferableBlock = operator.nextBlock();
@@ -79,7 +79,7 @@ public class LiteralValueOperatorTest {
// Given:
DataSchema schema = new DataSchema(new String[]{}, new ColumnDataType[]{});
List<List<RexExpression>> literals = ImmutableList.of(ImmutableList.of());
- LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getDefaultContext(), schema, literals);
+ LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getTracingContext(), schema, literals);
// When:
TransferableBlock transferableBlock = operator.nextBlock();
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java
index 2868f6526c..060ef3de7e 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperatorTest.java
@@ -35,6 +35,7 @@ import org.apache.pinot.query.routing.SharedMailboxInfos;
import org.apache.pinot.query.routing.StageMetadata;
import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.mockito.Mock;
@@ -120,7 +121,7 @@ public class MailboxReceiveOperatorTest {
@Test
public void shouldReceiveEosDirectlyFromSender() {
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
- when(_mailbox1.poll()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ when(_mailbox1.poll()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadata1);
@@ -134,7 +135,7 @@ public class MailboxReceiveOperatorTest {
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
Object[] row = new Object[]{1, 1};
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadata1);
@@ -165,11 +166,11 @@ public class MailboxReceiveOperatorTest {
@Test
public void shouldReceiveMailboxFromTwoServersOneNull() {
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
- when(_mailbox1.poll()).thenReturn(null, TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ when(_mailbox1.poll()).thenReturn(null, TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2);
Object[] row = new Object[]{1, 1};
when(_mailbox2.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadataBoth);
@@ -189,10 +190,10 @@ public class MailboxReceiveOperatorTest {
Object[] row3 = new Object[]{3, 3};
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row1),
- OperatorTestUtil.block(DATA_SCHEMA, row3), TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ OperatorTestUtil.block(DATA_SCHEMA, row3), TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2);
when(_mailbox2.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row2),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadataBoth);
@@ -217,7 +218,7 @@ public class MailboxReceiveOperatorTest {
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2);
Object[] row = new Object[]{3, 3};
when(_mailbox2.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadataBoth);
@@ -236,10 +237,10 @@ public class MailboxReceiveOperatorTest {
Object[] row3 = new Object[]{3, 3};
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1);
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row1),
- OperatorTestUtil.block(DATA_SCHEMA, row3), TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ OperatorTestUtil.block(DATA_SCHEMA, row3), TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_2))).thenReturn(_mailbox2);
when(_mailbox2.poll()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row2),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
OpChainExecutionContext context =
OperatorTestUtil.getOpChainContext(_mailboxService, Long.MAX_VALUE, _stageMetadataBoth);
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxSendOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxSendOperatorTest.java
index 92f7b2ce7a..e28d1dc477 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxSendOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MailboxSendOperatorTest.java
@@ -21,7 +21,6 @@ package org.apache.pinot.query.runtime.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.List;
-import java.util.Map;
import java.util.concurrent.TimeoutException;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.mailbox.MailboxService;
@@ -30,6 +29,7 @@ import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.exchange.BlockExchange;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
@@ -41,7 +41,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;
import static org.mockito.MockitoAnnotations.openMocks;
-import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
@@ -124,7 +124,8 @@ public class MailboxSendOperatorTest {
public void shouldSendEosBlock()
throws Exception {
// Given:
- TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock(
+ MultiStageQueryStats.emptyStats(SENDER_STAGE_ID));
when(_sourceOperator.nextBlock()).thenReturn(eosBlock);
// When:
@@ -145,7 +146,8 @@ public class MailboxSendOperatorTest {
OperatorTestUtil.block(new DataSchema(new String[]{}, new DataSchema.ColumnDataType[]{}));
TransferableBlock dataBlock2 =
OperatorTestUtil.block(new DataSchema(new String[]{}, new DataSchema.ColumnDataType[]{}));
- TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock(
+ MultiStageQueryStats.emptyStats(SENDER_STAGE_ID));
when(_sourceOperator.nextBlock()).thenReturn(dataBlock1, dataBlock2, eosBlock);
// When:
@@ -172,9 +174,10 @@ public class MailboxSendOperatorTest {
assertTrue(blocks.get(2).isSuccessfulEndOfStreamBlock(), "expected to send EOS block to exchange on third call");
// EOS block should contain statistics
- Map<String, OperatorStats> resultMetadata = blocks.get(2).getResultMetadata();
- assertEquals(resultMetadata.size(), 1);
- assertTrue(resultMetadata.containsKey(mailboxSendOperator.getOperatorId()));
+ MultiStageQueryStats resultMetadata = blocks.get(2).getQueryStats();
+ MultiStageQueryStats.StageStats stageStats = resultMetadata.getCurrentStats();
+ assertNotNull(stageStats, "expected to have stats for sender stage");
+ assertNotNull(stageStats.getOperatorStats(0));
}
@Test
@@ -200,6 +203,6 @@ public class MailboxSendOperatorTest {
OpChainExecutionContext context =
new OpChainExecutionContext(_mailboxService, 123L, Long.MAX_VALUE, ImmutableMap.of(), stageMetadata,
workerMetadata, null);
- return new MailboxSendOperator(context, _sourceOperator, _exchange, null, null, false);
+ return new MailboxSendOperator(context, _sourceOperator, statMap -> _exchange, null, null, false);
}
}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MinusOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MinusOperatorTest.java
index 53e81612ad..81f4b13699 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MinusOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MinusOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.routing.VirtualServerAddress;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
@@ -67,13 +67,13 @@ public class MinusOperatorTest {
Mockito.when(_leftOperator.nextBlock())
.thenReturn(OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"},
new Object[]{4, "DD"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{5, "EE"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
MinusOperator minusOperator =
- new MinusOperator(OperatorTestUtil.getDefaultContext(), ImmutableList.of(_leftOperator, _rightOperator),
+ new MinusOperator(OperatorTestUtil.getTracingContext(), ImmutableList.of(_leftOperator, _rightOperator),
schema);
TransferableBlock result = minusOperator.nextBlock();
@@ -97,14 +97,14 @@ public class MinusOperatorTest {
.thenReturn(OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"},
new Object[]{4, "DD"}, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{3, "CC"},
new Object[]{4, "DD"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
Mockito.when(_rightOperator.nextBlock()).thenReturn(
OperatorTestUtil.block(schema, new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{5, "EE"},
new Object[]{1, "AA"}, new Object[]{2, "BB"}, new Object[]{5, "EE"}))
- .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ .thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
MinusOperator minusOperator =
- new MinusOperator(OperatorTestUtil.getDefaultContext(), ImmutableList.of(_leftOperator, _rightOperator),
+ new MinusOperator(OperatorTestUtil.getTracingContext(), ImmutableList.of(_leftOperator, _rightOperator),
schema);
TransferableBlock result = minusOperator.nextBlock();
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
index f4cd38a42e..beb78740d8 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
@@ -22,15 +22,14 @@ import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.Collections;
import java.util.List;
-import java.util.Map;
import java.util.Stack;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
-import org.apache.calcite.rel.RelDistribution;
-import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
@@ -50,12 +49,15 @@ import org.apache.pinot.query.routing.SharedMailboxInfos;
import org.apache.pinot.query.routing.StageMetadata;
import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
+import org.apache.pinot.query.runtime.blocks.TransferableBlockTestUtils;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.exchange.BlockExchange;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
-import org.apache.pinot.spi.utils.CommonConstants;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
@@ -65,9 +67,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertTrue;
-import static org.testng.Assert.fail;
+import static org.testng.Assert.*;
public class OpChainTest {
@@ -111,7 +111,7 @@ public class OpChainTest {
}).when(_exchange).send(any(TransferableBlock.class));
when(_mailbox2.poll()).then(x -> {
if (_blockList.isEmpty()) {
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0);
}
return _blockList.remove(0);
});
@@ -132,134 +132,42 @@ public class OpChainTest {
_executor.shutdown();
}
- @Test
- public void testExecutionTimerStats() {
- when(_sourceOperator.nextBlock()).then(x -> {
- Thread.sleep(100);
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
- });
- OpChain opChain = new OpChain(OperatorTestUtil.getDefaultContext(), _sourceOperator);
- opChain.getStats().executing();
- opChain.getRoot().nextBlock();
- opChain.getStats().queued();
- assertTrue(opChain.getStats().getExecutionTime() >= 100);
-
- when(_sourceOperator.nextBlock()).then(x -> {
- Thread.sleep(20);
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
- });
- opChain = new OpChain(OperatorTestUtil.getDefaultContext(), _sourceOperator);
- opChain.getStats().executing();
- opChain.getRoot().nextBlock();
- opChain.getStats().queued();
- assertTrue(opChain.getStats().getExecutionTime() >= 20);
- assertTrue(opChain.getStats().getExecutionTime() < 100);
- }
-
@Test
public void testStatsCollectionTracingEnabled() {
- OpChainExecutionContext context = OperatorTestUtil.getDefaultContext();
+ OpChainExecutionContext context = OperatorTestUtil.getTracingContext();
DummyMultiStageOperator dummyMultiStageOperator = new DummyMultiStageOperator(context);
OpChain opChain = new OpChain(context, dummyMultiStageOperator);
- opChain.getStats().executing();
- opChain.getRoot().nextBlock();
- opChain.getStats().queued();
-
- assertTrue(opChain.getStats().getExecutionTime() >= 1000);
- assertEquals(opChain.getStats().getOperatorStatsMap().size(), 1);
- assertTrue(opChain.getStats().getOperatorStatsMap().containsKey(dummyMultiStageOperator.getOperatorId()));
+ TransferableBlock eosBlock = drainOpChain(opChain);
- Map<String, String> executionStats =
- opChain.getStats().getOperatorStatsMap().get(dummyMultiStageOperator.getOperatorId()).getExecutionStats();
+ assertTrue(eosBlock.isSuccessfulEndOfStreamBlock(), "Expected end of stream block to be successful");
+ MultiStageQueryStats queryStats = eosBlock.getQueryStats();
+ assertNotNull(queryStats, "Expected query stats to be non-null");
- long time = Long.parseLong(executionStats.get(DataTable.MetadataKey.OPERATOR_EXECUTION_TIME_MS.getName()));
- assertTrue(time >= 1000 && time <= 2000,
- "Expected " + DataTable.MetadataKey.OPERATOR_EXECUTION_TIME_MS + " to be in [1000, 2000] but found " + time);
+ @SuppressWarnings("unchecked")
+ StatMap<LiteralValueOperator.StatKey> lastOperatorStats =
+ (StatMap<LiteralValueOperator.StatKey>) queryStats.getCurrentStats().getLastOperatorStats();
+ assertNotEquals(lastOperatorStats.getLong(LiteralValueOperator.StatKey.EXECUTION_TIME_MS), 0L,
+ "Expected execution time to be non-zero");
}
@Test
public void testStatsCollectionTracingDisabled() {
- OpChainExecutionContext context = OperatorTestUtil.getDefaultContextWithTracingDisabled();
+ OpChainExecutionContext context = OperatorTestUtil.getNoTracingContext();
DummyMultiStageOperator dummyMultiStageOperator = new DummyMultiStageOperator(context);
OpChain opChain = new OpChain(context, dummyMultiStageOperator);
- opChain.getStats().executing();
- opChain.getRoot().nextBlock();
- opChain.getStats().queued();
-
- assertTrue(opChain.getStats().getExecutionTime() >= 1000);
- assertEquals(opChain.getStats().getOperatorStatsMap().size(), 0);
- }
-
- @Test
- public void testStatsCollectionTracingEnabledMultipleOperators() {
- long dummyOperatorWaitTime = 1000L;
-
- OpChainExecutionContext context = new OpChainExecutionContext(_mailboxService1, 123L, Long.MAX_VALUE,
- ImmutableMap.of(CommonConstants.Broker.Request.TRACE, "true"), _stageMetadata, _workerMetadata, null);
- Stack<MultiStageOperator> operators = getFullOpChain(context, dummyOperatorWaitTime);
-
- OpChain opChain = new OpChain(context, operators.peek());
- opChain.getStats().executing();
- while (!opChain.getRoot().nextBlock().isEndOfStreamBlock()) {
- // Drain the opchain
- }
- opChain.getStats().queued();
-
- OpChainExecutionContext secondStageContext = new OpChainExecutionContext(_mailboxService2, 123L, Long.MAX_VALUE,
- ImmutableMap.of(CommonConstants.Broker.Request.TRACE, "true"), _stageMetadata, _workerMetadata, null);
- MailboxReceiveOperator secondStageReceiveOp =
- new MailboxReceiveOperator(secondStageContext, RelDistribution.Type.BROADCAST_DISTRIBUTED, 1);
-
- assertTrue(opChain.getStats().getExecutionTime() >= dummyOperatorWaitTime);
- int numOperators = operators.size();
- assertEquals(opChain.getStats().getOperatorStatsMap().size(), numOperators);
- while (!operators.isEmpty()) {
- assertTrue(opChain.getStats().getOperatorStatsMap().containsKey(operators.pop().getOperatorId()));
- }
-
- while (!secondStageReceiveOp.nextBlock().isEndOfStreamBlock()) {
- // Drain the mailbox
- }
- assertEquals(secondStageContext.getStats().getOperatorStatsMap().size(), numOperators + 1);
- }
+ TransferableBlock eosBlock = drainOpChain(opChain);
- @Test
- public void testStatsCollectionTracingDisableMultipleOperators() {
- long dummyOperatorWaitTime = 1000L;
-
- OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService1, 123L, Long.MAX_VALUE, ImmutableMap.of(), _stageMetadata,
- _workerMetadata, null);
- Stack<MultiStageOperator> operators = getFullOpChain(context, dummyOperatorWaitTime);
-
- OpChain opChain = new OpChain(context, operators.peek());
- opChain.getStats().executing();
- opChain.getRoot().nextBlock();
- opChain.getStats().queued();
-
- OpChainExecutionContext secondStageContext =
- new OpChainExecutionContext(_mailboxService2, 123L, Long.MAX_VALUE, ImmutableMap.of(), _stageMetadata,
- _workerMetadata, null);
- MailboxReceiveOperator secondStageReceiveOp =
- new MailboxReceiveOperator(secondStageContext, RelDistribution.Type.BROADCAST_DISTRIBUTED, 1);
-
- assertTrue(opChain.getStats().getExecutionTime() >= dummyOperatorWaitTime);
- assertEquals(opChain.getStats().getOperatorStatsMap().size(), 2);
- assertTrue(opChain.getStats().getOperatorStatsMap().containsKey(operators.pop().getOperatorId()));
-
- while (!secondStageReceiveOp.nextBlock().isEndOfStreamBlock()) {
- // Drain the mailbox
- }
+ assertTrue(eosBlock.isSuccessfulEndOfStreamBlock(), "Expected end of stream block to be successful");
+ MultiStageQueryStats queryStats = eosBlock.getQueryStats();
+ assertNotNull(queryStats, "Expected query stats to be non-null");
- while (!operators.isEmpty()) {
- MultiStageOperator operator = operators.pop();
- if (operator.toExplainString().contains("SEND") || operator.toExplainString().contains("LEAF")) {
- assertTrue(opChain.getStats().getOperatorStatsMap().containsKey(operator.getOperatorId()));
- }
- }
- assertEquals(secondStageContext.getStats().getOperatorStatsMap().size(), 2);
+ @SuppressWarnings("unchecked")
+ StatMap<LiteralValueOperator.StatKey> lastOperatorStats =
+ (StatMap<LiteralValueOperator.StatKey>) queryStats.getCurrentStats().getLastOperatorStats();
+ assertNotEquals(lastOperatorStats.getLong(LiteralValueOperator.StatKey.EXECUTION_TIME_MS), 0L,
+ "Expected execution time to be collected");
}
private Stack<MultiStageOperator> getFullOpChain(OpChainExecutionContext context, long waitTimeInMillis) {
@@ -268,7 +176,7 @@ public class OpChainTest {
//Mailbox Receive Operator
try {
when(_mailbox1.poll()).thenReturn(OperatorTestUtil.block(upStreamSchema, new Object[]{1}),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
} catch (Exception e) {
fail("Exception while mocking mailbox receive: " + e.getMessage());
}
@@ -297,7 +205,7 @@ public class OpChainTest {
//Mailbox Send operator
MailboxSendOperator sendOperator =
- new MailboxSendOperator(context, dummyWaitOperator, _exchange, null, null, false);
+ new MailboxSendOperator(context, dummyWaitOperator, ignore -> _exchange, null, null, false);
operators.push(leafOp);
operators.push(transformOp);
@@ -319,12 +227,33 @@ public class OpChainTest {
return queryExecutor;
}
+ private TransferableBlock drainOpChain(OpChain opChain) {
+ TransferableBlock resultBlock = opChain.getRoot().nextBlock();
+ while (!resultBlock.isEndOfStreamBlock()) {
+ resultBlock = opChain.getRoot().nextBlock();
+ }
+ return resultBlock;
+ }
+
static class DummyMultiStageOperator extends MultiStageOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(DummyMultiStageOperator.class);
+ private final StatMap<LiteralValueOperator.StatKey> _statMap = new StatMap<>(LiteralValueOperator.StatKey.class);
public DummyMultiStageOperator(OpChainExecutionContext context) {
super(context);
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(LiteralValueOperator.StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(LiteralValueOperator.StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
protected TransferableBlock getNextBlock() {
try {
@@ -332,7 +261,17 @@ public class OpChainTest {
} catch (InterruptedException e) {
// IGNORE
}
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(MultiStageQueryStats.createLiteral(0, _statMap));
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.LITERAL;
+ }
+
+ @Override
+ public List<MultiStageOperator> getChildOperators() {
+ return Collections.emptyList();
}
@Nullable
@@ -343,8 +282,10 @@ public class OpChainTest {
}
static class DummyMultiStageCallableOperator extends MultiStageOperator {
+ private static final Logger LOGGER = LoggerFactory.getLogger(DummyMultiStageCallableOperator.class);
private final MultiStageOperator _upstream;
private final long _sleepTimeInMillis;
+ private final StatMap<TransformOperator.StatKey> _statMap = new StatMap<>(TransformOperator.StatKey.class);
public DummyMultiStageCallableOperator(OpChainExecutionContext context, MultiStageOperator upstream,
long sleepTimeInMillis) {
@@ -353,6 +294,17 @@ public class OpChainTest {
_sleepTimeInMillis = sleepTimeInMillis;
}
+ @Override
+ public void registerExecution(long time, int numRows) {
+ _statMap.merge(TransformOperator.StatKey.EXECUTION_TIME_MS, time);
+ _statMap.merge(TransformOperator.StatKey.EMITTED_ROWS, numRows);
+ }
+
+ @Override
+ protected Logger logger() {
+ return LOGGER;
+ }
+
@Override
public List<MultiStageOperator> getChildOperators() {
return ImmutableList.of(_upstream);
@@ -360,13 +312,26 @@ public class OpChainTest {
@Override
protected TransferableBlock getNextBlock() {
+ TransferableBlock block;
+
try {
Thread.sleep(_sleepTimeInMillis);
- _upstream.nextBlock();
+ do {
+ block = _upstream.nextBlock();
+ } while (block.isEndOfStreamBlock());
+
+ MultiStageQueryStats queryStats = block.getQueryStats();
+ assert queryStats != null;
+ queryStats.getCurrentStats().addLastOperator(getOperatorType(), _statMap);
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock(queryStats);
} catch (InterruptedException e) {
- // IGNORE
+ throw new RuntimeException(e);
}
- return TransferableBlockUtils.getEndOfStreamTransferableBlock();
+ }
+
+ @Override
+ public Type getOperatorType() {
+ return Type.TRANSFORM;
}
@Override
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java
index 5280724541..da4537cb19 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OperatorTestUtil.java
@@ -24,16 +24,19 @@ import java.util.Arrays;
import java.util.List;
import java.util.Map;
import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.routing.StageMetadata;
-import org.apache.pinot.query.routing.VirtualServerAddress;
+import org.apache.pinot.query.routing.StagePlan;
import org.apache.pinot.query.routing.WorkerMetadata;
import org.apache.pinot.query.runtime.blocks.TransferableBlock;
-import org.apache.pinot.query.runtime.operator.utils.OperatorUtils;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.apache.pinot.query.runtime.plan.server.ServerPlanRequestContext;
import org.apache.pinot.query.testutils.MockDataBlockOperatorFactory;
import org.apache.pinot.spi.utils.CommonConstants;
+import org.testng.Assert;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -52,10 +55,8 @@ public class OperatorTestUtil {
public static final String OP_1 = "op1";
public static final String OP_2 = "op2";
... 1516 lines suppressed ...
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org