From f1449ca665952869fb0d1da71db39d5fb6fb482a Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 13 Mar 2023 14:02:02 +0530 Subject: [PATCH 01/12] Refactor: Pass context object instead on individual arguments to operator --- .../runtime/operator/AggregateOperator.java | 7 ++ .../runtime/operator/FilterOperator.java | 6 ++ .../runtime/operator/HashJoinOperator.java | 7 ++ .../LeafStageTransferableBlockOperator.java | 6 ++ .../operator/LiteralValueOperator.java | 6 ++ .../operator/MailboxReceiveOperator.java | 7 ++ .../runtime/operator/MailboxSendOperator.java | 9 ++ .../pinot/query/runtime/operator/OpChain.java | 6 ++ .../query/runtime/operator/SortOperator.java | 8 ++ .../runtime/operator/TransformOperator.java | 7 ++ .../operator/WindowAggregateOperator.java | 11 +++ .../plan/OperatorExecutionContext.java | 82 +++++++++++++++++++ .../runtime/plan/PhysicalPlanVisitor.java | 39 ++++----- .../runtime/plan/PlanRequestContext.java | 4 + 14 files changed, 181 insertions(+), 24 deletions(-) create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java 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 eacd4ac49f83..7349d68d6318 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 @@ -39,6 +39,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.utils.AggregationUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.apache.pinot.segment.local.customobject.PinotFourthMoment; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -87,6 +88,12 @@ public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema requestId, stageId, virtualServerAddress); } + public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, + List groupSet, DataSchema inputSchema, OperatorExecutionContext context) { + this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.MERGERS, + context.getRequestId(), context.getStageId(), context.getServer()); + } + @VisibleForTesting AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, 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 0b819fdd4642..ee2f724d4986 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 @@ -30,6 +30,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -55,6 +56,11 @@ public class FilterOperator extends MultiStageOperator { private final DataSchema _dataSchema; private TransferableBlock _upstreamErrorBlock; + public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, + OperatorExecutionContext context) { + this(upstreamOperator, dataSchema, filter, context.getRequestId(), context.getStageId(), context.getServer()); + } + public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, long requestId, int stageId, VirtualServerAddress serverAddress) { super(requestId, stageId, serverAddress); 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 a813f35db36b..d1a02687666c 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 @@ -40,6 +40,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -88,6 +89,12 @@ public class HashJoinOperator extends MultiStageOperator { private KeySelector _leftKeySelector; private KeySelector _rightKeySelector; + public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, + DataSchema leftSchema, JoinNode node, OperatorExecutionContext context) { + this(leftTableOperator, rightTableOperator, leftSchema, node, context.getRequestId(), context.getStageId(), + context.getServer()); + } + public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, DataSchema leftSchema, JoinNode node, long requestId, int stageId, VirtualServerAddress serverAddress) { super(requestId, stageId, serverAddress); 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 d2b9c4db16af..d670834af9fc 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 @@ -41,6 +41,7 @@ import org.apache.pinot.core.query.selection.SelectionOperatorUtils; import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.TransferableBlock; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -67,6 +68,11 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator { private final DataSchema _desiredDataSchema; private int _currentIndex; + public LeafStageTransferableBlockOperator(List baseResultBlock, DataSchema dataSchema, + OperatorExecutionContext context) { + this(baseResultBlock, dataSchema, context.getRequestId(), context.getStageId(), context.getServer()); + } + public LeafStageTransferableBlockOperator(List baseResultBlock, DataSchema dataSchema, long requestId, int stageId, VirtualServerAddress serverAddress) { super(requestId, stageId, serverAddress); 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 36d3fe511fba..1fe8e8ffbb4e 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 @@ -28,6 +28,7 @@ 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.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -40,6 +41,11 @@ public class LiteralValueOperator extends MultiStageOperator { private final TransferableBlock _rexLiteralBlock; private boolean _isLiteralBlockReturned; + public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, + OperatorExecutionContext context) { + this(dataSchema, rexLiteralRows, context.getRequestId(), context.getStageId(), context.getServer()); + } + public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, long requestId, int stageId, VirtualServerAddress serverAddress) { super(requestId, stageId, serverAddress); 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 74bd60e10d9f..1cd22c84b409 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 @@ -36,6 +36,7 @@ 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.plan.OperatorExecutionContext; import org.apache.pinot.query.service.QueryConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -76,6 +77,12 @@ private static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, i receiverStageId); } + public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, + int receiverStageId, OperatorExecutionContext context) { + this(context.getMailboxService(), context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, + context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getTimeoutMs()); + } + // TODO: Move deadlineInNanoSeconds to OperatorContext. public MailboxReceiveOperator(MailboxService mailboxService, List sendingStageInstances, RelDistribution.Type exchangeType, VirtualServerAddress receiver, 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 f4e79765276e..e879c8ed7864 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 @@ -38,6 +38,7 @@ 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.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -77,6 +78,14 @@ public MailboxSendOperator(MailboxService mailboxService, jobId, senderStageId, receiverStageId, sendingServer, deadlineMs); } + public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, + KeySelector keySelector, int senderStageId, int receiverStageId, + OperatorExecutionContext context) { + this(context.getMailboxService(), dataTableBlockBaseOperator, + context.getMetadataMap().get(receiverStageId).getServerInstances(), exchangeType, keySelector, + context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getDeadlineMs()); + } + @VisibleForTesting MailboxSendOperator(MailboxService mailboxService, MultiStageOperator dataTableBlockBaseOperator, List receivingStageInstances, 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 6de4e920407c..4680a17b4c33 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 @@ -24,6 +24,7 @@ import org.apache.pinot.core.common.Operator; import org.apache.pinot.query.mailbox.MailboxIdentifier; import org.apache.pinot.query.runtime.blocks.TransferableBlock; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; /** @@ -44,6 +45,11 @@ public OpChain(MultiStageOperator root, List receivingMailbox _stats = new OpChainStats(_id.toString()); } + public OpChain(MultiStageOperator root, List receivingMailboxes, int virtualServerId, + OperatorExecutionContext context) { + this(root, receivingMailboxes, virtualServerId, context.getRequestId(), context.getStageId()); + } + public Operator getRoot() { return _root; } 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 2298dbb4c443..10f97e2dbe0e 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 @@ -34,6 +34,7 @@ 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.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -61,6 +62,13 @@ public SortOperator(MultiStageOperator upstreamOperator, List col SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, requestId, stageId, serverAddress); } + public SortOperator(MultiStageOperator upstreamOperator, List collationKeys, + List collationDirections, int fetch, int offset, DataSchema dataSchema, + OperatorExecutionContext context) { + this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, context.getRequestId(), + context.getStageId(), context.getServer()); + } + @VisibleForTesting SortOperator(MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, 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 03f54ee6f455..99fdbb09c24f 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 @@ -31,6 +31,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -55,6 +56,12 @@ public class TransformOperator extends MultiStageOperator { private final DataSchema _resultSchema; private TransferableBlock _upstreamErrorBlock; + public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, + DataSchema upstreamDataSchema, OperatorExecutionContext context) { + this(upstreamOperator, resultSchema, transforms, upstreamDataSchema, context.getRequestId(), context.getStageId(), + context.getServer()); + } + public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, DataSchema upstreamDataSchema, long requestId, int stageId, VirtualServerAddress serverAddress) { super(requestId, stageId, serverAddress); 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 55a93df43fb3..58b8665b009b 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 @@ -41,6 +41,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.utils.AggregationUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -99,6 +100,16 @@ public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, + List orderSet, List orderSetDirection, + List orderSetNullDirection, List aggCalls, int lowerBound, + int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, + DataSchema resultSchema, DataSchema inputSchema, OperatorExecutionContext context) { + this(inputOperator, groupSet, orderSet, orderSetDirection, orderSetNullDirection, aggCalls, lowerBound, + upperBound, windowFrameType, constants, resultSchema, inputSchema, AggregationUtils.Accumulator.MERGERS, + context.getRequestId(), context.getStageId(), context.getServer()); + } + @VisibleForTesting public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, List orderSet, List orderSetDirection, diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java new file mode 100644 index 000000000000..18ce9769ef7e --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java @@ -0,0 +1,82 @@ +/** + * 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 java.util.Map; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.planner.StageMetadata; +import org.apache.pinot.query.routing.VirtualServerAddress; +import org.apache.pinot.query.runtime.blocks.TransferableBlock; + + +public class OperatorExecutionContext { + private final MailboxService _mailboxService; + private final long _requestId; + private final int _stageId; + + private final VirtualServerAddress _server; + private final long _timeoutMs; + private final long _deadlineMs; + protected final Map _metadataMap; + + public OperatorExecutionContext(MailboxService mailboxService, long requestId, int stageId, + VirtualServerAddress server, long timeoutMs, long deadlineMs, Map metadataMap) { + _mailboxService = mailboxService; + _requestId = requestId; + _stageId = stageId; + _server = server; + _timeoutMs = timeoutMs; + _deadlineMs = deadlineMs; + _metadataMap = metadataMap; + } + + public OperatorExecutionContext(PlanRequestContext planRequestContext) { + this(planRequestContext.getMailboxService(), planRequestContext.getRequestId(), planRequestContext.getStageId(), + planRequestContext.getServer(), planRequestContext.getTimeoutMs(), planRequestContext.getDeadlineMs(), + planRequestContext.getMetadataMap()); + } + + public MailboxService getMailboxService() { + return _mailboxService; + } + + public long getRequestId() { + return _requestId; + } + + public int getStageId() { + return _stageId; + } + + public VirtualServerAddress getServer() { + return _server; + } + + public long getTimeoutMs() { + return _timeoutMs; + } + + public long getDeadlineMs() { + return _deadlineMs; + } + + public Map getMetadataMap() { + return _metadataMap; + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java index 8ff1c041d1ec..1098b35c23ff 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java @@ -18,8 +18,6 @@ */ package org.apache.pinot.query.runtime.plan; -import java.util.List; -import org.apache.pinot.query.planner.StageMetadata; import org.apache.pinot.query.planner.stage.AggregateNode; import org.apache.pinot.query.planner.stage.FilterNode; import org.apache.pinot.query.planner.stage.JoinNode; @@ -32,7 +30,6 @@ import org.apache.pinot.query.planner.stage.TableScanNode; import org.apache.pinot.query.planner.stage.ValueNode; import org.apache.pinot.query.planner.stage.WindowNode; -import org.apache.pinot.query.routing.VirtualServer; import org.apache.pinot.query.runtime.operator.AggregateOperator; import org.apache.pinot.query.runtime.operator.FilterOperator; import org.apache.pinot.query.runtime.operator.HashJoinOperator; @@ -59,17 +56,15 @@ public class PhysicalPlanVisitor implements StageNodeVisitor sendingInstances = context.getMetadataMap().get(node.getSenderStageId()).getServerInstances(); MailboxReceiveOperator mailboxReceiveOperator = - new MailboxReceiveOperator(context.getMailboxService(), sendingInstances, - node.getExchangeType(), context.getServer(), - context.getRequestId(), node.getSenderStageId(), node.getStageId(), context.getTimeoutMs()); + new MailboxReceiveOperator(node.getExchangeType(), node.getSenderStageId(), node.getStageId(), + context.getOperatorExecutionContext()); context.addReceivingMailboxes(mailboxReceiveOperator.getSendingMailbox()); return mailboxReceiveOperator; } @@ -77,18 +72,15 @@ public MultiStageOperator visitMailboxReceive(MailboxReceiveNode node, PlanReque @Override public MultiStageOperator visitMailboxSend(MailboxSendNode node, PlanRequestContext context) { MultiStageOperator nextOperator = node.getInputs().get(0).visit(this, context); - StageMetadata receivingStageMetadata = context.getMetadataMap().get(node.getReceiverStageId()); - return new MailboxSendOperator(context.getMailboxService(), nextOperator, - receivingStageMetadata.getServerInstances(), node.getExchangeType(), node.getPartitionKeySelector(), - context.getServer(), context.getRequestId(), node.getStageId(), node.getReceiverStageId(), - context.getDeadlineMs()); + return new MailboxSendOperator(nextOperator, node.getExchangeType(), node.getPartitionKeySelector(), + node.getStageId(), node.getReceiverStageId(), context.getOperatorExecutionContext()); } @Override public MultiStageOperator visitAggregate(AggregateNode node, PlanRequestContext context) { MultiStageOperator nextOperator = node.getInputs().get(0).visit(this, context); return new AggregateOperator(nextOperator, node.getDataSchema(), node.getAggCalls(), node.getGroupSet(), - node.getInputs().get(0).getDataSchema(), context._requestId, context._stageId, context.getServer()); + node.getInputs().get(0).getDataSchema(), context.getOperatorExecutionContext()); } @Override @@ -97,14 +89,14 @@ public MultiStageOperator visitWindow(WindowNode node, PlanRequestContext contex return new WindowAggregateOperator(nextOperator, node.getGroupSet(), node.getOrderSet(), node.getOrderSetDirection(), node.getOrderSetNullDirection(), node.getAggCalls(), node.getLowerBound(), node.getUpperBound(), node.getWindowFrameType(), node.getConstants(), node.getDataSchema(), - node.getInputs().get(0).getDataSchema(), context._requestId, context._stageId, context.getServer()); + node.getInputs().get(0).getDataSchema(), context.getOperatorExecutionContext()); } @Override public MultiStageOperator visitFilter(FilterNode node, PlanRequestContext context) { MultiStageOperator nextOperator = node.getInputs().get(0).visit(this, context); - return new FilterOperator(nextOperator, node.getDataSchema(), node.getCondition(), context.getRequestId(), - context.getStageId(), context.getServer()); + return new FilterOperator(nextOperator, node.getDataSchema(), node.getCondition(), + context.getOperatorExecutionContext()); } @Override @@ -115,22 +107,22 @@ public MultiStageOperator visitJoin(JoinNode node, PlanRequestContext context) { MultiStageOperator leftOperator = left.visit(this, context); MultiStageOperator rightOperator = right.visit(this, context); - return new HashJoinOperator(leftOperator, rightOperator, left.getDataSchema(), node, context.getRequestId(), - context.getStageId(), context.getServer()); + return new HashJoinOperator(leftOperator, rightOperator, left.getDataSchema(), node, + context.getOperatorExecutionContext()); } @Override public MultiStageOperator visitProject(ProjectNode node, PlanRequestContext context) { MultiStageOperator nextOperator = node.getInputs().get(0).visit(this, context); return new TransformOperator(nextOperator, node.getDataSchema(), node.getProjects(), - node.getInputs().get(0).getDataSchema(), context.getRequestId(), context.getStageId(), context.getServer()); + node.getInputs().get(0).getDataSchema(), context.getOperatorExecutionContext()); } @Override public MultiStageOperator visitSort(SortNode node, PlanRequestContext context) { MultiStageOperator nextOperator = node.getInputs().get(0).visit(this, context); return new SortOperator(nextOperator, node.getCollationKeys(), node.getCollationDirections(), node.getFetch(), - node.getOffset(), node.getDataSchema(), context.getRequestId(), context.getStageId(), context.getServer()); + node.getOffset(), node.getDataSchema(), context.getOperatorExecutionContext()); } @Override @@ -140,7 +132,6 @@ public MultiStageOperator visitTableScan(TableScanNode node, PlanRequestContext @Override public MultiStageOperator visitValue(ValueNode node, PlanRequestContext context) { - return new LiteralValueOperator(node.getDataSchema(), node.getLiteralRows(), context.getRequestId(), - context.getStageId(), context.getServer()); + return new LiteralValueOperator(node.getDataSchema(), node.getLiteralRows(), context.getOperatorExecutionContext()); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanRequestContext.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanRequestContext.java index dac0f9757587..584aa6357a62 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanRequestContext.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanRequestContext.java @@ -87,4 +87,8 @@ public void addReceivingMailboxes(List ids) { public List getReceivingMailboxes() { return ImmutableList.copyOf(_receivingMailboxes); } + + public OperatorExecutionContext getOperatorExecutionContext() { + return new OperatorExecutionContext(this); + } } From 7def37583c6a71d4679a9b1b6552e12aa6ece2ee Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 13 Mar 2023 15:12:15 +0530 Subject: [PATCH 02/12] Remove no-context operator constructor and replace with new constructor in tests --- .../runtime/operator/AggregateOperator.java | 7 - .../runtime/operator/FilterOperator.java | 8 +- .../runtime/operator/HashJoinOperator.java | 9 +- .../operator/LiteralValueOperator.java | 8 +- .../query/runtime/operator/SortOperator.java | 12 +- .../runtime/operator/TransformOperator.java | 8 +- .../operator/WindowAggregateOperator.java | 11 - .../operator/AggregateOperatorTest.java | 24 +- .../runtime/operator/FilterOperatorTest.java | 45 ++-- .../operator/HashJoinOperatorTest.java | 44 ++-- ...eafStageTransferableBlockOperatorTest.java | 26 +- .../operator/LiteralValueOperatorTest.java | 4 +- .../runtime/operator/OperatorTestUtil.java | 28 ++- .../runtime/operator/SortOperatorTest.java | 42 ++-- .../operator/TransformOperatorTest.java | 36 ++- .../operator/WindowAggregateOperatorTest.java | 234 +++++++++--------- 16 files changed, 276 insertions(+), 270 deletions(-) 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 7349d68d6318..a7c56f17228f 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 @@ -81,13 +81,6 @@ public class AggregateOperator extends MultiStageOperator { // aggCalls has to be a list of FunctionCall and cannot be null // groupSet has to be a list of InputRef and cannot be null // TODO: Add these two checks when we confirm we can handle error in upstream ctor call. - public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, - List groupSet, DataSchema inputSchema, long requestId, int stageId, - VirtualServerAddress virtualServerAddress) { - this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.AGG_MERGERS, - requestId, stageId, virtualServerAddress); - } - public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, OperatorExecutionContext context) { this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.MERGERS, 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 ee2f724d4986..7f7a4361e4f4 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 @@ -25,7 +25,6 @@ import org.apache.pinot.common.datablock.DataBlock; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; -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.operator.operands.TransformOperand; @@ -58,12 +57,7 @@ public class FilterOperator extends MultiStageOperator { public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, OperatorExecutionContext context) { - this(upstreamOperator, dataSchema, filter, context.getRequestId(), context.getStageId(), context.getServer()); - } - - public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, - long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + super(context.getRequestId(), context.getStageId(), context.getServer()); _upstreamOperator = upstreamOperator; _dataSchema = dataSchema; _filterOperand = TransformOperand.toTransformOperand(filter, dataSchema); 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 d1a02687666c..c807a7188a9e 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 @@ -35,7 +35,6 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.partitioning.KeySelector; import org.apache.pinot.query.planner.stage.JoinNode; -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.operator.operands.TransformOperand; @@ -91,13 +90,7 @@ public class HashJoinOperator extends MultiStageOperator { public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, DataSchema leftSchema, JoinNode node, OperatorExecutionContext context) { - this(leftTableOperator, rightTableOperator, leftSchema, node, context.getRequestId(), context.getStageId(), - context.getServer()); - } - - public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, - DataSchema leftSchema, JoinNode node, long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + super(context.getRequestId(), context.getStageId(), context.getServer()); Preconditions.checkState(SUPPORTED_JOIN_TYPES.contains(node.getJoinRelType()), "Join type: " + node.getJoinRelType() + " is not supported!"); _joinType = node.getJoinRelType(); 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 1fe8e8ffbb4e..dac8278497e2 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 @@ -25,7 +25,6 @@ import org.apache.pinot.common.datablock.DataBlock; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; -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.plan.OperatorExecutionContext; @@ -43,12 +42,7 @@ public class LiteralValueOperator extends MultiStageOperator { public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, OperatorExecutionContext context) { - this(dataSchema, rexLiteralRows, context.getRequestId(), context.getStageId(), context.getServer()); - } - - public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, - long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + super(context.getRequestId(), context.getStageId(), context.getServer()); _dataSchema = dataSchema; _rexLiteralBlock = constructBlock(rexLiteralRows); _isLiteralBlockReturned = false; 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 10f97e2dbe0e..6afc1bb747c1 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 @@ -55,18 +55,12 @@ public class SortOperator extends MultiStageOperator { private boolean _isSortedBlockConstructed; private TransferableBlock _upstreamErrorBlock; - public SortOperator(MultiStageOperator upstreamOperator, List collationKeys, - List collationDirections, int fetch, int offset, DataSchema dataSchema, - long requestId, int stageId, VirtualServerAddress serverAddress) { - this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, - SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, requestId, stageId, serverAddress); - } - public SortOperator(MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, OperatorExecutionContext context) { - this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, context.getRequestId(), - context.getStageId(), context.getServer()); + this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, + SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, context.getRequestId(), context.getStageId(), + context.getServer()); } @VisibleForTesting 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 99fdbb09c24f..bb57c41778cf 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 @@ -26,7 +26,6 @@ import org.apache.pinot.common.datablock.DataBlock; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; -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.operator.operands.TransformOperand; @@ -58,13 +57,8 @@ public class TransformOperator extends MultiStageOperator { public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, DataSchema upstreamDataSchema, OperatorExecutionContext context) { - this(upstreamOperator, resultSchema, transforms, upstreamDataSchema, context.getRequestId(), context.getStageId(), + super(context.getRequestId(), context.getStageId(), context.getServer()); - } - - public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, - DataSchema upstreamDataSchema, long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); Preconditions.checkState(!transforms.isEmpty(), "transform operand should not be empty."); Preconditions.checkState(resultSchema.size() == transforms.size(), "result schema size:" + resultSchema.size() + " doesn't match transform operand size:" + transforms.size()); 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 58b8665b009b..240d2bfd45db 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 @@ -89,17 +89,6 @@ public class WindowAggregateOperator extends MultiStageOperator { private boolean _readyToConstruct; private boolean _hasReturnedWindowAggregateBlock; - public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, - List orderSet, List orderSetDirection, - List orderSetNullDirection, List aggCalls, int lowerBound, - int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, - DataSchema resultSchema, DataSchema inputSchema, long requestId, int stageId, - VirtualServerAddress virtualServerAddress) { - this(inputOperator, groupSet, orderSet, orderSetDirection, orderSetNullDirection, aggCalls, lowerBound, - upperBound, windowFrameType, constants, resultSchema, inputSchema, AggregationUtils.Accumulator.MERGERS, - requestId, stageId, virtualServerAddress); - } - public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, List orderSet, List orderSetDirection, List orderSetNullDirection, List aggCalls, int lowerBound, 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 85c9b7c2a30e..e57f24ff2826 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 @@ -77,7 +77,8 @@ public void shouldHandleUpstreamErrorBlocks() { DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, INT}); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); // build @@ -97,7 +98,8 @@ public void shouldHandleEndOfStreamBlockWithNoOtherInputs() { DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, INT}); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block = operator.nextBlock(); @@ -119,7 +121,8 @@ public void shouldHandleUpstreamNoOpBlocksWhileConstructing() { .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); // build when reading NoOp block @@ -142,7 +145,8 @@ public void shouldAggregateSingleInputBlock() { .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -167,7 +171,8 @@ public void shouldAggregateSingleInputBlockWithLiteralInput() { .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -222,7 +227,8 @@ public void testGroupByAggregateWithHashCollision() { DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, INT}); AggregateOperator sum0GroupBy1 = new AggregateOperator(upstreamOperator, OperatorTestUtil.getDataSchema(OperatorTestUtil.OP_1), - Arrays.asList(agg), Arrays.asList(new RexExpression.InputRef(1)), inSchema, 1, 2, _serverAddress); + Arrays.asList(agg), Arrays.asList(new RexExpression.InputRef(1)), inSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = sum0GroupBy1.getNextBlock(); while (result.isNoOpBlock()) { result = sum0GroupBy1.getNextBlock(); @@ -245,7 +251,8 @@ public void shouldThrowOnUnknownAggFunction() { DataSchema inSchema = new DataSchema(new String[]{"unknown"}, new ColumnDataType[]{DOUBLE}); // When: - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); } @Test @@ -262,7 +269,8 @@ public void shouldReturnErrorBlockOnUnexpectedInputType() { .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); - AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group, inSchema, 1, 2, _serverAddress); + AggregateOperator operator = + new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block = operator.nextBlock(); 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 26e40cc4ab82..3ee4aa5cca30 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 @@ -25,7 +25,6 @@ import org.apache.pinot.common.exception.QueryException; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; -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.spi.data.FieldSpec; @@ -43,13 +42,9 @@ public class FilterOperatorTest { @Mock private MultiStageOperator _upstreamOperator; - @Mock - private VirtualServerAddress _serverAddress; - @BeforeMethod public void setUp() { _mocks = MockitoAnnotations.openMocks(this); - Mockito.when(_serverAddress.toString()).thenReturn(new VirtualServerAddress("mock", 80, 0).toString()); } @AfterMethod @@ -66,7 +61,8 @@ public void shouldPropagateUpstreamErrorBlock() { DataSchema inputSchema = new DataSchema(new String[]{"boolCol"}, new DataSchema.ColumnDataType[]{ DataSchema.ColumnDataType.BOOLEAN }); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock error = errorBlock.getDataBlock(); @@ -81,7 +77,8 @@ public void shouldPropagateUpstreamEOS() { DataSchema.ColumnDataType.INT }); Mockito.when(_upstreamOperator.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertTrue(dataBlock.isEndOfStreamBlock()); } @@ -94,7 +91,8 @@ public void shouldPropagateUpstreamNoop() { DataSchema.ColumnDataType.INT }); Mockito.when(_upstreamOperator.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertTrue(dataBlock.isNoOpBlock()); } @@ -109,7 +107,8 @@ public void shouldHandleTrueBooleanLiteralFilter() { Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{0}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -127,7 +126,8 @@ public void shouldHandleFalseBooleanLiteralFilter() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -142,7 +142,8 @@ public void shouldThrowOnNonBooleanTypeBooleanLiteral() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock data = errorBlock.getDataBlock(); @@ -157,7 +158,7 @@ public void shouldThrowOnNonBooleanTypeInputRef() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref0, 1, 2, _serverAddress); + FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref0, OperatorTestUtil.getDefaultContext()); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock data = errorBlock.getDataBlock(); @@ -172,7 +173,7 @@ public void shouldHandleBooleanInputRef() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1, true}, new Object[]{2, false})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref1, 1, 2, _serverAddress); + FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref1, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -192,7 +193,8 @@ public void shouldHandleAndFilter() { RexExpression.FunctionCall andCall = new RexExpression.FunctionCall(SqlKind.AND, FieldSpec.DataType.BOOLEAN, "AND", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, andCall, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, andCall, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -212,7 +214,8 @@ public void shouldHandleOrFilter() { RexExpression.FunctionCall orCall = new RexExpression.FunctionCall(SqlKind.OR, FieldSpec.DataType.BOOLEAN, "OR", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, orCall, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, orCall, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -234,7 +237,8 @@ public void shouldHandleNotFilter() { RexExpression.FunctionCall notCall = new RexExpression.FunctionCall(SqlKind.NOT, FieldSpec.DataType.BOOLEAN, "NOT", ImmutableList.of(new RexExpression.InputRef(0))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, notCall, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, notCall, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -253,7 +257,8 @@ public void shouldHandleGreaterThanFilter() { RexExpression.FunctionCall greaterThan = new RexExpression.FunctionCall(SqlKind.GREATER_THAN, FieldSpec.DataType.BOOLEAN, "greaterThan", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, greaterThan, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, greaterThan, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -273,7 +278,8 @@ public void shouldHandleBooleanFunction() { new RexExpression.FunctionCall(SqlKind.OTHER, FieldSpec.DataType.BOOLEAN, "startsWith", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(FieldSpec.DataType.STRING, "star"))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, startsWith, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, startsWith, OperatorTestUtil.getDefaultContext()); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -294,6 +300,7 @@ public void shouldThrowOnUnfoundFunction() { new RexExpression.FunctionCall(SqlKind.OTHER, FieldSpec.DataType.BOOLEAN, "startsWithError", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(FieldSpec.DataType.STRING, "star"))); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, startsWith, 1, 2, _serverAddress); + FilterOperator op = + new FilterOperator(_upstreamOperator, inputSchema, startsWith, OperatorTestUtil.getDefaultContext()); } } 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 c77107a58cd4..68de25a4a5f3 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 @@ -96,7 +96,7 @@ public void shouldHandleHashJoinKeyCollisionInnerJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); HashJoinOperator joinOnString = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = joinOnString.nextBlock(); while (result.isNoOpBlock()) { @@ -134,7 +134,7 @@ public void shouldHandleInnerJoinOnInt() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator joinOnInt = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = joinOnInt.nextBlock(); while (result.isNoOpBlock()) { result = joinOnInt.nextBlock(); @@ -169,7 +169,7 @@ public void shouldHandleJoinOnEmptySelector() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); HashJoinOperator joinOnInt = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = joinOnInt.nextBlock(); while (result.isNoOpBlock()) { result = joinOnInt.nextBlock(); @@ -210,7 +210,8 @@ public void shouldHandleLeftJoin() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -244,7 +245,8 @@ public void shouldPassLeftTableEOS() { List joinClauses = new ArrayList<>(); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -275,7 +277,8 @@ public void shouldHandleLeftJoinOneToN() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -310,7 +313,8 @@ public void shouldPassRightTableEOS() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -348,7 +352,8 @@ public void shouldHandleInequiJoinOnString() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -386,7 +391,8 @@ public void shouldHandleInequiJoinOnInt() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -421,7 +427,7 @@ public void shouldHandleRightJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.RIGHT, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator joinOnNum = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = joinOnNum.nextBlock(); while (result.isNoOpBlock()) { result = joinOnNum.nextBlock(); @@ -470,7 +476,8 @@ public void shouldHandleSemiJoin() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.SEMI, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -509,7 +516,8 @@ public void shouldHandleFullJoin() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.FULL, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -561,7 +569,8 @@ public void shouldHandleAntiJoin() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.ANTI, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -599,7 +608,8 @@ public void shouldPropagateRightTableError() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -632,7 +642,8 @@ public void shouldPropagateLeftTableError() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -668,7 +679,8 @@ public void shouldHandleNoOpBlock() { }); JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); - HashJoinOperator join = new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, 1, 2, _serverAddress); + HashJoinOperator join = + new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); TransferableBlock result = join.nextBlock(); // first no-op consumes first right data block. Assert.assertTrue(result.isNoOpBlock()); 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 a8332f22d222..f5a39d735e91 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 @@ -75,7 +75,7 @@ public void shouldReturnDataBlockThenMetadataBlock() { List resultsBlockList = Collections.singletonList(new InstanceResponseBlock( new SelectionResultsBlock(schema, Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -101,7 +101,7 @@ public void shouldHandleDesiredDataSchemaConversionCorrectly() { new SelectionResultsBlock(resultSchema, Arrays.asList(new Object[]{1, 1660000000000L}, new Object[]{0, 1600000000000L})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, desiredSchema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -123,7 +123,7 @@ public void shouldHandleCanonicalizationCorrectly() { new SelectionResultsBlock(schema, Arrays.asList(new Object[]{1, 1660000000000L}, new Object[]{0, 1600000000000L})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -148,7 +148,7 @@ public void shouldReturnMultipleDataBlockThenMetadataBlock() { queryContext), new InstanceResponseBlock(new SelectionResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock1 = operator.nextBlock(); @@ -178,7 +178,7 @@ public void shouldGetErrorBlockWhenInstanceResponseContainsError() { errorBlock, new InstanceResponseBlock(new SelectionResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -199,7 +199,7 @@ public void shouldReorderWhenQueryContextAskForNotInOrderGroupByAsDistinct() { new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(schema, Arrays.asList(new Record(new Object[]{1, "foo"}), new Record(new Object[]{2, "bar"})))), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -220,7 +220,7 @@ public void shouldParsedBlocksSuccessfullyWithDistinctQuery() { new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(schema, Arrays.asList(new Record(new Object[]{"foo", 1}), new Record(new Object[]{"bar", 2})))), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -244,7 +244,7 @@ public void shouldReorderWhenQueryContextAskForGroupByOutOfOrder() { List resultsBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -267,7 +267,7 @@ public void shouldNotErrorOutWhenQueryContextAskForGroupByOutOfOrderWithHaving() List resultsBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -286,7 +286,7 @@ public void shouldNotErrorOutWhenDealingWithAggregationResults() { List resultsBlockList = Collections.singletonList(new InstanceResponseBlock( new AggregationResultsBlock(queryContext.getAggregationFunctions(), Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -308,7 +308,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsSelecti List responseBlockList = Collections.singletonList( new InstanceResponseBlock(new SelectionResultsBlock(resultSchema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); TransferableBlock resultBlock = operator.nextBlock(); // Then: @@ -331,7 +331,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsDistinc new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(resultSchema, Collections.emptyList())), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); TransferableBlock resultBlock = operator.nextBlock(); // Then: @@ -353,7 +353,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsGroupBy List responseBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(resultSchema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, 1, 2, _serverAddress); + new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); TransferableBlock resultBlock = operator.nextBlock(); // Then: 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 749431e090ac..705afc0192ca 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 @@ -66,7 +66,7 @@ public void shouldReturnLiteralBlock() { List> literals = ImmutableList.of( ImmutableList.of(new RexExpression.Literal(DataType.STRING, "foo"), new RexExpression.Literal(DataType.INT, 1)), ImmutableList.of(new RexExpression.Literal(DataType.STRING, ""), new RexExpression.Literal(DataType.INT, 2))); - LiteralValueOperator operator = new LiteralValueOperator(schema, literals, 1, 2, _serverAddress); + LiteralValueOperator operator = new LiteralValueOperator(schema, literals, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock transferableBlock = operator.nextBlock(); @@ -82,7 +82,7 @@ public void shouldHandleEmptyLiteralRows() { // Given: DataSchema schema = new DataSchema(new String[]{}, new ColumnDataType[]{}); List> literals = ImmutableList.of(ImmutableList.of()); - LiteralValueOperator operator = new LiteralValueOperator(schema, literals, 1, 2, _serverAddress); + LiteralValueOperator operator = new LiteralValueOperator(schema, literals, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock transferableBlock = operator.nextBlock(); 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 39921e324d92..3b456f0b6765 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 @@ -19,18 +19,20 @@ package org.apache.pinot.query.runtime.operator; import java.util.Arrays; +import java.util.HashMap; import java.util.List; 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.plan.OperatorExecutionContext; import org.apache.pinot.query.testutils.MockDataBlockOperatorFactory; - public class OperatorTestUtil { // simple key-value collision schema/data test set: "Aa" and "BB" have same hash code in java. - private static final List> SIMPLE_KV_DATA_ROWS = Arrays.asList( - Arrays.asList(new Object[]{1, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}), - Arrays.asList(new Object[]{1, "AA"}, new Object[]{2, "Aa"})); + private static final List> SIMPLE_KV_DATA_ROWS = + Arrays.asList(Arrays.asList(new Object[]{1, "Aa"}, new Object[]{2, "BB"}, new Object[]{3, "BB"}), + Arrays.asList(new Object[]{1, "AA"}, new Object[]{2, "Aa"})); private static final MockDataBlockOperatorFactory MOCK_OPERATOR_FACTORY; public static final DataSchema SIMPLE_KV_DATA_SCHEMA = new DataSchema(new String[]{"foo", "bar"}, @@ -40,10 +42,8 @@ public class OperatorTestUtil { public static final String OP_2 = "op2"; static { - MOCK_OPERATOR_FACTORY = new MockDataBlockOperatorFactory() - .registerOperator(OP_1, SIMPLE_KV_DATA_SCHEMA) - .registerOperator(OP_2, SIMPLE_KV_DATA_SCHEMA) - .addRows(OP_1, SIMPLE_KV_DATA_ROWS.get(0)) + MOCK_OPERATOR_FACTORY = new MockDataBlockOperatorFactory().registerOperator(OP_1, SIMPLE_KV_DATA_SCHEMA) + .registerOperator(OP_2, SIMPLE_KV_DATA_SCHEMA).addRows(OP_1, SIMPLE_KV_DATA_ROWS.get(0)) .addRows(OP_2, SIMPLE_KV_DATA_ROWS.get(1)); } @@ -61,4 +61,16 @@ public static DataSchema getDataSchema(String operatorName) { public static TransferableBlock block(DataSchema schema, Object[]... rows) { return new TransferableBlock(Arrays.asList(rows), schema, DataBlock.Type.ROW); } + + public static OperatorExecutionContext getDefaultContext() { + VirtualServerAddress virtualServerAddress = new VirtualServerAddress("mock", 80, 0); + return new OperatorExecutionContext(null, 1, 2, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, + new HashMap<>()); + } + + public static OperatorExecutionContext getContext(long requestId, int stageId, + VirtualServerAddress virtualServerAddress) { + return new OperatorExecutionContext(null, requestId, stageId, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, + new HashMap<>()); + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 04c3b8c700b4..3b5166ef5ae1 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -69,7 +69,8 @@ public void shouldHandleUpstreamErrorBlock() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()) .thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!"))); @@ -87,7 +88,8 @@ public void shouldHandleUpstreamNoOpBlock() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); @@ -104,7 +106,8 @@ public void shouldCreateEmptyBlockOnUpstreamEOS() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -121,7 +124,8 @@ public void shouldConsumeAndSortInputOneBlockWithTwoRows() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -143,7 +147,8 @@ public void shouldConsumeAndSortOnNonZeroIdxCollation() { List collation = collation(1); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"ignored", "sort"}, new DataSchema.ColumnDataType[]{INT, INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{1, 2}, new Object[]{2, 1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -165,7 +170,8 @@ public void shouldConsumeAndSortInputOneBlockWithTwoRowsNonNumeric() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{STRING}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{"b"}, new Object[]{"a"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -187,7 +193,8 @@ public void shouldConsumeAndSortDescending() { List collation = collation(0); List directions = ImmutableList.of(Direction.DESCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -209,7 +216,8 @@ public void shouldOffsetSortInputOneBlockWithThreeRows() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 1, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 1, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -231,7 +239,8 @@ public void shouldOffsetLimitSortInputOneBlockWithThreeRows() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 1, 1, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 1, 1, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -273,7 +282,8 @@ public void shouldFetchAllWithNegativeFetch() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, -1, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, -1, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -293,7 +303,8 @@ public void shouldConsumeAndSortTwoInputBlocksWithOneRowEach() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2})) .thenReturn(block(schema, new Object[]{1})) @@ -316,7 +327,8 @@ public void shouldBreakTiesUsingSecondCollationKey() { List collation = collation(0, 1); List directions = ImmutableList.of(Direction.ASCENDING, Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"first", "second"}, new DataSchema.ColumnDataType[]{INT, INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()) .thenReturn(block(schema, new Object[]{1, 2}, new Object[]{1, 1}, new Object[]{1, 3})) @@ -340,7 +352,8 @@ public void shouldBreakTiesUsingSecondCollationKeyWithDifferentDirection() { List collation = collation(0, 1); List directions = ImmutableList.of(Direction.ASCENDING, Direction.DESCENDING); DataSchema schema = new DataSchema(new String[]{"first", "second"}, new DataSchema.ColumnDataType[]{INT, INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()) .thenReturn(block(schema, new Object[]{1, 2}, new Object[]{1, 1}, new Object[]{1, 3})) @@ -364,7 +377,8 @@ public void shouldHandleNoOpUpstreamBlockWhileConstructing() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 10, 0, schema, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2})) .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()).thenReturn(block(schema, new Object[]{1})) diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java index f990882f7426..4afc580cdbc0 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java @@ -26,7 +26,6 @@ import org.apache.pinot.common.exception.QueryException; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; -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.spi.data.FieldSpec; @@ -48,13 +47,9 @@ public class TransformOperatorTest { @Mock private MultiStageOperator _upstreamOp; - @Mock - private VirtualServerAddress _serverAddress; - @BeforeMethod public void setUp() { _mocks = MockitoAnnotations.openMocks(this); - Mockito.when(_serverAddress.toString()).thenReturn(new VirtualServerAddress("mock", 80, 0).toString()); } @AfterMethod @@ -76,8 +71,8 @@ public void shouldHandleRefTransform() { RexExpression.InputRef ref0 = new RexExpression.InputRef(0); RexExpression.InputRef ref1 = new RexExpression.InputRef(1); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0, ref1), upStreamSchema, 1, 2, - _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0, ref1), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); Assert.assertTrue(!result.isErrorBlock()); @@ -101,8 +96,8 @@ public void shouldHandleLiteralTransform() { RexExpression.Literal boolLiteral = new RexExpression.Literal(FieldSpec.DataType.BOOLEAN, true); RexExpression.Literal strLiteral = new RexExpression.Literal(FieldSpec.DataType.STRING, "str"); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, 1, - 2, _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); // Literal operands should just output original literals. Assert.assertTrue(!result.isErrorBlock()); @@ -132,8 +127,8 @@ public void shouldHandlePlusMinusFuncTransform() { DataSchema resultSchema = new DataSchema(new String[]{"plusR", "minusR"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.DOUBLE, DataSchema.ColumnDataType.DOUBLE}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, 1, 2, - _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); Assert.assertTrue(!result.isErrorBlock()); List resultRows = result.getContainer(); @@ -161,8 +156,8 @@ public void shouldThrowOnTypeMismatchFuncTransform() { DataSchema resultSchema = new DataSchema(new String[]{"plusR", "minusR"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.DOUBLE, DataSchema.ColumnDataType.DOUBLE}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, 1, 2, - _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); Assert.assertTrue(result.isErrorBlock()); @@ -182,8 +177,8 @@ public void shouldPropagateUpstreamError() { DataSchema resultSchema = new DataSchema(new String[]{"inCol", "strCol"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, 1, - 2, _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); Assert.assertTrue(result.isErrorBlock()); DataBlock data = result.getDataBlock(); @@ -206,8 +201,8 @@ public void testNoopBlock() { DataSchema resultSchema = new DataSchema(new String[]{"boolCol", "strCol"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BOOLEAN, DataSchema.ColumnDataType.STRING}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, 1, - 2, _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, + OperatorTestUtil.getDefaultContext()); TransferableBlock result = op.nextBlock(); // First block has two rows Assert.assertFalse(result.isErrorBlock()); @@ -238,8 +233,8 @@ public void testWrongNumTransform() { DataSchema upStreamSchema = new DataSchema(new String[]{"string1", "string2"}, new DataSchema.ColumnDataType[]{ DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.STRING }); - TransformOperator transform = - new TransformOperator(_upstreamOp, resultSchema, new ArrayList<>(), upStreamSchema, 1, 2, _serverAddress); + TransformOperator transform = new TransformOperator(_upstreamOp, resultSchema, new ArrayList<>(), upStreamSchema, + OperatorTestUtil.getDefaultContext()); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*doesn't match " @@ -252,6 +247,7 @@ public void testMismatchedSchemaOperandSize() { }); RexExpression.InputRef ref0 = new RexExpression.InputRef(0); TransformOperator transform = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0), upStreamSchema, 1, 2, _serverAddress); + new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0), upStreamSchema, + OperatorTestUtil.getDefaultContext()); } }; diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 42606e5422c6..04271d51a372 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -78,11 +78,12 @@ public void testShouldHandleUpstreamErrorBlocks() { .thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!"))); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, INT}); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); // build @@ -101,11 +102,12 @@ public void testShouldHandleEndOfStreamBlockWithNoOtherInputs() { Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, INT}); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block = operator.nextBlock(); @@ -126,11 +128,12 @@ public void testShouldHandleUpstreamNoOpBlocksWhileConstructing() { .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); // build when reading NoOp block @@ -156,11 +159,12 @@ public void testShouldHandleUpstreamNoOpBlocksWhileConstructingMultipleRows() { .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{1, 2})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); // build when reading NoOp block @@ -189,11 +193,12 @@ public void testShouldWindowAggregateOverSingleInputBlock() { Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -218,12 +223,13 @@ public void testShouldWindowAggregateOverSingleInputBlockWithSameOrderByKeys() { Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, order, - Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), - calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), - outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), + Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, + OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -246,12 +252,12 @@ public void testShouldWindowAggregateOverSingleInputBlockWithoutPartitionByKeys( Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, - Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, - _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -275,11 +281,12 @@ public void testShouldWindowAggregateOverSingleInputBlockWithLiteralInput() { Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, 3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -309,12 +316,12 @@ public void testShouldCallMergerWhenWindowAggregatingMultipleRows() { AggregationUtils.Merger merger = Mockito.mock(AggregationUtils.Merger.class); Mockito.when(merger.merge(Mockito.any(), Mockito.any())).thenReturn(12d); Mockito.when(merger.initialize(Mockito.any(), Mockito.any())).thenReturn(1d); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, - ImmutableMap.of("SUM", cdt -> merger), 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, ImmutableMap.of("SUM", cdt -> merger), 1, 2, _serverAddress); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) @@ -332,28 +339,27 @@ public void testShouldCallMergerWhenWindowAggregatingMultipleRows() { "Expected three columns (original two columns, agg literal value)"); } - @Test public void testPartitionByWindowAggregateWithHashCollision() { MultiStageOperator upstreamOperator = OperatorTestUtil.getOperator(OperatorTestUtil.OP_1); // Create an aggregation call with sum for first column and group by second column. RexExpression.FunctionCall agg = getSum(new RexExpression.InputRef(0)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, INT}); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator sum0PartitionBy1 = new WindowAggregateOperator(upstreamOperator, Arrays.asList(new RexExpression.InputRef(1)), Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Arrays.asList(agg), Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, - inSchema, 1, 2, _serverAddress); + inSchema, OperatorTestUtil.getDefaultContext()); TransferableBlock result = sum0PartitionBy1.getNextBlock(); while (result.isNoOpBlock()) { result = sum0PartitionBy1.getNextBlock(); } List resultRows = result.getContainer(); - List expectedRows = Arrays.asList(new Object[]{2, "BB", 5.0}, new Object[]{3, "BB", 5.0}, - new Object[]{1, "Aa", 1}); + List expectedRows = + Arrays.asList(new Object[]{2, "BB", 5.0}, new Object[]{3, "BB", 5.0}, new Object[]{1, "Aa", 1}); Assert.assertEquals(resultRows.size(), expectedRows.size()); Assert.assertEquals(resultRows.get(0), expectedRows.get(0)); Assert.assertEquals(resultRows.get(1), expectedRows.get(1)); @@ -371,9 +377,10 @@ public void testShouldThrowOnUnknownAggFunction() { DataSchema inSchema = new DataSchema(new String[]{"unknown"}, new DataSchema.ColumnDataType[]{DOUBLE}); // When: - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*Unexpected aggregation " @@ -388,9 +395,10 @@ public void testShouldThrowOnUnknownRankAggFunction() { DataSchema inSchema = new DataSchema(new String[]{"unknown"}, new DataSchema.ColumnDataType[]{DOUBLE}); // When: - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Order by is not yet " @@ -403,16 +411,16 @@ public void testShouldThrowOnNonEmptyOrderByKeysNotMatchingPartitionByKeys() { List order = ImmutableList.of(new RexExpression.InputRef(1)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, order, - Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), - calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), - outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), + Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, + OperatorTestUtil.getDefaultContext()); } @Test @@ -426,18 +434,18 @@ public void testShouldThrowOnNonEmptyOrderByKeysMatchingPartitionByKeysWithDiffe List order = ImmutableList.of(new RexExpression.InputRef(1)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "bar"})) .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{3, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, order, - Arrays.asList(RelFieldCollation.Direction.DESCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), - calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), - outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.DESCENDING), + Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, + OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) @@ -460,16 +468,15 @@ public void testShouldThrowOnCustomFramesRows() { List group = ImmutableList.of(new RexExpression.InputRef(0)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.ROW, Collections.emptyList(), outSchema, inSchema, 1, 2, - _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.ROW, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); } @Test @@ -480,16 +487,16 @@ public void testShouldNotThrowCurrentRowPartitionByOrderByOnSameKey() { List order = ImmutableList.of(new RexExpression.InputRef(1)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, order, - Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), - calls, Integer.MIN_VALUE, 0, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, - inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), + Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, 0, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, + OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block1 = operator.nextBlock(); @@ -512,16 +519,15 @@ public void testShouldThrowOnCustomFramesCustomPreceding() { List group = ImmutableList.of(new RexExpression.InputRef(0)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, 5, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, - _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, 5, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Only default frame is " @@ -533,15 +539,15 @@ public void testShouldThrowOnCustomFramesCustomFollowing() { List group = ImmutableList.of(new RexExpression.InputRef(0)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "foo"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, 5, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, 5, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); } @Test @@ -553,16 +559,16 @@ public void testShouldReturnErrorBlockOnUnexpectedInputType() { DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new DataSchema.ColumnDataType[]{INT, STRING}); // TODO: it is necessary to produce two values here, the operator only throws on second // (see the comment in WindowAggregate operator) - Mockito.when(_input.nextBlock()) - .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "metallica"})) + Mockito.when(_input.nextBlock()).thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "metallica"})) .thenReturn(OperatorTestUtil.block(inSchema, new Object[]{2, "pink floyd"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, - new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); - WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), - Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, 1, 2, _serverAddress); + DataSchema outSchema = + new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); + WindowAggregateOperator operator = + new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, + Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); // When: TransferableBlock block = operator.nextBlock(); From a28ebbb51b6d8c57cd3cb68f0f8e2b480976fc72 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 13 Mar 2023 15:22:43 +0530 Subject: [PATCH 03/12] Refactor MultiStageOperator to accept only context --- .../pinot/query/runtime/operator/AggregateOperator.java | 2 +- .../apache/pinot/query/runtime/operator/FilterOperator.java | 2 +- .../pinot/query/runtime/operator/HashJoinOperator.java | 2 +- .../pinot/query/runtime/operator/LiteralValueOperator.java | 2 +- .../pinot/query/runtime/operator/MultiStageOperator.java | 5 +++++ .../pinot/query/runtime/operator/TransformOperator.java | 3 +-- 6 files changed, 10 insertions(+), 6 deletions(-) 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 a7c56f17228f..7b1614a33b6f 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 @@ -83,7 +83,7 @@ public class AggregateOperator extends MultiStageOperator { // TODO: Add these two checks when we confirm we can handle error in upstream ctor call. public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, OperatorExecutionContext context) { - this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.MERGERS, + this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.AGG_MERGERS, context.getRequestId(), context.getStageId(), context.getServer()); } 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 7f7a4361e4f4..0ffc0ea9646f 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 @@ -57,7 +57,7 @@ public class FilterOperator extends MultiStageOperator { public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, OperatorExecutionContext context) { - super(context.getRequestId(), context.getStageId(), context.getServer()); + super(context); _upstreamOperator = upstreamOperator; _dataSchema = dataSchema; _filterOperand = TransformOperand.toTransformOperand(filter, dataSchema); 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 c807a7188a9e..ad479f6c7609 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 @@ -90,7 +90,7 @@ public class HashJoinOperator extends MultiStageOperator { public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, DataSchema leftSchema, JoinNode node, OperatorExecutionContext context) { - super(context.getRequestId(), context.getStageId(), context.getServer()); + super(context); Preconditions.checkState(SUPPORTED_JOIN_TYPES.contains(node.getJoinRelType()), "Join type: " + node.getJoinRelType() + " is not supported!"); _joinType = node.getJoinRelType(); 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 dac8278497e2..1acd38a4617e 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 @@ -42,7 +42,7 @@ public class LiteralValueOperator extends MultiStageOperator { public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, OperatorExecutionContext context) { - super(context.getRequestId(), context.getStageId(), context.getServer()); + super(context); _dataSchema = dataSchema; _rexLiteralBlock = constructBlock(rexLiteralRows); _isLiteralBlockReturned = false; 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 55fbc8830d34..d3415ad26acd 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 @@ -28,6 +28,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.utils.OperatorUtils; +import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; import org.apache.pinot.spi.exception.EarlyTerminationException; import org.apache.pinot.spi.trace.InvocationScope; import org.apache.pinot.spi.trace.Tracing; @@ -45,6 +46,10 @@ public abstract class MultiStageOperator implements Operator, protected final Map _operatorStatsMap; private final String _operatorId; + public MultiStageOperator(OperatorExecutionContext context) { + this(context.getRequestId(), context.getStageId(), context.getServer()); + } + public MultiStageOperator(long requestId, int stageId, VirtualServerAddress serverAddress) { _requestId = requestId; _stageId = stageId; 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 bb57c41778cf..a3a78aea6a0c 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 @@ -57,8 +57,7 @@ public class TransformOperator extends MultiStageOperator { public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, DataSchema upstreamDataSchema, OperatorExecutionContext context) { - super(context.getRequestId(), context.getStageId(), - context.getServer()); + super(context); Preconditions.checkState(!transforms.isEmpty(), "transform operand should not be empty."); Preconditions.checkState(resultSchema.size() == transforms.size(), "result schema size:" + resultSchema.size() + " doesn't match transform operand size:" + transforms.size()); From be073e4078b071825fea116f5ebafb243c615519 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Tue, 14 Mar 2023 20:05:49 +0530 Subject: [PATCH 04/12] Replace arguments with operator context in constructors used in only testing --- .../pinot/query/runtime/operator/AggregateOperator.java | 7 +++---- .../apache/pinot/query/runtime/operator/SortOperator.java | 7 +++---- .../query/runtime/operator/WindowAggregateOperator.java | 7 +++---- .../query/runtime/operator/AggregateOperatorTest.java | 4 ++-- .../pinot/query/runtime/operator/SortOperatorTest.java | 3 ++- .../runtime/operator/WindowAggregateOperatorTest.java | 3 ++- 6 files changed, 15 insertions(+), 16 deletions(-) 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 7b1614a33b6f..8c0bb42b6351 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 @@ -84,15 +84,14 @@ public class AggregateOperator extends MultiStageOperator { public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, OperatorExecutionContext context) { this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.AGG_MERGERS, - context.getRequestId(), context.getStageId(), context.getServer()); + context); } @VisibleForTesting AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, - Map> mergers, long requestId, int stageId, - VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + Map> mergers, OperatorExecutionContext context) { + super(context); _inputOperator = inputOperator; _groupSet = groupSet; _upstreamErrorBlock = null; 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 6afc1bb747c1..33853ae6d4d2 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 @@ -59,15 +59,14 @@ public SortOperator(MultiStageOperator upstreamOperator, List col List collationDirections, int fetch, int offset, DataSchema dataSchema, OperatorExecutionContext context) { this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, - SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, context.getRequestId(), context.getStageId(), - context.getServer()); + SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, context); } @VisibleForTesting SortOperator(MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, - int defaultHolderCapacity, long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + int defaultHolderCapacity, OperatorExecutionContext context) { + super(context); _upstreamOperator = upstreamOperator; _fetch = fetch; _offset = Math.max(offset, 0); 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 240d2bfd45db..3821fb065f8f 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 @@ -96,7 +96,7 @@ public WindowAggregateOperator(MultiStageOperator inputOperator, List orderSetNullDirection, List aggCalls, int lowerBound, int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, DataSchema resultSchema, DataSchema inputSchema, Map> mergers, long requestId, int stageId, - VirtualServerAddress virtualServerAddress) { - super(requestId, stageId, virtualServerAddress); + Function> mergers, OperatorExecutionContext context) { + super(context); boolean isPartitionByOnly = isPartitionByOnlyQuery(groupSet, orderSet, orderSetDirection, orderSetNullDirection); Preconditions.checkState(orderSet == null || orderSet.isEmpty() || isPartitionByOnly, 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 e57f24ff2826..ecf678e8c281 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 @@ -204,8 +204,8 @@ public void shouldCallMergerWhenAggregatingMultipleRows() { Mockito.when(merger.initialize(Mockito.any(), Mockito.any())).thenReturn(1d); DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); AggregateOperator operator = - new AggregateOperator(_input, outSchema, calls, group, inSchema, ImmutableMap.of("SUM", cdt -> merger), 1, 2, - _serverAddress); + new AggregateOperator(_input, outSchema, calls, group, inSchema, ImmutableMap.of("SUM", cdt -> merger), + OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 3b5166ef5ae1..6f311ad8e72d 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -261,7 +261,8 @@ public void shouldRespectDefaultLimit() { List collation = collation(0); List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); - SortOperator op = new SortOperator(_input, collation, directions, 0, 0, schema, 1, 1, 2, _serverAddress); + SortOperator op = + new SortOperator(_input, collation, directions, 0, 0, schema, 1, OperatorTestUtil.getDefaultContext()); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 04271d51a372..92b32383f182 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -321,7 +321,8 @@ public void testShouldCallMergerWhenWindowAggregatingMultipleRows() { WindowAggregateOperator operator = new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, ImmutableMap.of("SUM", cdt -> merger), 1, 2, _serverAddress); + Collections.emptyList(), outSchema, inSchema, ImmutableMap.of("SUM", cdt -> merger), + OperatorTestUtil.getDefaultContext()); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) From 9af4a2a840d3249f1bf6e5b0f4edb161e107a5fc Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Tue, 14 Mar 2023 20:42:08 +0530 Subject: [PATCH 05/12] Rename OperatorExecutionContext to OpChainExecutionContext --- .../runtime/operator/AggregateOperator.java | 7 +++---- .../runtime/operator/FilterOperator.java | 4 ++-- .../runtime/operator/HashJoinOperator.java | 4 ++-- .../LeafStageTransferableBlockOperator.java | 4 ++-- .../operator/LiteralValueOperator.java | 4 ++-- .../operator/MailboxReceiveOperator.java | 4 ++-- .../runtime/operator/MailboxSendOperator.java | 4 ++-- .../runtime/operator/MultiStageOperator.java | 4 ++-- .../pinot/query/runtime/operator/OpChain.java | 4 ++-- .../query/runtime/operator/SortOperator.java | 7 +++---- .../runtime/operator/TransformOperator.java | 4 ++-- .../operator/WindowAggregateOperator.java | 7 +++---- ...text.java => OpChainExecutionContext.java} | 14 ++++++++----- .../runtime/plan/PhysicalPlanVisitor.java | 20 +++++++++---------- .../runtime/plan/PlanRequestContext.java | 7 ++++--- .../runtime/operator/OperatorTestUtil.java | 10 +++++----- 16 files changed, 55 insertions(+), 53 deletions(-) rename pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/{OperatorExecutionContext.java => OpChainExecutionContext.java} (81%) 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 8c0bb42b6351..3b041c59b378 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 @@ -35,11 +35,10 @@ import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.core.data.table.Key; import org.apache.pinot.query.planner.logical.RexExpression; -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.operator.utils.AggregationUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.segment.local.customobject.PinotFourthMoment; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -82,7 +81,7 @@ public class AggregateOperator extends MultiStageOperator { // groupSet has to be a list of InputRef and cannot be null // TODO: Add these two checks when we confirm we can handle error in upstream ctor call. public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, - List groupSet, DataSchema inputSchema, OperatorExecutionContext context) { + List groupSet, DataSchema inputSchema, OpChainExecutionContext context) { this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.AGG_MERGERS, context); } @@ -90,7 +89,7 @@ public AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema @VisibleForTesting AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, List groupSet, DataSchema inputSchema, - Map> mergers, OperatorExecutionContext context) { + Map> mergers, OpChainExecutionContext context) { super(context); _inputOperator = inputOperator; _groupSet = groupSet; 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 0ffc0ea9646f..35c7915e8b06 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 @@ -29,7 +29,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -56,7 +56,7 @@ public class FilterOperator extends MultiStageOperator { private TransferableBlock _upstreamErrorBlock; public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, - OperatorExecutionContext context) { + OpChainExecutionContext context) { super(context); _upstreamOperator = upstreamOperator; _dataSchema = dataSchema; 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 ad479f6c7609..97a03ae25c08 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 @@ -39,7 +39,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -89,7 +89,7 @@ public class HashJoinOperator extends MultiStageOperator { private KeySelector _rightKeySelector; public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, - DataSchema leftSchema, JoinNode node, OperatorExecutionContext context) { + DataSchema leftSchema, JoinNode node, OpChainExecutionContext context) { super(context); Preconditions.checkState(SUPPORTED_JOIN_TYPES.contains(node.getJoinRelType()), "Join type: " + node.getJoinRelType() + " is not supported!"); 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 d670834af9fc..be607df51027 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 @@ -41,7 +41,7 @@ import org.apache.pinot.core.query.selection.SelectionOperatorUtils; import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.TransferableBlock; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -69,7 +69,7 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator { private int _currentIndex; public LeafStageTransferableBlockOperator(List baseResultBlock, DataSchema dataSchema, - OperatorExecutionContext context) { + OpChainExecutionContext context) { this(baseResultBlock, dataSchema, context.getRequestId(), context.getStageId(), context.getServer()); } 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 1acd38a4617e..af81e1e31d17 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 @@ -27,7 +27,7 @@ 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.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -41,7 +41,7 @@ public class LiteralValueOperator extends MultiStageOperator { private boolean _isLiteralBlockReturned; public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, - OperatorExecutionContext context) { + OpChainExecutionContext context) { super(context); _dataSchema = dataSchema; _rexLiteralBlock = constructBlock(rexLiteralRows); 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 1cd22c84b409..15d8b06df59c 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 @@ -36,7 +36,7 @@ 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.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.query.service.QueryConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -78,7 +78,7 @@ private static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, i } public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, - int receiverStageId, OperatorExecutionContext context) { + int receiverStageId, OpChainExecutionContext context) { this(context.getMailboxService(), context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getTimeoutMs()); } 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 e879c8ed7864..2edf1add67f2 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 @@ -38,7 +38,7 @@ 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.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -80,7 +80,7 @@ public MailboxSendOperator(MailboxService mailboxService, public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, KeySelector keySelector, int senderStageId, int receiverStageId, - OperatorExecutionContext context) { + OpChainExecutionContext context) { this(context.getMailboxService(), dataTableBlockBaseOperator, context.getMetadataMap().get(receiverStageId).getServerInstances(), exchangeType, keySelector, context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getDeadlineMs()); 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 d3415ad26acd..172a54a7b318 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 @@ -28,7 +28,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.utils.OperatorUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.spi.exception.EarlyTerminationException; import org.apache.pinot.spi.trace.InvocationScope; import org.apache.pinot.spi.trace.Tracing; @@ -46,7 +46,7 @@ public abstract class MultiStageOperator implements Operator, protected final Map _operatorStatsMap; private final String _operatorId; - public MultiStageOperator(OperatorExecutionContext context) { + public MultiStageOperator(OpChainExecutionContext context) { this(context.getRequestId(), context.getStageId(), context.getServer()); } 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 4680a17b4c33..9d80e7801bda 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 @@ -24,7 +24,7 @@ import org.apache.pinot.core.common.Operator; import org.apache.pinot.query.mailbox.MailboxIdentifier; import org.apache.pinot.query.runtime.blocks.TransferableBlock; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; /** @@ -46,7 +46,7 @@ public OpChain(MultiStageOperator root, List receivingMailbox } public OpChain(MultiStageOperator root, List receivingMailboxes, int virtualServerId, - OperatorExecutionContext context) { + OpChainExecutionContext context) { this(root, receivingMailboxes, virtualServerId, context.getRequestId(), context.getStageId()); } 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 33853ae6d4d2..e5e45fa1fa23 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 @@ -31,10 +31,9 @@ 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.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -57,7 +56,7 @@ public class SortOperator extends MultiStageOperator { public SortOperator(MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, - OperatorExecutionContext context) { + OpChainExecutionContext context) { this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, context); } @@ -65,7 +64,7 @@ public SortOperator(MultiStageOperator upstreamOperator, List col @VisibleForTesting SortOperator(MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, - int defaultHolderCapacity, OperatorExecutionContext context) { + int defaultHolderCapacity, OpChainExecutionContext context) { super(context); _upstreamOperator = upstreamOperator; _fetch = fetch; 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 a3a78aea6a0c..0b20a70dbd47 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 @@ -30,7 +30,7 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.query.runtime.operator.operands.TransformOperand; import org.apache.pinot.query.runtime.operator.utils.FunctionInvokeUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -56,7 +56,7 @@ public class TransformOperator extends MultiStageOperator { private TransferableBlock _upstreamErrorBlock; public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, - DataSchema upstreamDataSchema, OperatorExecutionContext context) { + DataSchema upstreamDataSchema, OpChainExecutionContext context) { super(context); Preconditions.checkState(!transforms.isEmpty(), "transform operand should not be empty."); Preconditions.checkState(resultSchema.size() == transforms.size(), 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 3821fb065f8f..3599860fcc7f 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 @@ -37,11 +37,10 @@ import org.apache.pinot.core.data.table.Key; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.stage.WindowNode; -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.operator.utils.AggregationUtils; -import org.apache.pinot.query.runtime.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -93,7 +92,7 @@ public WindowAggregateOperator(MultiStageOperator inputOperator, List orderSet, List orderSetDirection, List orderSetNullDirection, List aggCalls, int lowerBound, int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, - DataSchema resultSchema, DataSchema inputSchema, OperatorExecutionContext context) { + DataSchema resultSchema, DataSchema inputSchema, OpChainExecutionContext context) { this(inputOperator, groupSet, orderSet, orderSetDirection, orderSetNullDirection, aggCalls, lowerBound, upperBound, windowFrameType, constants, resultSchema, inputSchema, AggregationUtils.Accumulator.MERGERS, context); @@ -105,7 +104,7 @@ public WindowAggregateOperator(MultiStageOperator inputOperator, List orderSetNullDirection, List aggCalls, int lowerBound, int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, DataSchema resultSchema, DataSchema inputSchema, Map> mergers, OperatorExecutionContext context) { + Function> mergers, OpChainExecutionContext context) { super(context); boolean isPartitionByOnly = isPartitionByOnlyQuery(groupSet, orderSet, orderSetDirection, orderSetNullDirection); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java similarity index 81% rename from pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java rename to pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java index 18ce9769ef7e..1f500ce8b28f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OperatorExecutionContext.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainExecutionContext.java @@ -25,17 +25,21 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock; -public class OperatorExecutionContext { +/** + * The {@code OpChainExecutionContext} class contains the information derived from the PlanRequestContext. + * Members of this class should not be changed once initialized. + * This information is then used by the OpChain to create the Operators for a query. + */ +public class OpChainExecutionContext { private final MailboxService _mailboxService; private final long _requestId; private final int _stageId; - private final VirtualServerAddress _server; private final long _timeoutMs; private final long _deadlineMs; - protected final Map _metadataMap; + private final Map _metadataMap; - public OperatorExecutionContext(MailboxService mailboxService, long requestId, int stageId, + public OpChainExecutionContext(MailboxService mailboxService, long requestId, int stageId, VirtualServerAddress server, long timeoutMs, long deadlineMs, Map metadataMap) { _mailboxService = mailboxService; _requestId = requestId; @@ -46,7 +50,7 @@ public OperatorExecutionContext(MailboxService mailboxService _metadataMap = metadataMap; } - public OperatorExecutionContext(PlanRequestContext planRequestContext) { + public OpChainExecutionContext(PlanRequestContext planRequestContext) { this(planRequestContext.getMailboxService(), planRequestContext.getRequestId(), planRequestContext.getStageId(), planRequestContext.getServer(), planRequestContext.getTimeoutMs(), planRequestContext.getDeadlineMs(), planRequestContext.getMetadataMap()); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java index 1098b35c23ff..53ff0ff658e8 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java @@ -57,14 +57,14 @@ public class PhysicalPlanVisitor implements StageNodeVisitor _metadataMap; protected final List _receivingMailboxes = new ArrayList<>(); - + private final OpChainExecutionContext _opChainExecutionContext; public PlanRequestContext(MailboxService mailboxService, long requestId, int stageId, long timeoutMs, long deadlineMs, VirtualServerAddress server, Map metadataMap) { @@ -50,6 +50,7 @@ public PlanRequestContext(MailboxService mailboxService, long _deadlineMs = deadlineMs; _server = server; _metadataMap = metadataMap; + _opChainExecutionContext = new OpChainExecutionContext(this); } public long getRequestId() { @@ -88,7 +89,7 @@ public List getReceivingMailboxes() { return ImmutableList.copyOf(_receivingMailboxes); } - public OperatorExecutionContext getOperatorExecutionContext() { - return new OperatorExecutionContext(this); + public OpChainExecutionContext getOpChainExecutionContext() { + return _opChainExecutionContext; } } 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 3b456f0b6765..cd1f3ffe8dd2 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 @@ -25,7 +25,7 @@ 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.plan.OperatorExecutionContext; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.query.testutils.MockDataBlockOperatorFactory; public class OperatorTestUtil { @@ -62,15 +62,15 @@ public static TransferableBlock block(DataSchema schema, Object[]... rows) { return new TransferableBlock(Arrays.asList(rows), schema, DataBlock.Type.ROW); } - public static OperatorExecutionContext getDefaultContext() { + public static OpChainExecutionContext getDefaultContext() { VirtualServerAddress virtualServerAddress = new VirtualServerAddress("mock", 80, 0); - return new OperatorExecutionContext(null, 1, 2, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, + return new OpChainExecutionContext(null, 1, 2, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); } - public static OperatorExecutionContext getContext(long requestId, int stageId, + public static OpChainExecutionContext getContext(long requestId, int stageId, VirtualServerAddress virtualServerAddress) { - return new OperatorExecutionContext(null, requestId, stageId, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, + return new OpChainExecutionContext(null, requestId, stageId, virtualServerAddress, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); } } From d7ca02622e97c34bb8d0fb68bf98c58089c480c2 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Tue, 14 Mar 2023 23:44:37 +0530 Subject: [PATCH 06/12] Refactor MailboxSendOperator for context --- .../pinot/query/runtime/QueryRunner.java | 13 ++-- .../LeafStageTransferableBlockOperator.java | 7 +-- .../runtime/operator/MailboxSendOperator.java | 32 ++++------ .../operator/MailboxSendOperatorTest.java | 61 ++++++++++++------- 4 files changed, 58 insertions(+), 55 deletions(-) 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 a25c2e0b54ff..2393a1d5dec8 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 @@ -54,6 +54,7 @@ import org.apache.pinot.query.runtime.operator.MailboxSendOperator; import org.apache.pinot.query.runtime.operator.OpChain; import org.apache.pinot.query.runtime.plan.DistributedStagePlan; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.query.runtime.plan.PhysicalPlanVisitor; import org.apache.pinot.query.runtime.plan.PlanRequestContext; import org.apache.pinot.query.runtime.plan.ServerRequestPlanVisitor; @@ -209,12 +210,12 @@ private void runLeafStage(DistributedStagePlan distributedStagePlan, Map baseResultBlock, DataSchema dataSchema, OpChainExecutionContext context) { - this(baseResultBlock, dataSchema, context.getRequestId(), context.getStageId(), context.getServer()); - } - - public LeafStageTransferableBlockOperator(List baseResultBlock, DataSchema dataSchema, - long requestId, int stageId, VirtualServerAddress serverAddress) { - super(requestId, stageId, serverAddress); + super(context); _baseResultBlock = baseResultBlock; _desiredDataSchema = dataSchema; _errorBlock = baseResultBlock.stream().filter(e -> !e.getExceptions().isEmpty()).findFirst().orElse(null); 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 2edf1add67f2..a92a73212937 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 @@ -69,32 +69,22 @@ interface MailboxIdGenerator { MailboxIdentifier generate(VirtualServer server); } - public MailboxSendOperator(MailboxService mailboxService, - MultiStageOperator dataTableBlockBaseOperator, List receivingStageInstances, - RelDistribution.Type exchangeType, KeySelector keySelector, - VirtualServerAddress sendingServer, long jobId, int senderStageId, int receiverStageId, long deadlineMs) { - this(mailboxService, dataTableBlockBaseOperator, receivingStageInstances, exchangeType, keySelector, - server -> toMailboxId(server, jobId, senderStageId, receiverStageId, sendingServer), BlockExchange::getExchange, - jobId, senderStageId, receiverStageId, sendingServer, deadlineMs); - } - public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, KeySelector keySelector, int senderStageId, int receiverStageId, OpChainExecutionContext context) { - this(context.getMailboxService(), dataTableBlockBaseOperator, - context.getMetadataMap().get(receiverStageId).getServerInstances(), exchangeType, keySelector, - context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getDeadlineMs()); + this(dataTableBlockBaseOperator, exchangeType, keySelector, (server) -> toMailboxId(server, context.getRequestId(), + senderStageId, receiverStageId, context.getServer()), BlockExchange::getExchange, receiverStageId, context); } @VisibleForTesting - MailboxSendOperator(MailboxService mailboxService, - MultiStageOperator dataTableBlockBaseOperator, List receivingStageInstances, - RelDistribution.Type exchangeType, KeySelector keySelector, - MailboxIdGenerator mailboxIdGenerator, BlockExchangeFactory blockExchangeFactory, long jobId, int senderStageId, - int receiverStageId, VirtualServerAddress serverAddress, long deadlineMs) { - super(jobId, senderStageId, serverAddress); + MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, + KeySelector keySelector, MailboxIdGenerator mailboxIdGenerator, + BlockExchangeFactory blockExchangeFactory, int receiverStageId, OpChainExecutionContext context) { + super(context); _dataTableBlockBaseOperator = dataTableBlockBaseOperator; - + MailboxService mailboxService = context.getMailboxService(); + List receivingStageInstances = + context.getMetadataMap().get(receiverStageId).getServerInstances(); List receivingMailboxes; if (exchangeType == RelDistribution.Type.SINGLETON) { // TODO: this logic should be moved into SingletonExchange @@ -121,8 +111,8 @@ public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDis } BlockSplitter splitter = TransferableBlockUtils::splitBlock; - _exchange = blockExchangeFactory.build(mailboxService, receivingMailboxes, exchangeType, keySelector, splitter, - deadlineMs); + _exchange = blockExchangeFactory.build(context.getMailboxService(), receivingMailboxes, exchangeType, keySelector, splitter, + context.getDeadlineMs()); Preconditions.checkState(SUPPORTED_EXCHANGE_TYPE.contains(exchangeType), String.format("Exchange type '%s' is not supported yet", exchangeType)); 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 ea7df4982996..d8184ba7e149 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 @@ -20,17 +20,21 @@ import com.google.common.collect.ImmutableList; import java.util.Arrays; +import java.util.Collections; +import java.util.Map; import org.apache.calcite.rel.RelDistribution; import org.apache.pinot.common.datablock.DataBlock; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.mailbox.JsonMailboxIdentifier; import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.planner.StageMetadata; import org.apache.pinot.query.planner.partitioning.KeySelector; import org.apache.pinot.query.routing.VirtualServer; 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.operator.exchange.BlockExchange; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.Mockito; @@ -65,8 +69,7 @@ public class MailboxSendOperatorTest { public void setUp() { _mocks = MockitoAnnotations.openMocks(this); Mockito.when(_exchangeFactory.build(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), - Mockito.anyLong())) - .thenReturn(_exchange); + Mockito.anyLong())).thenReturn(_exchange); Mockito.when(_server.getHostname()).thenReturn("mock"); Mockito.when(_server.getQueryMailboxPort()).thenReturn(0); @@ -84,11 +87,14 @@ public void shouldSwallowNoOpBlockFromUpstream() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: - MailboxSendOperator operator = new MailboxSendOperator(_mailboxService, _input, ImmutableList.of(_server), - RelDistribution.Type.HASH_DISTRIBUTED, _selector, +// OpChainExecutionContext context = +// OperatorTestUtil.getContext(1, DEFAULT_SENDER_STAGE_ID, new VirtualServerAddress(_server)); + OpChainExecutionContext context = getOpChainContext(deadlineMs); + + MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, 1, DEFAULT_SENDER_STAGE_ID, DEFAULT_RECEIVER_STAGE_ID, - new VirtualServerAddress(_server), deadlineMs); + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); // When: @@ -104,11 +110,11 @@ public void shouldSendErrorBlock() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: - MailboxSendOperator operator = new MailboxSendOperator(_mailboxService, _input, ImmutableList.of(_server), - RelDistribution.Type.HASH_DISTRIBUTED, _selector, + OpChainExecutionContext context = getOpChainContext(deadlineMs); + + MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, 1, DEFAULT_SENDER_STAGE_ID, DEFAULT_RECEIVER_STAGE_ID, - new VirtualServerAddress(_server), deadlineMs); + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock errorBlock = TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!")); Mockito.when(_input.nextBlock()).thenReturn(errorBlock); @@ -125,11 +131,11 @@ public void shouldSendErrorBlockWhenInputThrows() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: - MailboxSendOperator operator = new MailboxSendOperator(_mailboxService, _input, ImmutableList.of(_server), - RelDistribution.Type.HASH_DISTRIBUTED, _selector, + OpChainExecutionContext context = getOpChainContext(deadlineMs); + + MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, 1, DEFAULT_SENDER_STAGE_ID, DEFAULT_RECEIVER_STAGE_ID, - new VirtualServerAddress(_server), deadlineMs); + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); Mockito.when(_input.nextBlock()).thenThrow(new RuntimeException("foo!")); ArgumentCaptor captor = ArgumentCaptor.forClass(TransferableBlock.class); @@ -147,11 +153,12 @@ public void shouldSendEosBlock() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: - MailboxSendOperator operator = new MailboxSendOperator(_mailboxService, _input, ImmutableList.of(_server), - RelDistribution.Type.HASH_DISTRIBUTED, _selector, + OpChainExecutionContext context = getOpChainContext(deadlineMs); + + MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, 1, DEFAULT_SENDER_STAGE_ID, DEFAULT_RECEIVER_STAGE_ID, - new VirtualServerAddress(_server), deadlineMs); + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock(); Mockito.when(_input.nextBlock()).thenReturn(eosBlock); @@ -168,11 +175,11 @@ public void shouldSendDataBlock() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: - MailboxSendOperator operator = new MailboxSendOperator(_mailboxService, _input, ImmutableList.of(_server), - RelDistribution.Type.HASH_DISTRIBUTED, _selector, + OpChainExecutionContext context = getOpChainContext(deadlineMs); + + MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, 1, DEFAULT_SENDER_STAGE_ID, DEFAULT_RECEIVER_STAGE_ID, - new VirtualServerAddress(_server), deadlineMs); + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock dataBlock = block(new DataSchema(new String[]{}, new DataSchema.ColumnDataType[]{})); Mockito.when(_input.nextBlock()).thenReturn(dataBlock) .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); @@ -189,4 +196,14 @@ public void shouldSendDataBlock() private static TransferableBlock block(DataSchema schema, Object[]... rows) { return new TransferableBlock(Arrays.asList(rows), schema, DataBlock.Type.ROW); } + + private OpChainExecutionContext getOpChainContext(long deadlineMs) { + StageMetadata stageMetadata = new StageMetadata(); + stageMetadata.setServerInstances(Collections.singletonList(_server)); + Map stageMetadataMap = Collections.singletonMap(DEFAULT_RECEIVER_STAGE_ID, stageMetadata); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, 1, DEFAULT_SENDER_STAGE_ID, new VirtualServerAddress(_server), + deadlineMs, deadlineMs, stageMetadataMap); + return context; + } } From 404144357fd03ca2761d6bf6efe95951b263d436 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Wed, 15 Mar 2023 00:50:39 +0530 Subject: [PATCH 07/12] Remove redundant constructors and use context based constructor for MailboxReceiveOperator --- .../operator/MailboxReceiveOperator.java | 17 +- .../service/dispatch/QueryDispatcher.java | 15 +- .../operator/MailboxReceiveOperatorTest.java | 227 ++++++++++-------- 3 files changed, 150 insertions(+), 109 deletions(-) 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 15d8b06df59c..6268fe3dc604 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 @@ -79,19 +79,20 @@ private static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, i public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, OpChainExecutionContext context) { - this(context.getMailboxService(), context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, - context.getServer(), context.getRequestId(), senderStageId, receiverStageId, context.getTimeoutMs()); + this(context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId, receiverStageId, context); } // TODO: Move deadlineInNanoSeconds to OperatorContext. - public MailboxReceiveOperator(MailboxService mailboxService, - List sendingStageInstances, RelDistribution.Type exchangeType, VirtualServerAddress receiver, - long jobId, int senderStageId, int receiverStageId, Long timeoutMs) { - super(jobId, senderStageId, receiver); - _mailboxService = mailboxService; + public MailboxReceiveOperator(List sendingStageInstances, + RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, OpChainExecutionContext context) { + super(context); + _mailboxService = context.getMailboxService(); + VirtualServerAddress receiver = context.getServer(); + long jobId = context.getRequestId(); Preconditions.checkState(SUPPORTED_EXCHANGE_TYPES.contains(exchangeType), "Exchange/Distribution type: " + exchangeType + " is not supported!"); - long timeoutNano = (timeoutMs != null ? timeoutMs : QueryConfig.DEFAULT_MAILBOX_TIMEOUT_MS) * 1_000_000L; + //TODO: Should I be using the Long.MAX_VALUE here instead of -1 or previous null check? + long timeoutNano = (context.getTimeoutMs() != Long.MAX_VALUE ? context.getTimeoutMs() : QueryConfig.DEFAULT_MAILBOX_TIMEOUT_MS) * 1_000_000L; _deadlineTimestampNano = timeoutNano + System.nanoTime(); _exchangeType = exchangeType; 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 65aab54df073..219f4078a1a6 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 @@ -57,6 +57,7 @@ import org.apache.pinot.query.runtime.operator.OperatorStats; import org.apache.pinot.query.runtime.operator.utils.OperatorUtils; import org.apache.pinot.query.runtime.plan.DistributedStagePlan; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.query.runtime.plan.serde.QueryPlanSerDeUtils; import org.apache.pinot.query.service.QueryConfig; import org.roaringbitmap.RoaringBitmap; @@ -173,10 +174,10 @@ int submit(long requestId, QueryPlan queryPlan, long timeoutMs, Map mailboxService, Map statsAggregatorMap) { MailboxReceiveNode reduceNode = (MailboxReceiveNode) queryPlan.getQueryStageMap().get(reduceStageId); - MailboxReceiveOperator mailboxReceiveOperator = createReduceStageOperator(mailboxService, - queryPlan.getStageMetadataMap().get(reduceNode.getSenderStageId()).getServerInstances(), requestId, - reduceNode.getSenderStageId(), reduceStageId, reduceNode.getDataSchema(), - new VirtualServerAddress(mailboxService.getHostname(), mailboxService.getMailboxPort(), 0), timeoutMs); + MailboxReceiveOperator mailboxReceiveOperator = + createReduceStageOperator(mailboxService, queryPlan.getStageMetadataMap(), requestId, + reduceNode.getSenderStageId(), reduceStageId, reduceNode.getDataSchema(), + new VirtualServerAddress(mailboxService.getHostname(), mailboxService.getMailboxPort(), 0), timeoutMs); List resultDataBlocks = reduceMailboxReceive(mailboxReceiveOperator, timeoutMs, statsAggregatorMap, queryPlan); return toResultTable(resultDataBlocks, queryPlan.getQueryResultFields(), @@ -275,12 +276,12 @@ private static DataSchema toResultSchema(DataSchema inputSchema, List mailboxService, - List sendingInstances, long jobId, int stageId, int reducerStageId, DataSchema dataSchema, + Map stageMetadataMap, long jobId, int stageId, int reducerStageId, DataSchema dataSchema, VirtualServerAddress server, long timeoutMs) { + OpChainExecutionContext context = new OpChainExecutionContext(mailboxService, jobId, stageId, server, timeoutMs, timeoutMs, stageMetadataMap); // timeout is set for reduce stage MailboxReceiveOperator mailboxReceiveOperator = - new MailboxReceiveOperator(mailboxService, sendingInstances, - RelDistribution.Type.RANDOM_DISTRIBUTED, server, jobId, stageId, reducerStageId, timeoutMs); + new MailboxReceiveOperator(RelDistribution.Type.RANDOM_DISTRIBUTED, stageId, reducerStageId, context); return mailboxReceiveOperator; } 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 d325612ff140..ef2826c7f71d 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 @@ -20,6 +20,7 @@ import com.google.common.collect.ImmutableList; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import org.apache.calcite.rel.RelDistribution; import org.apache.pinot.common.datablock.MetadataBlock; @@ -32,6 +33,7 @@ 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.plan.OpChainExecutionContext; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.MockitoAnnotations; @@ -79,9 +81,11 @@ public void tearDown() public void shouldTimeoutOnExtraLongSleep() throws InterruptedException { // shorter timeoutMs should result in error. + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 10L, 10L, + new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(_mailboxService, new ArrayList<>(), RelDistribution.Type.SINGLETON, _testAddr, 456, - 789, DEFAULT_RECEIVER_STAGE_ID, 10L); + new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); Thread.sleep(200L); TransferableBlock mailbox = receiveOp.nextBlock(); Assert.assertTrue(mailbox.isErrorBlock()); @@ -89,15 +93,15 @@ public void shouldTimeoutOnExtraLongSleep() Assert.assertTrue(errorBlock.getExceptions().containsKey(QueryException.EXECUTION_TIMEOUT_ERROR_CODE)); // longer timeout or default timeout (10s) doesn't result in error. - receiveOp = - new MailboxReceiveOperator(_mailboxService, new ArrayList<>(), RelDistribution.Type.SINGLETON, _testAddr, 456, - 789, DEFAULT_RECEIVER_STAGE_ID, 2000L); + context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 2000L, 2000L, + new HashMap<>()); + receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); - receiveOp = - new MailboxReceiveOperator(_mailboxService, new ArrayList<>(), RelDistribution.Type.SINGLETON, _testAddr, 456, - 789, DEFAULT_RECEIVER_STAGE_ID, null); + context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); @@ -116,8 +120,12 @@ public void shouldThrowReceiveSingletonFromMultiMatchMailboxServer() { Mockito.when(_server2.getHostname()).thenReturn("singleton"); Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, _testAddr, 456, 789, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, 456, 789, + context); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*RANGE_DISTRIBUTED.*") @@ -131,8 +139,12 @@ public void shouldThrowRangeDistributionNotSupported() { Mockito.when(_server2.getHostname()).thenReturn("singleton"); Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.RANGE_DISTRIBUTED, _testAddr, 456, 789, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.RANGE_DISTRIBUTED, 456, + 789, context); } @Test @@ -156,8 +168,13 @@ public void shouldReceiveSingletonNoMatchMailboxServer() { String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); // Receive end of stream block directly when there is no match. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); @@ -184,14 +201,19 @@ public void shouldReceiveSingletonCloseMailbox() { String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(serverHost, server2port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(serverHost, server2port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(true); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); + // Receive end of stream block directly when mailbox is close. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); } @@ -218,16 +240,19 @@ public void shouldReceiveSingletonNullMailbox() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(serverHost, server2port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(serverHost, server2port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); // Receive null mailbox during timeout. Mockito.when(_mailbox.receive()).thenReturn(null); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); // Receive NoOpBlock. Assert.assertTrue(receiveOp.nextBlock().isNoOpBlock()); } @@ -254,15 +279,18 @@ public void shouldReceiveEosDirectlyFromSender() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(serverHost, server2port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(serverHost, server2port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); // Receive EosBloc. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); } @@ -289,17 +317,20 @@ public void shouldReceiveSingletonMailbox() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(serverHost, server2port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(serverHost, server2port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Object[] expRow = new Object[]{1, 1}; DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); Mockito.when(_mailbox.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -328,16 +359,19 @@ public void shouldReceiveSingletonErrorMailbox() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(serverHost, server2port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(serverHost, server2port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Exception e = new Exception("errorBlock"); Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getErrorTransferableBlock(e)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.SINGLETON, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock(); @@ -363,24 +397,25 @@ public void shouldReceiveMailboxFromTwoServersOneClose() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId1 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server1Host, server1Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server1Host, server1Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(true); - JsonMailboxIdentifier expectedMailboxId2 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server2Host, server2Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server2Host, server2Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2); Mockito.when(_mailbox2.isClosed()).thenReturn(false); Object[] expRow = new Object[]{1, 1}; DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.HASH_DISTRIBUTED, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -406,25 +441,26 @@ public void shouldReceiveMailboxFromTwoServersOneNull() String toHost = "toHost"; VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); - JsonMailboxIdentifier expectedMailboxId1 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server1Host, server1Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server1Host, server1Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Mockito.when(_mailbox.receive()).thenReturn(null); - JsonMailboxIdentifier expectedMailboxId2 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server2Host, server2Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server2Host, server2Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2); Mockito.when(_mailbox2.isClosed()).thenReturn(false); Object[] expRow = new Object[]{1, 1}; DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.HASH_DISTRIBUTED, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -451,10 +487,8 @@ public void shouldReceiveMailboxFromTwoServers() VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); - JsonMailboxIdentifier expectedMailboxId1 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server1Host, server1Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server1Host, server1Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Object[] expRow1 = new Object[]{1, 1}; @@ -464,15 +498,18 @@ public void shouldReceiveMailboxFromTwoServers() TransferableBlockUtils.getEndOfStreamTransferableBlock()); Object[] expRow3 = new Object[]{3, 3}; - JsonMailboxIdentifier expectedMailboxId2 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server2Host, server2Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server2Host, server2Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2); Mockito.when(_mailbox2.isClosed()).thenReturn(false); Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.HASH_DISTRIBUTED, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); // Receive first block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); @@ -511,25 +548,26 @@ public void shouldGetReceptionReceiveErrorMailbox() VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); - JsonMailboxIdentifier expectedMailboxId1 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server1Host, server1Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server1Host, server1Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Mockito.when(_mailbox.receive()) .thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("mailboxError"))); Object[] expRow3 = new Object[]{3, 3}; - JsonMailboxIdentifier expectedMailboxId2 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server2Host, server2Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server2Host, server2Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2); Mockito.when(_mailbox2.isClosed()).thenReturn(false); Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.HASH_DISTRIBUTED, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); // Receive error block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); @@ -557,24 +595,25 @@ public void shouldThrowReceiveWhenOneServerReceiveThrowException() VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0); DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT}); - JsonMailboxIdentifier expectedMailboxId1 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server1Host, server1Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server1Host, server1Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox); Mockito.when(_mailbox.isClosed()).thenReturn(false); Mockito.when(_mailbox.receive()).thenThrow(new Exception("mailboxError")); Object[] expRow3 = new Object[]{3, 3}; - JsonMailboxIdentifier expectedMailboxId2 = - new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), - new VirtualServerAddress(server2Host, server2Port, 0), - toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); + JsonMailboxIdentifier expectedMailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", jobId, stageId), + new VirtualServerAddress(server2Host, server2Port, 0), toAddress, stageId, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2); Mockito.when(_mailbox2.isClosed()).thenReturn(false); Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3)); - MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(_mailboxService, ImmutableList.of(_server1, _server2), - RelDistribution.Type.HASH_DISTRIBUTED, toAddress, jobId, stageId, DEFAULT_RECEIVER_STAGE_ID, null); + OpChainExecutionContext context = + new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE, + Long.MAX_VALUE, new HashMap<>()); + + MailboxReceiveOperator receiveOp = + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, + DEFAULT_RECEIVER_STAGE_ID, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock(), "server-1 should have returned an error-block"); } From ad2e3d5a6fb5ba057bb8ccb6b5a1e522879f310b Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Wed, 15 Mar 2023 00:58:49 +0530 Subject: [PATCH 08/12] Remove no context constructor from MultiStageOperator --- .../runtime/operator/MultiStageOperator.java | 18 ++++++------------ .../query/runtime/operator/OperatorStats.java | 6 ++++++ 2 files changed, 12 insertions(+), 12 deletions(-) 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 172a54a7b318..fcf083354e60 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 @@ -39,24 +39,18 @@ public abstract class MultiStageOperator implements Operator, private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(MultiStageOperator.class); // TODO: Move to OperatorContext class. - protected final long _requestId; - protected final int _stageId; - protected final VirtualServerAddress _serverAddress; protected final OperatorStats _operatorStats; protected final Map _operatorStatsMap; private final String _operatorId; + private final OpChainExecutionContext _context; public MultiStageOperator(OpChainExecutionContext context) { - this(context.getRequestId(), context.getStageId(), context.getServer()); - } - - public MultiStageOperator(long requestId, int stageId, VirtualServerAddress serverAddress) { - _requestId = requestId; - _stageId = stageId; - _operatorStats = new OperatorStats(requestId, stageId, serverAddress, toExplainString()); - _serverAddress = serverAddress; + _context = context; + _operatorStats = + new OperatorStats(_context, toExplainString()); _operatorStatsMap = new HashMap<>(); - _operatorId = Joiner.on("_").join(toExplainString(), _requestId, _stageId, _serverAddress); + _operatorId = + Joiner.on("_").join(toExplainString(), _context.getRequestId(), _context.getStageId(), _context.getServer()); } public Map getOperatorStatsMap() { 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 index 715be163bb85..cb01623ee490 100644 --- 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 @@ -25,6 +25,7 @@ import org.apache.pinot.common.datatable.DataTable; import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.operator.utils.OperatorUtils; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; public class OperatorStats { @@ -42,6 +43,11 @@ public class OperatorStats { private long _startTimeMs = -1; private final Map _executionStats; + public OperatorStats(OpChainExecutionContext context, String operatorType) { + this(context.getRequestId(), context.getStageId(), context.getServer(), operatorType); + } + + //TODO: remove this constructor after the context constructor can be used in serialization and deserialization public OperatorStats(long requestId, int stageId, VirtualServerAddress serverAddress, String operatorType) { _stageId = stageId; _requestId = requestId; From dcdd8a0a164ee2887d3500965c79159f51267e97 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Wed, 15 Mar 2023 13:16:56 +0530 Subject: [PATCH 09/12] Use boxed timeout value in mailbox receive --- .../operator/MailboxReceiveOperator.java | 13 ++++---- .../service/dispatch/QueryDispatcher.java | 3 +- .../operator/MailboxReceiveOperatorTest.java | 32 +++++++++---------- 3 files changed, 25 insertions(+), 23 deletions(-) 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 6268fe3dc604..1622834d266e 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 @@ -77,22 +77,23 @@ private static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, i receiverStageId); } - public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, - int receiverStageId, OpChainExecutionContext context) { - this(context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId, receiverStageId, context); + public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, + OpChainExecutionContext context) { + this(context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId, receiverStageId, + context.getTimeoutMs(), context); } // TODO: Move deadlineInNanoSeconds to OperatorContext. + //TODO: Remove boxed timeoutMs value from here and use long deadlineMs from context. public MailboxReceiveOperator(List sendingStageInstances, - RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, OpChainExecutionContext context) { + RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, Long timeoutMs, OpChainExecutionContext context) { super(context); _mailboxService = context.getMailboxService(); VirtualServerAddress receiver = context.getServer(); long jobId = context.getRequestId(); Preconditions.checkState(SUPPORTED_EXCHANGE_TYPES.contains(exchangeType), "Exchange/Distribution type: " + exchangeType + " is not supported!"); - //TODO: Should I be using the Long.MAX_VALUE here instead of -1 or previous null check? - long timeoutNano = (context.getTimeoutMs() != Long.MAX_VALUE ? context.getTimeoutMs() : QueryConfig.DEFAULT_MAILBOX_TIMEOUT_MS) * 1_000_000L; + long timeoutNano = (timeoutMs != null ? timeoutMs : QueryConfig.DEFAULT_MAILBOX_TIMEOUT_MS) * 1_000_000L; _deadlineTimestampNano = timeoutNano + System.nanoTime(); _exchangeType = exchangeType; 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 219f4078a1a6..0cf14a29f6fe 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 @@ -278,7 +278,8 @@ private static DataSchema toResultSchema(DataSchema inputSchema, List mailboxService, Map stageMetadataMap, long jobId, int stageId, int reducerStageId, DataSchema dataSchema, VirtualServerAddress server, long timeoutMs) { - OpChainExecutionContext context = new OpChainExecutionContext(mailboxService, jobId, stageId, server, timeoutMs, timeoutMs, stageMetadataMap); + OpChainExecutionContext context = + new OpChainExecutionContext(mailboxService, jobId, stageId, server, timeoutMs, timeoutMs, stageMetadataMap); // timeout is set for reduce stage MailboxReceiveOperator mailboxReceiveOperator = new MailboxReceiveOperator(RelDistribution.Type.RANDOM_DISTRIBUTED, stageId, reducerStageId, context); 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 ef2826c7f71d..953a73395888 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 @@ -85,7 +85,7 @@ public void shouldTimeoutOnExtraLongSleep() new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 10L, 10L, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); + new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 10L, context); Thread.sleep(200L); TransferableBlock mailbox = receiveOp.nextBlock(); Assert.assertTrue(mailbox.isErrorBlock()); @@ -95,13 +95,13 @@ public void shouldTimeoutOnExtraLongSleep() // longer timeout or default timeout (10s) doesn't result in error. context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 2000L, 2000L, new HashMap<>()); - receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); + receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 2000L, context); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); - receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, context); + receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, null, context); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); @@ -124,7 +124,7 @@ public void shouldThrowReceiveSingletonFromMultiMatchMailboxServer() { new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, 456, 789, + new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, 456, 789, null, context); } @@ -144,7 +144,7 @@ public void shouldThrowRangeDistributionNotSupported() { Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.RANGE_DISTRIBUTED, 456, - 789, context); + 789,null, context); } @Test @@ -174,7 +174,7 @@ public void shouldReceiveSingletonNoMatchMailboxServer() { MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive end of stream block directly when there is no match. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); @@ -212,7 +212,7 @@ public void shouldReceiveSingletonCloseMailbox() { MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive end of stream block directly when mailbox is close. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); @@ -252,7 +252,7 @@ public void shouldReceiveSingletonNullMailbox() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive NoOpBlock. Assert.assertTrue(receiveOp.nextBlock().isNoOpBlock()); } @@ -290,7 +290,7 @@ public void shouldReceiveEosDirectlyFromSender() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive EosBloc. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); } @@ -330,7 +330,7 @@ public void shouldReceiveSingletonMailbox() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -371,7 +371,7 @@ public void shouldReceiveSingletonErrorMailbox() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock(); @@ -415,7 +415,7 @@ public void shouldReceiveMailboxFromTwoServersOneClose() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -460,7 +460,7 @@ public void shouldReceiveMailboxFromTwoServersOneNull() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -509,7 +509,7 @@ public void shouldReceiveMailboxFromTwoServers() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive first block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); @@ -567,7 +567,7 @@ public void shouldGetReceptionReceiveErrorMailbox() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); // Receive error block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); @@ -613,7 +613,7 @@ public void shouldThrowReceiveWhenOneServerReceiveThrowException() MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, context); + DEFAULT_RECEIVER_STAGE_ID, null, context); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock(), "server-1 should have returned an error-block"); } From f87d99149ec672892070fa9c9dfb9f0f3527eb3d Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Wed, 15 Mar 2023 13:28:12 +0530 Subject: [PATCH 10/12] Fix linting --- .../java/org/apache/pinot/query/runtime/QueryRunner.java | 5 +++-- .../pinot/query/runtime/operator/AggregateOperator.java | 3 ++- .../runtime/operator/LeafStageTransferableBlockOperator.java | 1 - .../pinot/query/runtime/operator/MailboxReceiveOperator.java | 4 ++-- .../pinot/query/runtime/operator/MailboxSendOperator.java | 5 +++-- .../pinot/query/runtime/operator/MultiStageOperator.java | 1 - .../query/runtime/operator/MailboxReceiveOperatorTest.java | 2 +- .../query/runtime/operator/MailboxSendOperatorTest.java | 2 +- 8 files changed, 12 insertions(+), 11 deletions(-) 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 2393a1d5dec8..81a1ee0881ac 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 @@ -210,8 +210,9 @@ private void runLeafStage(DistributedStagePlan distributedStagePlan, Map aggCalls, List groupSet, DataSchema inputSchema, - Map> mergers, OpChainExecutionContext context) { + Map> mergers, + OpChainExecutionContext context) { super(context); _inputOperator = inputOperator; _groupSet = groupSet; 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 0c8ba803cb34..dafae3cf2103 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 @@ -39,7 +39,6 @@ import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; import org.apache.pinot.core.query.selection.SelectionOperatorUtils; -import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.TransferableBlock; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.slf4j.Logger; 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 1622834d266e..27ea3daff54b 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 @@ -85,8 +85,8 @@ public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStage // TODO: Move deadlineInNanoSeconds to OperatorContext. //TODO: Remove boxed timeoutMs value from here and use long deadlineMs from context. - public MailboxReceiveOperator(List sendingStageInstances, - RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, Long timeoutMs, OpChainExecutionContext context) { + public MailboxReceiveOperator(List sendingStageInstances, RelDistribution.Type exchangeType, + int senderStageId, int receiverStageId, Long timeoutMs, OpChainExecutionContext context) { super(context); _mailboxService = context.getMailboxService(); VirtualServerAddress receiver = context.getServer(); 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 a92a73212937..decfb225ad40 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 @@ -111,8 +111,9 @@ public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDis } BlockSplitter splitter = TransferableBlockUtils::splitBlock; - _exchange = blockExchangeFactory.build(context.getMailboxService(), receivingMailboxes, exchangeType, keySelector, splitter, - context.getDeadlineMs()); + _exchange = + blockExchangeFactory.build(context.getMailboxService(), receivingMailboxes, exchangeType, keySelector, splitter, + context.getDeadlineMs()); Preconditions.checkState(SUPPORTED_EXCHANGE_TYPE.contains(exchangeType), String.format("Exchange type '%s' is not supported yet", exchangeType)); 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 fcf083354e60..8f97f57d4625 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 @@ -24,7 +24,6 @@ import java.util.Map; import org.apache.pinot.common.datatable.DataTable; import org.apache.pinot.core.common.Operator; -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.operator.utils.OperatorUtils; 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 953a73395888..bece567f160d 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 @@ -144,7 +144,7 @@ public void shouldThrowRangeDistributionNotSupported() { Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.RANGE_DISTRIBUTED, 456, - 789,null, context); + 789, null, context); } @Test 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 d8184ba7e149..ac57ef5af521 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 @@ -199,7 +199,7 @@ private static TransferableBlock block(DataSchema schema, Object[]... rows) { private OpChainExecutionContext getOpChainContext(long deadlineMs) { StageMetadata stageMetadata = new StageMetadata(); - stageMetadata.setServerInstances(Collections.singletonList(_server)); + stageMetadata.setServerInstances(ImmutableList.of(_server)); Map stageMetadataMap = Collections.singletonMap(DEFAULT_RECEIVER_STAGE_ID, stageMetadata); OpChainExecutionContext context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_SENDER_STAGE_ID, new VirtualServerAddress(_server), From d76b01143483ebefc3a5ffcbd0581f187482e421 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Thu, 16 Mar 2023 14:54:09 +0530 Subject: [PATCH 11/12] Make context as first argument --- .../pinot/query/runtime/QueryRunner.java | 6 +- .../runtime/operator/AggregateOperator.java | 15 +-- .../runtime/operator/FilterOperator.java | 4 +- .../runtime/operator/HashJoinOperator.java | 4 +- .../LeafStageTransferableBlockOperator.java | 4 +- .../operator/LiteralValueOperator.java | 4 +- .../operator/MailboxReceiveOperator.java | 12 +- .../runtime/operator/MailboxSendOperator.java | 17 +-- .../pinot/query/runtime/operator/OpChain.java | 5 +- .../query/runtime/operator/SortOperator.java | 14 +- .../runtime/operator/TransformOperator.java | 4 +- .../operator/WindowAggregateOperator.java | 19 ++- .../runtime/plan/PhysicalPlanVisitor.java | 41 +++--- .../service/dispatch/QueryDispatcher.java | 2 +- .../operator/AggregateOperatorTest.java | 25 ++-- .../runtime/operator/FilterOperatorTest.java | 28 ++-- .../operator/HashJoinOperatorTest.java | 32 ++--- ...eafStageTransferableBlockOperatorTest.java | 26 ++-- .../operator/LiteralValueOperatorTest.java | 4 +- .../operator/MailboxReceiveOperatorTest.java | 59 ++++---- .../operator/MailboxSendOperatorTest.java | 35 ++--- .../runtime/operator/SortOperatorTest.java | 30 ++--- .../operator/TransformOperatorTest.java | 33 ++--- .../operator/WindowAggregateOperatorTest.java | 126 +++++++++--------- 24 files changed, 275 insertions(+), 274 deletions(-) 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 81a1ee0881ac..9436ec4bf08f 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 @@ -213,10 +213,10 @@ private void runLeafStage(DistributedStagePlan distributedStagePlan, Map aggCalls, - List groupSet, DataSchema inputSchema, OpChainExecutionContext context) { - this(inputOperator, dataSchema, aggCalls, groupSet, inputSchema, AggregateOperator.AggregateAccumulator.AGG_MERGERS, - context); + public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, DataSchema dataSchema, + List aggCalls, List groupSet, DataSchema inputSchema) { + this(context, inputOperator, dataSchema, aggCalls, groupSet, inputSchema, + AggregateOperator.AggregateAccumulator.AGG_MERGERS); } @VisibleForTesting - AggregateOperator(MultiStageOperator inputOperator, DataSchema dataSchema, List aggCalls, - List groupSet, DataSchema inputSchema, - Map> mergers, - OpChainExecutionContext context) { + AggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, DataSchema dataSchema, + List aggCalls, List groupSet, DataSchema inputSchema, + Map> mergers) { super(context); _inputOperator = inputOperator; _groupSet = groupSet; 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 35c7915e8b06..2851aa528af7 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 @@ -55,8 +55,8 @@ public class FilterOperator extends MultiStageOperator { private final DataSchema _dataSchema; private TransferableBlock _upstreamErrorBlock; - public FilterOperator(MultiStageOperator upstreamOperator, DataSchema dataSchema, RexExpression filter, - OpChainExecutionContext context) { + public FilterOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, DataSchema dataSchema, + RexExpression filter) { super(context); _upstreamOperator = upstreamOperator; _dataSchema = dataSchema; 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 97a03ae25c08..9bc8c28407b6 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 @@ -88,8 +88,8 @@ public class HashJoinOperator extends MultiStageOperator { private KeySelector _leftKeySelector; private KeySelector _rightKeySelector; - public HashJoinOperator(MultiStageOperator leftTableOperator, MultiStageOperator rightTableOperator, - DataSchema leftSchema, JoinNode node, OpChainExecutionContext context) { + public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftTableOperator, + MultiStageOperator rightTableOperator, DataSchema leftSchema, JoinNode node) { super(context); Preconditions.checkState(SUPPORTED_JOIN_TYPES.contains(node.getJoinRelType()), "Join type: " + node.getJoinRelType() + " is not supported!"); 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 dafae3cf2103..9af6e8819d7c 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 @@ -67,8 +67,8 @@ public class LeafStageTransferableBlockOperator extends MultiStageOperator { private final DataSchema _desiredDataSchema; private int _currentIndex; - public LeafStageTransferableBlockOperator(List baseResultBlock, DataSchema dataSchema, - OpChainExecutionContext context) { + public LeafStageTransferableBlockOperator(OpChainExecutionContext context, + List baseResultBlock, DataSchema dataSchema) { super(context); _baseResultBlock = baseResultBlock; _desiredDataSchema = dataSchema; 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 af81e1e31d17..029e58e9a2a5 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 @@ -40,8 +40,8 @@ public class LiteralValueOperator extends MultiStageOperator { private final TransferableBlock _rexLiteralBlock; private boolean _isLiteralBlockReturned; - public LiteralValueOperator(DataSchema dataSchema, List> rexLiteralRows, - OpChainExecutionContext context) { + public LiteralValueOperator(OpChainExecutionContext context, DataSchema dataSchema, + List> rexLiteralRows) { super(context); _dataSchema = dataSchema; _rexLiteralBlock = constructBlock(rexLiteralRows); 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 27ea3daff54b..520392985e2c 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 @@ -77,16 +77,16 @@ private static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, i receiverStageId); } - public MailboxReceiveOperator(RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, - OpChainExecutionContext context) { - this(context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId, receiverStageId, - context.getTimeoutMs(), context); + public MailboxReceiveOperator(OpChainExecutionContext context, RelDistribution.Type exchangeType, int senderStageId, + int receiverStageId) { + this(context, context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId, + receiverStageId, context.getTimeoutMs()); } // TODO: Move deadlineInNanoSeconds to OperatorContext. //TODO: Remove boxed timeoutMs value from here and use long deadlineMs from context. - public MailboxReceiveOperator(List sendingStageInstances, RelDistribution.Type exchangeType, - int senderStageId, int receiverStageId, Long timeoutMs, OpChainExecutionContext context) { + public MailboxReceiveOperator(OpChainExecutionContext context, List sendingStageInstances, + RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, Long timeoutMs) { super(context); _mailboxService = context.getMailboxService(); VirtualServerAddress receiver = context.getServer(); 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 decfb225ad40..940e668cdc7a 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 @@ -69,17 +69,18 @@ interface MailboxIdGenerator { MailboxIdentifier generate(VirtualServer server); } - public MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, - KeySelector keySelector, int senderStageId, int receiverStageId, - OpChainExecutionContext context) { - this(dataTableBlockBaseOperator, exchangeType, keySelector, (server) -> toMailboxId(server, context.getRequestId(), - senderStageId, receiverStageId, context.getServer()), BlockExchange::getExchange, receiverStageId, context); + public MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator dataTableBlockBaseOperator, + RelDistribution.Type exchangeType, KeySelector keySelector, int senderStageId, + int receiverStageId) { + this(context, dataTableBlockBaseOperator, exchangeType, keySelector, + (server) -> toMailboxId(server, context.getRequestId(), senderStageId, receiverStageId, context.getServer()), + BlockExchange::getExchange, receiverStageId); } @VisibleForTesting - MailboxSendOperator(MultiStageOperator dataTableBlockBaseOperator, RelDistribution.Type exchangeType, - KeySelector keySelector, MailboxIdGenerator mailboxIdGenerator, - BlockExchangeFactory blockExchangeFactory, int receiverStageId, OpChainExecutionContext context) { + MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator dataTableBlockBaseOperator, + RelDistribution.Type exchangeType, KeySelector keySelector, + MailboxIdGenerator mailboxIdGenerator, BlockExchangeFactory blockExchangeFactory, int receiverStageId) { super(context); _dataTableBlockBaseOperator = dataTableBlockBaseOperator; MailboxService mailboxService = context.getMailboxService(); 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 9d80e7801bda..84fac943c20a 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 @@ -45,9 +45,8 @@ public OpChain(MultiStageOperator root, List receivingMailbox _stats = new OpChainStats(_id.toString()); } - public OpChain(MultiStageOperator root, List receivingMailboxes, int virtualServerId, - OpChainExecutionContext context) { - this(root, receivingMailboxes, virtualServerId, context.getRequestId(), context.getStageId()); + public OpChain(OpChainExecutionContext context, MultiStageOperator root, List receivingMailboxes) { + this(root, receivingMailboxes, context.getServer().virtualId(), context.getRequestId(), context.getStageId()); } public Operator getRoot() { 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 e5e45fa1fa23..9eeea863a9ea 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 @@ -54,17 +54,17 @@ public class SortOperator extends MultiStageOperator { private boolean _isSortedBlockConstructed; private TransferableBlock _upstreamErrorBlock; - public SortOperator(MultiStageOperator upstreamOperator, List collationKeys, - List collationDirections, int fetch, int offset, DataSchema dataSchema, - OpChainExecutionContext context) { - this(upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, - SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY, context); + public SortOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, + List collationKeys, List collationDirections, int fetch, int offset, + DataSchema dataSchema) { + this(context, upstreamOperator, collationKeys, collationDirections, fetch, offset, dataSchema, + SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY); } @VisibleForTesting - SortOperator(MultiStageOperator upstreamOperator, List collationKeys, + SortOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, List collationKeys, List collationDirections, int fetch, int offset, DataSchema dataSchema, - int defaultHolderCapacity, OpChainExecutionContext context) { + int defaultHolderCapacity) { super(context); _upstreamOperator = upstreamOperator; _fetch = fetch; 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 0b20a70dbd47..e53bcfdac232 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 @@ -55,8 +55,8 @@ public class TransformOperator extends MultiStageOperator { private final DataSchema _resultSchema; private TransferableBlock _upstreamErrorBlock; - public TransformOperator(MultiStageOperator upstreamOperator, DataSchema resultSchema, List transforms, - DataSchema upstreamDataSchema, OpChainExecutionContext context) { + public TransformOperator(OpChainExecutionContext context, MultiStageOperator upstreamOperator, + DataSchema resultSchema, List transforms, DataSchema upstreamDataSchema) { super(context); Preconditions.checkState(!transforms.isEmpty(), "transform operand should not be empty."); Preconditions.checkState(resultSchema.size() == transforms.size(), 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 3599860fcc7f..0d30dd394d27 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 @@ -88,23 +88,22 @@ public class WindowAggregateOperator extends MultiStageOperator { private boolean _readyToConstruct; private boolean _hasReturnedWindowAggregateBlock; - public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, - List orderSet, List orderSetDirection, + public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, + List groupSet, List orderSet, List orderSetDirection, List orderSetNullDirection, List aggCalls, int lowerBound, int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, - DataSchema resultSchema, DataSchema inputSchema, OpChainExecutionContext context) { - this(inputOperator, groupSet, orderSet, orderSetDirection, orderSetNullDirection, aggCalls, lowerBound, - upperBound, windowFrameType, constants, resultSchema, inputSchema, AggregationUtils.Accumulator.MERGERS, - context); + DataSchema resultSchema, DataSchema inputSchema) { + this(context, inputOperator, groupSet, orderSet, orderSetDirection, orderSetNullDirection, aggCalls, lowerBound, + upperBound, windowFrameType, constants, resultSchema, inputSchema, AggregationUtils.Accumulator.MERGERS); } @VisibleForTesting - public WindowAggregateOperator(MultiStageOperator inputOperator, List groupSet, - List orderSet, List orderSetDirection, + public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator inputOperator, + List groupSet, List orderSet, List orderSetDirection, List orderSetNullDirection, List aggCalls, int lowerBound, int upperBound, WindowNode.WindowFrameType windowFrameType, List constants, - DataSchema resultSchema, DataSchema inputSchema, Map> mergers, OpChainExecutionContext context) { + DataSchema resultSchema, DataSchema inputSchema, + Map> mergers) { super(context); boolean isPartitionByOnly = isPartitionByOnlyQuery(groupSet, orderSet, orderSetDirection, orderSetNullDirection); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java index 53ff0ff658e8..4a9892a4dd85 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PhysicalPlanVisitor.java @@ -56,15 +56,14 @@ public class PhysicalPlanVisitor implements StageNodeVisitor merger), - OperatorTestUtil.getDefaultContext()); + new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, inSchema, + ImmutableMap.of("SUM", cdt -> merger)); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) @@ -225,10 +225,9 @@ public void testGroupByAggregateWithHashCollision() { // Create an aggregation call with sum for first column and group by second column. RexExpression.FunctionCall agg = getSum(new RexExpression.InputRef(0)); DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, INT}); - AggregateOperator sum0GroupBy1 = - new AggregateOperator(upstreamOperator, OperatorTestUtil.getDataSchema(OperatorTestUtil.OP_1), - Arrays.asList(agg), Arrays.asList(new RexExpression.InputRef(1)), inSchema, - OperatorTestUtil.getDefaultContext()); + AggregateOperator sum0GroupBy1 = new AggregateOperator(OperatorTestUtil.getDefaultContext(), upstreamOperator, + OperatorTestUtil.getDataSchema(OperatorTestUtil.OP_1), Arrays.asList(agg), + Arrays.asList(new RexExpression.InputRef(1)), inSchema); TransferableBlock result = sum0GroupBy1.getNextBlock(); while (result.isNoOpBlock()) { result = sum0GroupBy1.getNextBlock(); @@ -252,7 +251,7 @@ public void shouldThrowOnUnknownAggFunction() { // When: AggregateOperator operator = - new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); + new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, inSchema); } @Test @@ -270,7 +269,7 @@ public void shouldReturnErrorBlockOnUnexpectedInputType() { DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); AggregateOperator operator = - new AggregateOperator(_input, outSchema, calls, group, inSchema, OperatorTestUtil.getDefaultContext()); + new AggregateOperator(OperatorTestUtil.getDefaultContext(), _input, outSchema, calls, group, inSchema); // When: TransferableBlock block = operator.nextBlock(); 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 3ee4aa5cca30..ba02d44fe174 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 @@ -62,7 +62,7 @@ public void shouldPropagateUpstreamErrorBlock() { DataSchema.ColumnDataType.BOOLEAN }); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock error = errorBlock.getDataBlock(); @@ -78,7 +78,7 @@ public void shouldPropagateUpstreamEOS() { }); Mockito.when(_upstreamOperator.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertTrue(dataBlock.isEndOfStreamBlock()); } @@ -92,7 +92,7 @@ public void shouldPropagateUpstreamNoop() { }); Mockito.when(_upstreamOperator.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertTrue(dataBlock.isNoOpBlock()); } @@ -108,7 +108,7 @@ public void shouldHandleTrueBooleanLiteralFilter() { .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{0}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -127,7 +127,7 @@ public void shouldHandleFalseBooleanLiteralFilter() { Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -143,7 +143,7 @@ public void shouldThrowOnNonBooleanTypeBooleanLiteral() { Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, booleanLiteral, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, booleanLiteral); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock data = errorBlock.getDataBlock(); @@ -158,7 +158,7 @@ public void shouldThrowOnNonBooleanTypeInputRef() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1}, new Object[]{2})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref0, OperatorTestUtil.getDefaultContext()); + FilterOperator op = new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, ref0); TransferableBlock errorBlock = op.getNextBlock(); Assert.assertTrue(errorBlock.isErrorBlock()); DataBlock data = errorBlock.getDataBlock(); @@ -173,7 +173,7 @@ public void shouldHandleBooleanInputRef() { }); Mockito.when(_upstreamOperator.nextBlock()) .thenReturn(OperatorTestUtil.block(inputSchema, new Object[]{1, true}, new Object[]{2, false})); - FilterOperator op = new FilterOperator(_upstreamOperator, inputSchema, ref1, OperatorTestUtil.getDefaultContext()); + FilterOperator op = new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, ref1); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -194,7 +194,7 @@ public void shouldHandleAndFilter() { ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, andCall, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, andCall); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -215,7 +215,7 @@ public void shouldHandleOrFilter() { ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, orCall, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, orCall); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -238,7 +238,7 @@ public void shouldHandleNotFilter() { ImmutableList.of(new RexExpression.InputRef(0))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, notCall, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, notCall); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -258,7 +258,7 @@ public void shouldHandleGreaterThanFilter() { new RexExpression.FunctionCall(SqlKind.GREATER_THAN, FieldSpec.DataType.BOOLEAN, "greaterThan", ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, greaterThan, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, greaterThan); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -279,7 +279,7 @@ public void shouldHandleBooleanFunction() { ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(FieldSpec.DataType.STRING, "star"))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, startsWith, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _upstreamOperator, inputSchema, startsWith); TransferableBlock dataBlock = op.getNextBlock(); Assert.assertFalse(dataBlock.isErrorBlock()); List result = dataBlock.getContainer(); @@ -301,6 +301,6 @@ public void shouldThrowOnUnfoundFunction() { ImmutableList.of(new RexExpression.InputRef(0), new RexExpression.Literal(FieldSpec.DataType.STRING, "star"))); FilterOperator op = - new FilterOperator(_upstreamOperator, inputSchema, startsWith, OperatorTestUtil.getDefaultContext()); + new FilterOperator(OperatorTestUtil.getDefaultContext(), _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 68de25a4a5f3..c9a48f9ba93b 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 @@ -96,7 +96,7 @@ public void shouldHandleHashJoinKeyCollisionInnerJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); HashJoinOperator joinOnString = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = joinOnString.nextBlock(); while (result.isNoOpBlock()) { @@ -134,7 +134,7 @@ public void shouldHandleInnerJoinOnInt() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator joinOnInt = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = joinOnInt.nextBlock(); while (result.isNoOpBlock()) { result = joinOnInt.nextBlock(); @@ -169,7 +169,7 @@ public void shouldHandleJoinOnEmptySelector() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); HashJoinOperator joinOnInt = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = joinOnInt.nextBlock(); while (result.isNoOpBlock()) { result = joinOnInt.nextBlock(); @@ -211,7 +211,7 @@ public void shouldHandleLeftJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -246,7 +246,7 @@ public void shouldPassLeftTableEOS() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -278,7 +278,7 @@ public void shouldHandleLeftJoinOneToN() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.LEFT, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -314,7 +314,7 @@ public void shouldPassRightTableEOS() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -353,7 +353,7 @@ public void shouldHandleInequiJoinOnString() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -392,7 +392,7 @@ public void shouldHandleInequiJoinOnInt() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(new ArrayList<>(), new ArrayList<>()), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -427,7 +427,7 @@ public void shouldHandleRightJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.RIGHT, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator joinOnNum = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = joinOnNum.nextBlock(); while (result.isNoOpBlock()) { result = joinOnNum.nextBlock(); @@ -477,7 +477,7 @@ public void shouldHandleSemiJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.SEMI, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -517,7 +517,7 @@ public void shouldHandleFullJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.FULL, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -570,7 +570,7 @@ public void shouldHandleAntiJoin() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.ANTI, getJoinKeys(Arrays.asList(1), Arrays.asList(1)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { result = join.nextBlock(); @@ -609,7 +609,7 @@ public void shouldPropagateRightTableError() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -643,7 +643,7 @@ public void shouldPropagateLeftTableError() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); while (result.isNoOpBlock()) { @@ -680,7 +680,7 @@ public void shouldHandleNoOpBlock() { JoinNode node = new JoinNode(1, resultSchema, leftSchema, rightSchema, JoinRelType.INNER, getJoinKeys(Arrays.asList(0), Arrays.asList(0)), joinClauses); HashJoinOperator join = - new HashJoinOperator(_leftOperator, _rightOperator, leftSchema, node, OperatorTestUtil.getDefaultContext()); + new HashJoinOperator(OperatorTestUtil.getDefaultContext(), _leftOperator, _rightOperator, leftSchema, node); TransferableBlock result = join.nextBlock(); // first no-op consumes first right data block. Assert.assertTrue(result.isNoOpBlock()); 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 f5a39d735e91..9d04d07df93d 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 @@ -75,7 +75,7 @@ public void shouldReturnDataBlockThenMetadataBlock() { List resultsBlockList = Collections.singletonList(new InstanceResponseBlock( new SelectionResultsBlock(schema, Arrays.asList(new Object[]{"foo", 1}, new Object[]{"", 2})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -101,7 +101,7 @@ public void shouldHandleDesiredDataSchemaConversionCorrectly() { new SelectionResultsBlock(resultSchema, Arrays.asList(new Object[]{1, 1660000000000L}, new Object[]{0, 1600000000000L})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, desiredSchema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -123,7 +123,7 @@ public void shouldHandleCanonicalizationCorrectly() { new SelectionResultsBlock(schema, Arrays.asList(new Object[]{1, 1660000000000L}, new Object[]{0, 1600000000000L})), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -148,7 +148,7 @@ public void shouldReturnMultipleDataBlockThenMetadataBlock() { queryContext), new InstanceResponseBlock(new SelectionResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock1 = operator.nextBlock(); @@ -178,7 +178,7 @@ public void shouldGetErrorBlockWhenInstanceResponseContainsError() { errorBlock, new InstanceResponseBlock(new SelectionResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -199,7 +199,7 @@ public void shouldReorderWhenQueryContextAskForNotInOrderGroupByAsDistinct() { new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(schema, Arrays.asList(new Record(new Object[]{1, "foo"}), new Record(new Object[]{2, "bar"})))), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -220,7 +220,7 @@ public void shouldParsedBlocksSuccessfullyWithDistinctQuery() { new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(schema, Arrays.asList(new Record(new Object[]{"foo", 1}), new Record(new Object[]{"bar", 2})))), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -244,7 +244,7 @@ public void shouldReorderWhenQueryContextAskForGroupByOutOfOrder() { List resultsBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -267,7 +267,7 @@ public void shouldNotErrorOutWhenQueryContextAskForGroupByOutOfOrderWithHaving() List resultsBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(schema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -286,7 +286,7 @@ public void shouldNotErrorOutWhenDealingWithAggregationResults() { List resultsBlockList = Collections.singletonList(new InstanceResponseBlock( new AggregationResultsBlock(queryContext.getAggregationFunctions(), Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(resultsBlockList, schema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), resultsBlockList, schema); // When: TransferableBlock resultBlock = operator.nextBlock(); @@ -308,7 +308,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsSelecti List responseBlockList = Collections.singletonList( new InstanceResponseBlock(new SelectionResultsBlock(resultSchema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), responseBlockList, desiredSchema); TransferableBlock resultBlock = operator.nextBlock(); // Then: @@ -331,7 +331,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsDistinc new DistinctResultsBlock(mock(DistinctAggregationFunction.class), new DistinctTable(resultSchema, Collections.emptyList())), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), responseBlockList, desiredSchema); TransferableBlock resultBlock = operator.nextBlock(); // Then: @@ -353,7 +353,7 @@ public void shouldNotErrorOutWhenIncorrectDataSchemaProvidedWithEmptyRowsGroupBy List responseBlockList = Collections.singletonList( new InstanceResponseBlock(new GroupByResultsBlock(resultSchema, Collections.emptyList()), queryContext)); LeafStageTransferableBlockOperator operator = - new LeafStageTransferableBlockOperator(responseBlockList, desiredSchema, OperatorTestUtil.getDefaultContext()); + new LeafStageTransferableBlockOperator(OperatorTestUtil.getDefaultContext(), responseBlockList, desiredSchema); TransferableBlock resultBlock = operator.nextBlock(); // Then: 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 705afc0192ca..8276c647aacd 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 @@ -66,7 +66,7 @@ public void shouldReturnLiteralBlock() { List> literals = ImmutableList.of( ImmutableList.of(new RexExpression.Literal(DataType.STRING, "foo"), new RexExpression.Literal(DataType.INT, 1)), ImmutableList.of(new RexExpression.Literal(DataType.STRING, ""), new RexExpression.Literal(DataType.INT, 2))); - LiteralValueOperator operator = new LiteralValueOperator(schema, literals, OperatorTestUtil.getDefaultContext()); + LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getDefaultContext(), schema, literals); // When: TransferableBlock transferableBlock = operator.nextBlock(); @@ -82,7 +82,7 @@ public void shouldHandleEmptyLiteralRows() { // Given: DataSchema schema = new DataSchema(new String[]{}, new ColumnDataType[]{}); List> literals = ImmutableList.of(ImmutableList.of()); - LiteralValueOperator operator = new LiteralValueOperator(schema, literals, OperatorTestUtil.getDefaultContext()); + LiteralValueOperator operator = new LiteralValueOperator(OperatorTestUtil.getDefaultContext(), 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 bece567f160d..ebd712a0c452 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 @@ -85,7 +85,7 @@ public void shouldTimeoutOnExtraLongSleep() new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 10L, 10L, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 10L, context); + new MailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 10L); Thread.sleep(200L); TransferableBlock mailbox = receiveOp.nextBlock(); Assert.assertTrue(mailbox.isErrorBlock()); @@ -95,13 +95,13 @@ public void shouldTimeoutOnExtraLongSleep() // longer timeout or default timeout (10s) doesn't result in error. context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 2000L, 2000L, new HashMap<>()); - receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 2000L, context); + receiveOp = new MailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 2000L); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); - receiveOp = new MailboxReceiveOperator(new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, null, context); + receiveOp = new MailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, null); Thread.sleep(200L); mailbox = receiveOp.nextBlock(); Assert.assertFalse(mailbox.isErrorBlock()); @@ -124,8 +124,8 @@ public void shouldThrowReceiveSingletonFromMultiMatchMailboxServer() { new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, 456, 789, null, - context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, 456, + 789, null); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*RANGE_DISTRIBUTED.*") @@ -142,9 +142,8 @@ public void shouldThrowRangeDistributionNotSupported() { OpChainExecutionContext context = new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE, Long.MAX_VALUE, new HashMap<>()); - MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.RANGE_DISTRIBUTED, 456, - 789, null, context); + MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), + RelDistribution.Type.RANGE_DISTRIBUTED, 456, 789, null); } @Test @@ -173,8 +172,8 @@ public void shouldReceiveSingletonNoMatchMailboxServer() { Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive end of stream block directly when there is no match. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); @@ -211,8 +210,8 @@ public void shouldReceiveSingletonCloseMailbox() { Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive end of stream block directly when mailbox is close. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); @@ -251,8 +250,8 @@ public void shouldReceiveSingletonNullMailbox() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive NoOpBlock. Assert.assertTrue(receiveOp.nextBlock().isNoOpBlock()); } @@ -289,8 +288,8 @@ public void shouldReceiveEosDirectlyFromSender() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive EosBloc. Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock()); } @@ -329,8 +328,8 @@ public void shouldReceiveSingletonMailbox() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -370,8 +369,8 @@ public void shouldReceiveSingletonErrorMailbox() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock(); @@ -414,8 +413,8 @@ public void shouldReceiveMailboxFromTwoServersOneClose() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -459,8 +458,8 @@ public void shouldReceiveMailboxFromTwoServersOneNull() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); Assert.assertEquals(resultRows.size(), 1); @@ -508,8 +507,8 @@ public void shouldReceiveMailboxFromTwoServers() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive first block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); List resultRows = receivedBlock.getContainer(); @@ -566,8 +565,8 @@ public void shouldGetReceptionReceiveErrorMailbox() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); // Receive error block from first server. TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock()); @@ -612,8 +611,8 @@ public void shouldThrowReceiveWhenOneServerReceiveThrowException() Long.MAX_VALUE, new HashMap<>()); MailboxReceiveOperator receiveOp = - new MailboxReceiveOperator(ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, stageId, - DEFAULT_RECEIVER_STAGE_ID, null, context); + new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, + stageId, DEFAULT_RECEIVER_STAGE_ID, null); TransferableBlock receivedBlock = receiveOp.nextBlock(); Assert.assertTrue(receivedBlock.isErrorBlock(), "server-1 should have returned an error-block"); } 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 ac57ef5af521..3f7a6ae4c4f9 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 @@ -91,9 +91,10 @@ public void shouldSwallowNoOpBlockFromUpstream() // OperatorTestUtil.getContext(1, DEFAULT_SENDER_STAGE_ID, new VirtualServerAddress(_server)); OpChainExecutionContext context = getOpChainContext(deadlineMs); - MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, - server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + MailboxSendOperator operator = + new MailboxSendOperator(context, _input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, + server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); @@ -112,9 +113,10 @@ public void shouldSendErrorBlock() // Given: OpChainExecutionContext context = getOpChainContext(deadlineMs); - MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, - server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + MailboxSendOperator operator = + new MailboxSendOperator(context, _input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, + server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID); TransferableBlock errorBlock = TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!")); Mockito.when(_input.nextBlock()).thenReturn(errorBlock); @@ -133,9 +135,10 @@ public void shouldSendErrorBlockWhenInputThrows() // Given: OpChainExecutionContext context = getOpChainContext(deadlineMs); - MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, - server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + MailboxSendOperator operator = + new MailboxSendOperator(context, _input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, + server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID); Mockito.when(_input.nextBlock()).thenThrow(new RuntimeException("foo!")); ArgumentCaptor captor = ArgumentCaptor.forClass(TransferableBlock.class); @@ -155,9 +158,10 @@ public void shouldSendEosBlock() // Given: OpChainExecutionContext context = getOpChainContext(deadlineMs); - MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, - server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + MailboxSendOperator operator = + new MailboxSendOperator(context, _input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, + server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID); TransferableBlock eosBlock = TransferableBlockUtils.getEndOfStreamTransferableBlock(); Mockito.when(_input.nextBlock()).thenReturn(eosBlock); @@ -177,9 +181,10 @@ public void shouldSendDataBlock() // Given: OpChainExecutionContext context = getOpChainContext(deadlineMs); - MailboxSendOperator operator = new MailboxSendOperator(_input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, - server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, - DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID, context); + MailboxSendOperator operator = + new MailboxSendOperator(context, _input, RelDistribution.Type.HASH_DISTRIBUTED, _selector, + server -> new JsonMailboxIdentifier("123", "0@from:1", "0@to:2", DEFAULT_SENDER_STAGE_ID, + DEFAULT_RECEIVER_STAGE_ID), _exchangeFactory, DEFAULT_RECEIVER_STAGE_ID); TransferableBlock dataBlock = block(new DataSchema(new String[]{}, new DataSchema.ColumnDataType[]{})); Mockito.when(_input.nextBlock()).thenReturn(dataBlock) .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 6f311ad8e72d..11e525221d18 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -70,7 +70,7 @@ public void shouldHandleUpstreamErrorBlock() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()) .thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!"))); @@ -89,7 +89,7 @@ public void shouldHandleUpstreamNoOpBlock() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); @@ -107,7 +107,7 @@ public void shouldCreateEmptyBlockOnUpstreamEOS() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -125,7 +125,7 @@ public void shouldConsumeAndSortInputOneBlockWithTwoRows() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -148,7 +148,7 @@ public void shouldConsumeAndSortOnNonZeroIdxCollation() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"ignored", "sort"}, new DataSchema.ColumnDataType[]{INT, INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{1, 2}, new Object[]{2, 1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -171,7 +171,7 @@ public void shouldConsumeAndSortInputOneBlockWithTwoRowsNonNumeric() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{STRING}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{"b"}, new Object[]{"a"})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -194,7 +194,7 @@ public void shouldConsumeAndSortDescending() { List directions = ImmutableList.of(Direction.DESCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -217,7 +217,7 @@ public void shouldOffsetSortInputOneBlockWithThreeRows() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 1, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 1, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -240,7 +240,7 @@ public void shouldOffsetLimitSortInputOneBlockWithThreeRows() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 1, 1, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 1, 1, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -262,7 +262,7 @@ public void shouldRespectDefaultLimit() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 0, 0, schema, 1, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 0, 0, schema, 1); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -284,7 +284,7 @@ public void shouldFetchAllWithNegativeFetch() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, -1, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, -1, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2}, new Object[]{1}, new Object[]{3})) .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); @@ -305,7 +305,7 @@ public void shouldConsumeAndSortTwoInputBlocksWithOneRowEach() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2})) .thenReturn(block(schema, new Object[]{1})) @@ -329,7 +329,7 @@ public void shouldBreakTiesUsingSecondCollationKey() { List directions = ImmutableList.of(Direction.ASCENDING, Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"first", "second"}, new DataSchema.ColumnDataType[]{INT, INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()) .thenReturn(block(schema, new Object[]{1, 2}, new Object[]{1, 1}, new Object[]{1, 3})) @@ -354,7 +354,7 @@ public void shouldBreakTiesUsingSecondCollationKeyWithDifferentDirection() { List directions = ImmutableList.of(Direction.ASCENDING, Direction.DESCENDING); DataSchema schema = new DataSchema(new String[]{"first", "second"}, new DataSchema.ColumnDataType[]{INT, INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()) .thenReturn(block(schema, new Object[]{1, 2}, new Object[]{1, 1}, new Object[]{1, 3})) @@ -379,7 +379,7 @@ public void shouldHandleNoOpUpstreamBlockWhileConstructing() { List directions = ImmutableList.of(Direction.ASCENDING); DataSchema schema = new DataSchema(new String[]{"sort"}, new DataSchema.ColumnDataType[]{INT}); SortOperator op = - new SortOperator(_input, collation, directions, 10, 0, schema, OperatorTestUtil.getDefaultContext()); + new SortOperator(OperatorTestUtil.getDefaultContext(), _input, collation, directions, 10, 0, schema); Mockito.when(_input.nextBlock()).thenReturn(block(schema, new Object[]{2})) .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()).thenReturn(block(schema, new Object[]{1})) diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java index 4afc580cdbc0..46279f03a23f 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/TransformOperatorTest.java @@ -71,8 +71,8 @@ public void shouldHandleRefTransform() { RexExpression.InputRef ref0 = new RexExpression.InputRef(0); RexExpression.InputRef ref1 = new RexExpression.InputRef(1); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0, ref1), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(ref0, ref1), upStreamSchema); TransferableBlock result = op.nextBlock(); Assert.assertTrue(!result.isErrorBlock()); @@ -96,8 +96,8 @@ public void shouldHandleLiteralTransform() { RexExpression.Literal boolLiteral = new RexExpression.Literal(FieldSpec.DataType.BOOLEAN, true); RexExpression.Literal strLiteral = new RexExpression.Literal(FieldSpec.DataType.STRING, "str"); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema); TransferableBlock result = op.nextBlock(); // Literal operands should just output original literals. Assert.assertTrue(!result.isErrorBlock()); @@ -127,8 +127,8 @@ public void shouldHandlePlusMinusFuncTransform() { DataSchema resultSchema = new DataSchema(new String[]{"plusR", "minusR"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.DOUBLE, DataSchema.ColumnDataType.DOUBLE}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema); TransferableBlock result = op.nextBlock(); Assert.assertTrue(!result.isErrorBlock()); List resultRows = result.getContainer(); @@ -156,8 +156,8 @@ public void shouldThrowOnTypeMismatchFuncTransform() { DataSchema resultSchema = new DataSchema(new String[]{"plusR", "minusR"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.DOUBLE, DataSchema.ColumnDataType.DOUBLE}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(plus01, minus01), upStreamSchema); TransferableBlock result = op.nextBlock(); Assert.assertTrue(result.isErrorBlock()); @@ -177,8 +177,8 @@ public void shouldPropagateUpstreamError() { DataSchema resultSchema = new DataSchema(new String[]{"inCol", "strCol"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema); TransferableBlock result = op.nextBlock(); Assert.assertTrue(result.isErrorBlock()); DataBlock data = result.getDataBlock(); @@ -201,8 +201,8 @@ public void testNoopBlock() { DataSchema resultSchema = new DataSchema(new String[]{"boolCol", "strCol"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BOOLEAN, DataSchema.ColumnDataType.STRING}); TransformOperator op = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), + _upstreamOp, resultSchema, ImmutableList.of(boolLiteral, strLiteral), upStreamSchema); TransferableBlock result = op.nextBlock(); // First block has two rows Assert.assertFalse(result.isErrorBlock()); @@ -233,8 +233,9 @@ public void testWrongNumTransform() { DataSchema upStreamSchema = new DataSchema(new String[]{"string1", "string2"}, new DataSchema.ColumnDataType[]{ DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.STRING }); - TransformOperator transform = new TransformOperator(_upstreamOp, resultSchema, new ArrayList<>(), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + TransformOperator transform = + new TransformOperator(OperatorTestUtil.getDefaultContext(), _upstreamOp, resultSchema, new ArrayList<>(), + upStreamSchema); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*doesn't match " @@ -247,7 +248,7 @@ public void testMismatchedSchemaOperandSize() { }); RexExpression.InputRef ref0 = new RexExpression.InputRef(0); TransformOperator transform = - new TransformOperator(_upstreamOp, resultSchema, ImmutableList.of(ref0), upStreamSchema, - OperatorTestUtil.getDefaultContext()); + new TransformOperator(OperatorTestUtil.getDefaultContext(), _upstreamOp, resultSchema, ImmutableList.of(ref0), + upStreamSchema); } }; diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 92b32383f182..ea63f8573125 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -81,9 +81,9 @@ public void testShouldHandleUpstreamErrorBlocks() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); // build @@ -105,9 +105,9 @@ public void testShouldHandleEndOfStreamBlockWithNoOtherInputs() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block = operator.nextBlock(); @@ -131,9 +131,9 @@ public void testShouldHandleUpstreamNoOpBlocksWhileConstructing() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); // build when reading NoOp block @@ -162,9 +162,9 @@ public void testShouldHandleUpstreamNoOpBlocksWhileConstructingMultipleRows() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); // build when reading NoOp block @@ -196,9 +196,9 @@ public void testShouldWindowAggregateOverSingleInputBlock() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); @@ -226,10 +226,10 @@ public void testShouldWindowAggregateOverSingleInputBlockWithSameOrderByKeys() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), - Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, - OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, order, + Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), + calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), + outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); @@ -255,9 +255,9 @@ public void testShouldWindowAggregateOverSingleInputBlockWithoutPartitionByKeys( DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, + Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); @@ -284,9 +284,9 @@ public void testShouldWindowAggregateOverSingleInputBlockWithLiteralInput() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block1 = operator.nextBlock(); @@ -319,10 +319,10 @@ public void testShouldCallMergerWhenWindowAggregatingMultipleRows() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, ImmutableMap.of("SUM", cdt -> merger), - OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, + ImmutableMap.of("SUM", cdt -> merger)); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) @@ -349,10 +349,10 @@ public void testPartitionByWindowAggregateWithHashCollision() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, INT, DOUBLE}); WindowAggregateOperator sum0PartitionBy1 = - new WindowAggregateOperator(upstreamOperator, Arrays.asList(new RexExpression.InputRef(1)), - Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Arrays.asList(agg), - Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, - inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), upstreamOperator, + Arrays.asList(new RexExpression.InputRef(1)), Collections.emptyList(), Collections.emptyList(), + Collections.emptyList(), Arrays.asList(agg), Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); TransferableBlock result = sum0PartitionBy1.getNextBlock(); while (result.isNoOpBlock()) { @@ -379,9 +379,9 @@ public void testShouldThrowOnUnknownAggFunction() { // When: WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*Unexpected aggregation " @@ -397,9 +397,9 @@ public void testShouldThrowOnUnknownRankAggFunction() { // When: WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Order by is not yet " @@ -418,10 +418,10 @@ public void testShouldThrowOnNonEmptyOrderByKeysNotMatchingPartitionByKeys() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), - Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, - OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, order, + Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), + calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), + outSchema, inSchema); } @Test @@ -443,10 +443,10 @@ public void testShouldThrowOnNonEmptyOrderByKeysMatchingPartitionByKeysWithDiffe DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.DESCENDING), - Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, - OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, order, + Arrays.asList(RelFieldCollation.Direction.DESCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), + calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), + outSchema, inSchema); // When: TransferableBlock resultBlock = operator.nextBlock(); // (output result) @@ -475,9 +475,9 @@ public void testShouldThrowOnCustomFramesRows() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.ROW, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.ROW, Collections.emptyList(), outSchema, inSchema); } @Test @@ -494,10 +494,10 @@ public void testShouldNotThrowCurrentRowPartitionByOrderByOnSameKey() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, order, Arrays.asList(RelFieldCollation.Direction.ASCENDING), - Arrays.asList(RelFieldCollation.NullDirection.LAST), calls, Integer.MIN_VALUE, 0, - WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema, - OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, order, + Arrays.asList(RelFieldCollation.Direction.ASCENDING), Arrays.asList(RelFieldCollation.NullDirection.LAST), + calls, Integer.MIN_VALUE, 0, WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, + inSchema); // When: TransferableBlock block1 = operator.nextBlock(); @@ -526,9 +526,9 @@ public void testShouldThrowOnCustomFramesCustomPreceding() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, 5, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, 5, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); } @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Only default frame is " @@ -546,9 +546,9 @@ public void testShouldThrowOnCustomFramesCustomFollowing() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, 5, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, 5, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); } @Test @@ -567,9 +567,9 @@ public void testShouldReturnErrorBlockOnUnexpectedInputType() { DataSchema outSchema = new DataSchema(new String[]{"group", "arg", "sum"}, new DataSchema.ColumnDataType[]{INT, STRING, DOUBLE}); WindowAggregateOperator operator = - new WindowAggregateOperator(_input, group, Collections.emptyList(), Collections.emptyList(), - Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowFrameType.RANGE, - Collections.emptyList(), outSchema, inSchema, OperatorTestUtil.getDefaultContext()); + new WindowAggregateOperator(OperatorTestUtil.getDefaultContext(), _input, group, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), calls, Integer.MIN_VALUE, Integer.MAX_VALUE, + WindowNode.WindowFrameType.RANGE, Collections.emptyList(), outSchema, inSchema); // When: TransferableBlock block = operator.nextBlock(); From 09f58ca3e56fd583cf8a190dd9728f7d4b64a229 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Thu, 16 Mar 2023 15:01:33 +0530 Subject: [PATCH 12/12] Remove comments --- .../pinot/query/runtime/operator/MailboxSendOperatorTest.java | 2 -- 1 file changed, 2 deletions(-) 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 3f7a6ae4c4f9..13735ed883d7 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 @@ -87,8 +87,6 @@ public void shouldSwallowNoOpBlockFromUpstream() throws Exception { long deadlineMs = System.currentTimeMillis() + 10_000; // Given: -// OpChainExecutionContext context = -// OperatorTestUtil.getContext(1, DEFAULT_SENDER_STAGE_ID, new VirtualServerAddress(_server)); OpChainExecutionContext context = getOpChainContext(deadlineMs); MailboxSendOperator operator =