You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ja...@apache.org on 2023/04/23 01:34:47 UTC
[pinot] branch master updated: Enhance mailbox receive operator (#10669)
This is an automated email from the ASF dual-hosted git repository.
jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 866c796bd5 Enhance mailbox receive operator (#10669)
866c796bd5 is described below
commit 866c796bd56cf846b654f29f024f3e610557b2c7
Author: Xiaotian (Jackie) Jiang <17...@users.noreply.github.com>
AuthorDate: Sat Apr 22 18:34:41 2023 -0700
Enhance mailbox receive operator (#10669)
---
.../pinot/query/mailbox/GrpcMailboxService.java | 11 +-
.../pinot/query/mailbox/GrpcReceivingMailbox.java | 10 +-
.../query/mailbox/InMemoryMailboxService.java | 11 +-
.../query/mailbox/InMemoryReceivingMailbox.java | 11 +-
.../apache/pinot/query/mailbox/MailboxService.java | 3 +-
.../query/mailbox/MultiplexingMailboxService.java | 10 +-
.../pinot/query/mailbox/ReceivingMailbox.java | 2 +
.../operator/BaseMailboxReceiveOperator.java | 85 +-
.../runtime/operator/MailboxReceiveOperator.java | 103 +-
.../query/runtime/operator/MultiStageOperator.java | 7 +-
.../operator/SortedMailboxReceiveOperator.java | 168 ++--
.../query/runtime/plan/PhysicalPlanVisitor.java | 8 +-
.../query/service/dispatch/QueryDispatcher.java | 4 +-
.../operator/MailboxReceiveOperatorTest.java | 739 +++++---------
.../pinot/query/runtime/operator/OpChainTest.java | 19 +-
.../operator/SortedMailboxReceiveOperatorTest.java | 1050 ++++++--------------
16 files changed, 752 insertions(+), 1489 deletions(-)
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcMailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcMailboxService.java
index 9db144a48a..b2b6f1e1aa 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcMailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcMailboxService.java
@@ -121,7 +121,7 @@ public class GrpcMailboxService implements MailboxService<TransferableBlock> {
public ReceivingMailbox<TransferableBlock> getReceivingMailbox(MailboxIdentifier mailboxId) {
try {
return _receivingMailboxCache.get(mailboxId,
- () -> new GrpcReceivingMailbox(_gotMailCallback));
+ () -> new GrpcReceivingMailbox(mailboxId, _gotMailCallback));
} catch (ExecutionException e) {
LOGGER.error(String.format("Error getting receiving mailbox: %s", mailboxId), e);
throw new RuntimeException(e);
@@ -140,12 +140,11 @@ public class GrpcMailboxService implements MailboxService<TransferableBlock> {
* </p>
*/
@Override
- public void releaseReceivingMailbox(MailboxIdentifier mailboxId) {
- GrpcReceivingMailbox receivingMailbox = _receivingMailboxCache.getIfPresent(mailboxId);
- if (receivingMailbox != null && !receivingMailbox.isClosed()) {
- receivingMailbox.cancel();
+ public void releaseReceivingMailbox(ReceivingMailbox<TransferableBlock> mailbox) {
+ if (!mailbox.isClosed()) {
+ mailbox.cancel();
}
- _receivingMailboxCache.invalidate(mailboxId);
+ _receivingMailboxCache.invalidate(mailbox.getId());
}
@Override
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcReceivingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcReceivingMailbox.java
index 96add5d702..85ca408c10 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcReceivingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcReceivingMailbox.java
@@ -47,13 +47,16 @@ import org.slf4j.LoggerFactory;
public class GrpcReceivingMailbox implements ReceivingMailbox<TransferableBlock> {
private static final Logger LOGGER = LoggerFactory.getLogger(GrpcReceivingMailbox.class);
private static final long DEFAULT_MAILBOX_INIT_TIMEOUT = 100L;
+
+ private final MailboxIdentifier _id;
private final Consumer<MailboxIdentifier> _gotMailCallback;
private final CountDownLatch _initializationLatch;
private MailboxContentStreamObserver _contentStreamObserver;
private StreamObserver<Mailbox.MailboxStatus> _statusStreamObserver;
- public GrpcReceivingMailbox(Consumer<MailboxIdentifier> gotMailCallback) {
+ public GrpcReceivingMailbox(MailboxIdentifier id, Consumer<MailboxIdentifier> gotMailCallback) {
+ _id = id;
_gotMailCallback = gotMailCallback;
_initializationLatch = new CountDownLatch(1);
}
@@ -68,6 +71,11 @@ public class GrpcReceivingMailbox implements ReceivingMailbox<TransferableBlock>
return _gotMailCallback;
}
+ @Override
+ public MailboxIdentifier getId() {
+ return _id;
+ }
+
/**
* Polls the underlying channel and converts the received data into a TransferableBlock. This may return null in the
* following cases:
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java
index ec156973be..aa98c78a3c 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java
@@ -89,7 +89,7 @@ public class InMemoryMailboxService implements MailboxService<TransferableBlock>
public ReceivingMailbox<TransferableBlock> getReceivingMailbox(MailboxIdentifier mailboxId) {
Preconditions.checkState(mailboxId.isLocal(), "Cannot use in-memory mailbox service for non-local transport");
try {
- return _receivingMailboxCache.get(mailboxId, InMemoryReceivingMailbox::new);
+ return _receivingMailboxCache.get(mailboxId, () -> new InMemoryReceivingMailbox(mailboxId));
} catch (ExecutionException e) {
LOGGER.error(String.format("Error getting in-memory receiving mailbox=%s", mailboxId), e);
throw new RuntimeException(e);
@@ -97,11 +97,10 @@ public class InMemoryMailboxService implements MailboxService<TransferableBlock>
}
@Override
- public void releaseReceivingMailbox(MailboxIdentifier mailboxId) {
- InMemoryReceivingMailbox receivingMailbox = _receivingMailboxCache.getIfPresent(mailboxId);
- if (receivingMailbox != null) {
- receivingMailbox.cancel();
- _receivingMailboxCache.invalidate(mailboxId);
+ public void releaseReceivingMailbox(ReceivingMailbox<TransferableBlock> mailbox) {
+ if (!mailbox.isClosed()) {
+ mailbox.cancel();
}
+ _receivingMailboxCache.invalidate(mailbox.getId());
}
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryReceivingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryReceivingMailbox.java
index 7ea76fe427..11c4b3ff09 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryReceivingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryReceivingMailbox.java
@@ -23,16 +23,25 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock;
public class InMemoryReceivingMailbox implements ReceivingMailbox<TransferableBlock> {
+
+ private final MailboxIdentifier _id;
+
private InMemoryTransferStream _transferStream;
private volatile boolean _closed = false;
- public InMemoryReceivingMailbox() {
+ public InMemoryReceivingMailbox(MailboxIdentifier id) {
+ _id = id;
}
public void init(InMemoryTransferStream transferStream) {
_transferStream = transferStream;
}
+ @Override
+ public MailboxIdentifier getId() {
+ return _id;
+ }
+
@Override
public TransferableBlock receive()
throws Exception {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
index b4d16b0a85..845d6bf96f 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MailboxService.java
@@ -92,7 +92,6 @@ public interface MailboxService<T> {
* the {@link MailboxService} to use this method. It can use any internal method it needs to do the clean-up.
* </li>
* </ol>
- * @param mailboxId
*/
- void releaseReceivingMailbox(MailboxIdentifier mailboxId);
+ void releaseReceivingMailbox(ReceivingMailbox<T> mailbox);
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MultiplexingMailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MultiplexingMailboxService.java
index c4f9c6fdec..cf786885ae 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MultiplexingMailboxService.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/MultiplexingMailboxService.java
@@ -80,12 +80,12 @@ public class MultiplexingMailboxService implements MailboxService<TransferableBl
}
@Override
- public void releaseReceivingMailbox(MailboxIdentifier mailboxId) {
- if (mailboxId.isLocal()) {
- _inMemoryMailboxService.releaseReceivingMailbox(mailboxId);
- return;
+ public void releaseReceivingMailbox(ReceivingMailbox<TransferableBlock> mailbox) {
+ if (mailbox instanceof InMemoryReceivingMailbox) {
+ _inMemoryMailboxService.releaseReceivingMailbox(mailbox);
+ } else {
+ _grpcMailboxService.releaseReceivingMailbox(mailbox);
}
- _grpcMailboxService.releaseReceivingMailbox(mailboxId);
}
public static MultiplexingMailboxService newInstance(String hostname, int port,
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
index 922004501f..72ebec9ea3 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/ReceivingMailbox.java
@@ -33,6 +33,8 @@ import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
*/
public interface ReceivingMailbox<T> {
+ MailboxIdentifier getId();
+
/**
* Returns a unit of data. Implementations are allowed to return null, in which case {@link MailboxReceiveOperator}
* will assume that this mailbox doesn't have any data to return and it will instead poll the other mailbox (if any).
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
index 0237c11393..2d73ee18de 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseMailboxReceiveOperator.java
@@ -19,21 +19,23 @@
package org.apache.pinot.query.runtime.operator;
import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
+import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.Deque;
import java.util.List;
import java.util.Set;
+import java.util.stream.Collectors;
import org.apache.calcite.rel.RelDistribution;
import org.apache.pinot.query.mailbox.JsonMailboxIdentifier;
import org.apache.pinot.query.mailbox.MailboxIdentifier;
import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.mailbox.ReceivingMailbox;
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.plan.OpChainExecutionContext;
-import org.apache.pinot.query.service.QueryConfig;
/**
@@ -54,83 +56,72 @@ public abstract class BaseMailboxReceiveOperator extends MultiStageOperator {
protected final MailboxService<TransferableBlock> _mailboxService;
protected final RelDistribution.Type _exchangeType;
- protected final List<MailboxIdentifier> _sendingMailbox;
- protected final long _deadlineTimestampNano;
- protected int _serverIdx;
- protected TransferableBlock _upstreamErrorBlock;
-
- protected static MailboxIdentifier toMailboxId(VirtualServer sender, long jobId, int senderStageId,
- int receiverStageId, VirtualServerAddress receiver) {
- return new JsonMailboxIdentifier(
- String.format("%s_%s", jobId, senderStageId),
- new VirtualServerAddress(sender),
- receiver,
- senderStageId,
- receiverStageId);
- }
+ protected final List<MailboxIdentifier> _mailboxIds;
+ protected final Deque<ReceivingMailbox<TransferableBlock>> _mailboxes;
public BaseMailboxReceiveOperator(OpChainExecutionContext context, List<VirtualServer> sendingStageInstances,
- RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, Long timeoutMs) {
+ RelDistribution.Type exchangeType, int senderStageId, int receiverStageId) {
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;
- _deadlineTimestampNano = timeoutNano + System.nanoTime();
-
+ Preconditions.checkState(SUPPORTED_EXCHANGE_TYPES.contains(exchangeType), "Unsupported exchange type: %s",
+ exchangeType);
_exchangeType = exchangeType;
- if (_exchangeType == RelDistribution.Type.SINGLETON) {
+
+ long requestId = context.getRequestId();
+ VirtualServerAddress receiver = context.getServer();
+ if (exchangeType == RelDistribution.Type.SINGLETON) {
VirtualServer singletonInstance = null;
for (VirtualServer serverInstance : sendingStageInstances) {
if (serverInstance.getHostname().equals(_mailboxService.getHostname())
&& serverInstance.getQueryMailboxPort() == _mailboxService.getMailboxPort()) {
- Preconditions.checkState(singletonInstance == null, "multiple instance found for singleton exchange type!");
+ Preconditions.checkState(singletonInstance == null, "Multiple instances found for SINGLETON exchange type");
singletonInstance = serverInstance;
}
}
-
- if (singletonInstance == null) {
- // TODO: fix WorkerManager assignment, this should not happen if we properly assign workers.
- // see: https://github.com/apache/pinot/issues/9611
- _sendingMailbox = Collections.emptyList();
- } else {
- _sendingMailbox =
- Collections.singletonList(toMailboxId(singletonInstance, jobId, senderStageId, receiverStageId, receiver));
- }
+ Preconditions.checkState(singletonInstance != null, "Failed to find instance for SINGLETON exchange type");
+ _mailboxIds = Collections.singletonList(
+ toMailboxId(singletonInstance, requestId, senderStageId, receiverStageId, receiver));
} else {
- _sendingMailbox = new ArrayList<>(sendingStageInstances.size());
+ _mailboxIds = new ArrayList<>(sendingStageInstances.size());
for (VirtualServer instance : sendingStageInstances) {
- _sendingMailbox.add(toMailboxId(instance, jobId, senderStageId, receiverStageId, receiver));
+ _mailboxIds.add(toMailboxId(instance, requestId, senderStageId, receiverStageId, receiver));
}
}
- _upstreamErrorBlock = null;
- _serverIdx = 0;
+ _mailboxes = _mailboxIds.stream().map(_mailboxService::getReceivingMailbox)
+ .collect(Collectors.toCollection(ArrayDeque::new));
+ }
+
+ protected static MailboxIdentifier toMailboxId(VirtualServer sender, long requestId, int senderStageId,
+ int receiverStageId, VirtualServerAddress receiver) {
+ return new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), new VirtualServerAddress(sender),
+ receiver, senderStageId, receiverStageId);
}
- public List<MailboxIdentifier> getSendingMailbox() {
- return _sendingMailbox;
+ public List<MailboxIdentifier> getMailboxIds() {
+ return _mailboxIds;
}
@Override
public List<MultiStageOperator> getChildOperators() {
- return ImmutableList.of();
+ return Collections.emptyList();
}
@Override
public void close() {
super.close();
- for (MailboxIdentifier sendingMailbox : _sendingMailbox) {
- _mailboxService.releaseReceivingMailbox(sendingMailbox);
- }
+ releaseRemainingMailboxes();
}
@Override
public void cancel(Throwable t) {
super.cancel(t);
- for (MailboxIdentifier sendingMailbox : _sendingMailbox) {
- _mailboxService.releaseReceivingMailbox(sendingMailbox);
+ releaseRemainingMailboxes();
+ }
+
+ protected void releaseRemainingMailboxes() {
+ ReceivingMailbox<TransferableBlock> mailbox;
+ while ((mailbox = _mailboxes.poll()) != null) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
}
}
}
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 b2e70c1dcf..5afcfc2aa3 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
@@ -18,19 +18,15 @@
*/
package org.apache.pinot.query.runtime.operator;
+import com.google.common.annotations.VisibleForTesting;
import java.util.List;
-import java.util.Map;
-import javax.annotation.Nullable;
import org.apache.calcite.rel.RelDistribution;
import org.apache.pinot.common.exception.QueryException;
-import org.apache.pinot.query.mailbox.MailboxIdentifier;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
import org.apache.pinot.query.routing.VirtualServer;
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.slf4j.Logger;
-import org.slf4j.LoggerFactory;
/**
@@ -38,23 +34,22 @@ import org.slf4j.LoggerFactory;
* {@link MultiStageOperator#getNextBlock()}()} API.
*/
public class MailboxReceiveOperator extends BaseMailboxReceiveOperator {
- private static final Logger LOGGER = LoggerFactory.getLogger(MailboxReceiveOperator.class);
private static final String EXPLAIN_NAME = "MAILBOX_RECEIVE";
+ private TransferableBlock _errorBlock;
+
public MailboxReceiveOperator(OpChainExecutionContext context, RelDistribution.Type exchangeType, int senderStageId,
int receiverStageId) {
this(context, context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, senderStageId,
- receiverStageId, context.getTimeoutMs());
+ receiverStageId);
}
- // TODO: Move deadlineInNanoSeconds to OperatorContext.
- // TODO: Remove boxed timeoutMs value from here and use long deadlineMs from context.
+ @VisibleForTesting
public MailboxReceiveOperator(OpChainExecutionContext context, List<VirtualServer> sendingStageInstances,
- RelDistribution.Type exchangeType, int senderStageId, int receiverStageId, Long timeoutMs) {
- super(context, sendingStageInstances, exchangeType, senderStageId, receiverStageId, timeoutMs);
+ RelDistribution.Type exchangeType, int senderStageId, int receiverStageId) {
+ super(context, sendingStageInstances, exchangeType, senderStageId, receiverStageId);
}
- @Nullable
@Override
public String toExplainString() {
return EXPLAIN_NAME;
@@ -62,59 +57,53 @@ public class MailboxReceiveOperator extends BaseMailboxReceiveOperator {
@Override
protected TransferableBlock getNextBlock() {
- if (_upstreamErrorBlock != null) {
- return _upstreamErrorBlock;
- } else if (System.nanoTime() >= _deadlineTimestampNano) {
- return TransferableBlockUtils.getErrorTransferableBlock(QueryException.EXECUTION_TIMEOUT_ERROR);
+ if (_errorBlock != null) {
+ return _errorBlock;
+ }
+ if (System.currentTimeMillis() > _context.getDeadlineMs()) {
+ _errorBlock = TransferableBlockUtils.getErrorTransferableBlock(QueryException.EXECUTION_TIMEOUT_ERROR);
+ return _errorBlock;
}
- int startingIdx = _serverIdx;
- int openMailboxCount = 0;
- int eosMailboxCount = 0;
- // For all non-singleton distribution, we poll from every instance to check mailbox content.
- // TODO: Fix wasted CPU cycles on waiting for servers that are not supposed to give content.
- for (int i = 0; i < _sendingMailbox.size(); i++) {
- // this implements a round-robin mailbox iterator, so we don't starve any mailboxes
- _serverIdx = (startingIdx + i) % _sendingMailbox.size();
- MailboxIdentifier mailboxId = _sendingMailbox.get(_serverIdx);
+ // Poll from every mailbox in round-robin fashion:
+ // - Return the first content block
+ // - If no content block found but there are mailboxes not finished, return no-op block
+ // - If all content blocks are already returned, return end-of-stream block
+ int numMailboxes = _mailboxes.size();
+ for (int i = 0; i < numMailboxes; i++) {
+ ReceivingMailbox<TransferableBlock> mailbox = _mailboxes.remove();
+ if (mailbox.isClosed()) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ continue;
+ }
try {
- ReceivingMailbox<TransferableBlock> mailbox = _mailboxService.getReceivingMailbox(mailboxId);
- if (!mailbox.isClosed()) {
- openMailboxCount++;
- TransferableBlock block = mailbox.receive();
- // Get null block when pulling times out from mailbox.
- if (block != null) {
- if (block.isErrorBlock()) {
- _upstreamErrorBlock =
- TransferableBlockUtils.getErrorTransferableBlock(block.getDataBlock().getExceptions());
- return _upstreamErrorBlock;
- }
- if (!block.isEndOfStreamBlock()) {
- return block;
- } else {
- if (_opChainStats != null && !block.getResultMetadata().isEmpty()) {
- for (Map.Entry<String, OperatorStats> entry : block.getResultMetadata().entrySet()) {
- _opChainStats.getOperatorStatsMap().compute(entry.getKey(), (_key, _value) -> entry.getValue());
- }
- }
- eosMailboxCount++;
- }
+ TransferableBlock block = mailbox.receive();
+
+ // Release the mailbox when the block is end-of-stream
+ if (block != null && block.isEndOfStreamBlock()) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ if (block.isErrorBlock()) {
+ _errorBlock = block;
+ return _errorBlock;
}
+ _opChainStats.getOperatorStatsMap().putAll(block.getResultMetadata());
+ continue;
+ }
+
+ // Add the mailbox back to the queue if the block is not end-of-stream
+ _mailboxes.add(mailbox);
+ if (block != null) {
+ return block;
}
} catch (Exception e) {
- return TransferableBlockUtils.getErrorTransferableBlock(
- new RuntimeException(String.format("Error polling mailbox=%s", mailboxId), e));
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ _errorBlock = TransferableBlockUtils.getErrorTransferableBlock(
+ new RuntimeException("Caught exception while polling from mailbox: " + mailbox.getId(), e));
+ return _errorBlock;
}
}
- // there are two conditions in which we should return EOS: (1) there were
- // no mailboxes to open (this shouldn't happen because the second condition
- // should be hit first, but is defensive) (2) every mailbox that was opened
- // returned an EOS block. in every other scenario, there are mailboxes that
- // are not yet exhausted and we should wait for more data to be available
- TransferableBlock block =
- openMailboxCount > 0 && openMailboxCount > eosMailboxCount ? TransferableBlockUtils.getNoOpTransferableBlock()
- : TransferableBlockUtils.getEndOfStreamTransferableBlock();
- return block;
+ return _mailboxes.isEmpty() ? TransferableBlockUtils.getEndOfStreamTransferableBlock()
+ : TransferableBlockUtils.getNoOpTransferableBlock();
}
}
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 6b9f630eb5..84f0f995da 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
@@ -26,14 +26,15 @@ 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;
+import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
public abstract class MultiStageOperator implements Operator<TransferableBlock>, AutoCloseable {
- private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(MultiStageOperator.class);
+ private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageOperator.class);
- private final String _operatorId;
- private final OpChainExecutionContext _context;
+ protected final OpChainExecutionContext _context;
+ protected final String _operatorId;
protected final OpChainStats _opChainStats;
public MultiStageOperator(OpChainExecutionContext context) {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
index b565ed2d4c..3a7fc78a78 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java
@@ -21,8 +21,6 @@ package org.apache.pinot.query.runtime.operator;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.List;
-import java.util.Map;
-import java.util.PriorityQueue;
import javax.annotation.Nullable;
import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelFieldCollation;
@@ -30,7 +28,6 @@ import org.apache.commons.collections.CollectionUtils;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.exception.QueryException;
import org.apache.pinot.common.utils.DataSchema;
-import org.apache.pinot.query.mailbox.MailboxIdentifier;
import org.apache.pinot.query.mailbox.ReceivingMailbox;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.routing.VirtualServer;
@@ -38,8 +35,6 @@ import org.apache.pinot.query.runtime.blocks.TransferableBlock;
import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils;
import org.apache.pinot.query.runtime.operator.utils.SortUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
/**
@@ -50,43 +45,34 @@ import org.slf4j.LoggerFactory;
* resorting via the PriorityQueue.
*/
public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator {
- private static final Logger LOGGER = LoggerFactory.getLogger(SortedMailboxReceiveOperator.class);
private static final String EXPLAIN_NAME = "SORTED_MAILBOX_RECEIVE";
+ private final DataSchema _dataSchema;
private final List<RexExpression> _collationKeys;
private final List<RelFieldCollation.Direction> _collationDirections;
private final boolean _isSortOnSender;
- private final boolean _isSortOnReceiver;
- private final DataSchema _dataSchema;
- private final PriorityQueue<Object[]> _priorityQueue;
+ private final List<Object[]> _rows = new ArrayList<>();
+
+ private TransferableBlock _errorBlock;
private boolean _isSortedBlockConstructed;
public SortedMailboxReceiveOperator(OpChainExecutionContext context, RelDistribution.Type exchangeType,
- List<RexExpression> collationKeys, List<RelFieldCollation.Direction> collationDirections, boolean isSortOnSender,
- boolean isSortOnReceiver, DataSchema dataSchema, int senderStageId, int receiverStageId) {
- this(context, context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, collationKeys,
- collationDirections, isSortOnSender, isSortOnReceiver, dataSchema, senderStageId,
- receiverStageId, context.getTimeoutMs());
+ DataSchema dataSchema, List<RexExpression> collationKeys, List<RelFieldCollation.Direction> collationDirections,
+ boolean isSortOnSender, int senderStageId, int receiverStageId) {
+ this(context, context.getMetadataMap().get(senderStageId).getServerInstances(), exchangeType, dataSchema,
+ collationKeys, collationDirections, isSortOnSender, senderStageId, receiverStageId);
}
- // TODO: Move deadlineInNanoSeconds to OperatorContext.
- // TODO: Remove boxed timeoutMs value from here and use long deadlineMs from context.
public SortedMailboxReceiveOperator(OpChainExecutionContext context, List<VirtualServer> sendingStageInstances,
- RelDistribution.Type exchangeType, List<RexExpression> collationKeys,
- List<RelFieldCollation.Direction> collationDirections, boolean isSortOnSender, boolean isSortOnReceiver,
- DataSchema dataSchema, int senderStageId, int receiverStageId, Long timeoutMs) {
- super(context, sendingStageInstances, exchangeType, senderStageId, receiverStageId, timeoutMs);
+ RelDistribution.Type exchangeType, DataSchema dataSchema, List<RexExpression> collationKeys,
+ List<RelFieldCollation.Direction> collationDirections, boolean isSortOnSender, int senderStageId,
+ int receiverStageId) {
+ super(context, sendingStageInstances, exchangeType, senderStageId, receiverStageId);
+ Preconditions.checkState(!CollectionUtils.isEmpty(collationKeys), "Collation keys must be set");
+ _dataSchema = dataSchema;
_collationKeys = collationKeys;
_collationDirections = collationDirections;
_isSortOnSender = isSortOnSender;
- _isSortOnReceiver = isSortOnReceiver;
- _dataSchema = dataSchema;
- Preconditions.checkState(!CollectionUtils.isEmpty(collationKeys) && isSortOnReceiver,
- "Collation keys should exist and sorting must be enabled otherwise use non-sorted MailboxReceiveOperator");
- // No need to switch the direction since all rows will be stored in the priority queue without applying limits
- _priorityQueue = new PriorityQueue<>(new SortUtils.SortComparator(collationKeys, collationDirections,
- dataSchema, false, false));
- _isSortedBlockConstructed = false;
}
@Nullable
@@ -97,104 +83,70 @@ public class SortedMailboxReceiveOperator extends BaseMailboxReceiveOperator {
@Override
protected TransferableBlock getNextBlock() {
- if (_upstreamErrorBlock != null) {
- return _upstreamErrorBlock;
- } else if (System.nanoTime() >= _deadlineTimestampNano) {
- return TransferableBlockUtils.getErrorTransferableBlock(QueryException.EXECUTION_TIMEOUT_ERROR);
+ if (_errorBlock != null) {
+ return _errorBlock;
+ }
+ if (System.currentTimeMillis() > _context.getDeadlineMs()) {
+ _errorBlock = TransferableBlockUtils.getErrorTransferableBlock(QueryException.EXECUTION_TIMEOUT_ERROR);
+ return _errorBlock;
}
- int startingIdx = _serverIdx;
- int openMailboxCount = 0;
- int eosMailboxCount = 0;
- boolean foundNonNullTransferableBlock;
- // For all non-singleton distribution, we poll from every instance to check mailbox content.
- // Since this operator needs to wait for all incoming data before it can send the data to the
- // upstream operators, this operator will keep trying to poll from all mailboxes until it only
- // receives null blocks or EOS for all mailboxes. This operator does not need to yield itself
- // for backpressure at the moment but once support is added for k-way merge when input data is
- // sorted this operator will have to return some data blocks without waiting for all the data.
- // TODO: Fix wasted CPU cycles on waiting for servers that are not supposed to give content.
- do {
- // Reset the following for each loop
- foundNonNullTransferableBlock = false;
- openMailboxCount = 0;
- eosMailboxCount = 0;
- for (int i = 0; i < _sendingMailbox.size(); i++) {
- // this implements a round-robin mailbox iterator, so we don't starve any mailboxes
- _serverIdx = (startingIdx + i) % _sendingMailbox.size();
- MailboxIdentifier mailboxId = _sendingMailbox.get(_serverIdx);
- try {
- ReceivingMailbox<TransferableBlock> mailbox = _mailboxService.getReceivingMailbox(mailboxId);
- if (!mailbox.isClosed()) {
- openMailboxCount++;
- TransferableBlock block = mailbox.receive();
- // Get null block when pulling times out from mailbox.
- if (block != null) {
- foundNonNullTransferableBlock = true;
- if (block.isErrorBlock()) {
- _upstreamErrorBlock =
- TransferableBlockUtils.getErrorTransferableBlock(block.getDataBlock().getExceptions());
- return _upstreamErrorBlock;
- }
- if (!block.isEndOfStreamBlock()) {
- // Add rows to the PriorityQueue to order them
- List<Object[]> container = block.getContainer();
- _priorityQueue.addAll(container);
- } else {
- if (_opChainStats != null && !block.getResultMetadata().isEmpty()) {
- for (Map.Entry<String, OperatorStats> entry : block.getResultMetadata().entrySet()) {
- _opChainStats.getOperatorStatsMap().compute(entry.getKey(), (_key, _value) -> entry.getValue());
- }
- }
- eosMailboxCount++;
- }
- }
+ // Since this operator needs to wait for all incoming data before it can send the data to the upstream operators,
+ // this operator will keep polling from all mailboxes until it receives null block or all blocks are collected.
+ // TODO: Use k-way merge when input data is sorted, and return blocks without waiting for all the data.
+ while (!_mailboxes.isEmpty()) {
+ ReceivingMailbox<TransferableBlock> mailbox = _mailboxes.remove();
+ if (mailbox.isClosed()) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ continue;
+ }
+ try {
+ TransferableBlock block = mailbox.receive();
+
+ // Release the mailbox when the block is end-of-stream
+ if (block != null && block.isEndOfStreamBlock()) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ if (block.isErrorBlock()) {
+ _errorBlock = block;
+ return _errorBlock;
}
- } catch (Exception e) {
- return TransferableBlockUtils.getErrorTransferableBlock(
- new RuntimeException(String.format("Error polling mailbox=%s", mailboxId), e));
+ _opChainStats.getOperatorStatsMap().putAll(block.getResultMetadata());
+ continue;
}
- }
- } while (foundNonNullTransferableBlock && ((openMailboxCount > 0) && (openMailboxCount > eosMailboxCount)));
- if (((openMailboxCount == 0) || (openMailboxCount == eosMailboxCount))
- && (!CollectionUtils.isEmpty(_priorityQueue)) && !_isSortedBlockConstructed) {
- // Some data is present in the PriorityQueue, these need to be sent upstream
- List<Object[]> rows = new ArrayList<>();
- while (_priorityQueue.size() > 0) {
- Object[] row = _priorityQueue.poll();
- rows.add(row);
+ // Add the mailbox back to the queue if the block is not end-of-stream
+ _mailboxes.add(mailbox);
+ if (block != null) {
+ _rows.addAll(block.getContainer());
+ } else {
+ return TransferableBlockUtils.getNoOpTransferableBlock();
+ }
+ } catch (Exception e) {
+ _mailboxService.releaseReceivingMailbox(mailbox);
+ _errorBlock = TransferableBlockUtils.getErrorTransferableBlock(
+ new RuntimeException("Caught exception while polling from mailbox: " + mailbox.getId(), e));
+ return _errorBlock;
}
- _isSortedBlockConstructed = true;
- return new TransferableBlock(rows, _dataSchema, DataBlock.Type.ROW);
}
- // there are two conditions in which we should return EOS: (1) there were
- // no mailboxes to open (this shouldn't happen because the second condition
- // should be hit first, but is defensive) (2) every mailbox that was opened
- // returned an EOS block. in every other scenario, there are mailboxes that
- // are not yet exhausted and we should wait for more data to be available
- TransferableBlock block =
- openMailboxCount > 0 && openMailboxCount > eosMailboxCount ? TransferableBlockUtils.getNoOpTransferableBlock()
- : TransferableBlockUtils.getEndOfStreamTransferableBlock();
- return block;
- }
-
- private void cleanUpResources() {
- if (_priorityQueue != null) {
- _priorityQueue.clear();
+ if (!_isSortedBlockConstructed && !_rows.isEmpty()) {
+ _rows.sort(new SortUtils.SortComparator(_collationKeys, _collationDirections, _dataSchema, false, false));
+ _isSortedBlockConstructed = true;
+ return new TransferableBlock(_rows, _dataSchema, DataBlock.Type.ROW);
+ } else {
+ return TransferableBlockUtils.getEndOfStreamTransferableBlock();
}
}
@Override
public void close() {
super.close();
- cleanUpResources();
+ _rows.clear();
}
@Override
public void cancel(Throwable t) {
super.cancel(t);
- cleanUpResources();
+ _rows.clear();
}
}
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 c323a6074d..d8f9a8c7e0 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
@@ -66,15 +66,15 @@ public class PhysicalPlanVisitor implements StageNodeVisitor<MultiStageOperator,
if (node.isSortOnReceiver()) {
SortedMailboxReceiveOperator sortedMailboxReceiveOperator =
new SortedMailboxReceiveOperator(context.getOpChainExecutionContext(), node.getExchangeType(),
- node.getCollationKeys(), node.getCollationDirections(), node.isSortOnSender(), node.isSortOnReceiver(),
- node.getDataSchema(), node.getSenderStageId(), node.getStageId());
- context.addReceivingMailboxes(sortedMailboxReceiveOperator.getSendingMailbox());
+ node.getDataSchema(), node.getCollationKeys(), node.getCollationDirections(), node.isSortOnSender(),
+ node.getSenderStageId(), node.getStageId());
+ context.addReceivingMailboxes(sortedMailboxReceiveOperator.getMailboxIds());
return sortedMailboxReceiveOperator;
} else {
MailboxReceiveOperator mailboxReceiveOperator =
new MailboxReceiveOperator(context.getOpChainExecutionContext(), node.getExchangeType(),
node.getSenderStageId(), node.getStageId());
- context.addReceivingMailboxes(mailboxReceiveOperator.getSendingMailbox());
+ context.addReceivingMailboxes(mailboxReceiveOperator.getMailboxIds());
return mailboxReceiveOperator;
}
}
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 6d113c9fa2..2d0ab6c275 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
@@ -182,8 +182,8 @@ public class QueryDispatcher {
VirtualServerAddress server =
new VirtualServerAddress(mailboxService.getHostname(), mailboxService.getMailboxPort(), 0);
OpChainExecutionContext context =
- new OpChainExecutionContext(mailboxService, requestId, reduceStageId, server, timeoutMs, timeoutMs,
- queryPlan.getStageMetadataMap(), traceEnabled);
+ new OpChainExecutionContext(mailboxService, requestId, reduceStageId, server, timeoutMs,
+ System.currentTimeMillis() + timeoutMs, queryPlan.getStageMetadataMap(), traceEnabled);
MailboxReceiveOperator mailboxReceiveOperator =
createReduceStageOperator(reduceNode.getSenderStageId(), reduceStageId, reduceNode.getDataSchema(), context);
List<DataBlock> resultDataBlocks =
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 cb67b22127..015800111c 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
@@ -18,11 +18,9 @@
*/
package org.apache.pinot.query.runtime.operator;
-import com.google.common.collect.ImmutableList;
-import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.List;
-import java.util.Random;
import org.apache.calcite.rel.RelDistribution;
import org.apache.pinot.common.datablock.MetadataBlock;
import org.apache.pinot.common.exception.QueryException;
@@ -36,42 +34,46 @@ 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;
-import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
public class MailboxReceiveOperatorTest {
-
- private static final int DEFAULT_RECEIVER_STAGE_ID = 10;
+ private static final VirtualServerAddress RECEIVER_ADDRESS = new VirtualServerAddress("localhost", 123, 0);
+ private static final DataSchema DATA_SCHEMA =
+ new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
private AutoCloseable _mocks;
-
- @Mock
- private ReceivingMailbox<TransferableBlock> _mailbox;
-
- @Mock
- private ReceivingMailbox<TransferableBlock> _mailbox2;
-
@Mock
private MailboxService<TransferableBlock> _mailboxService;
@Mock
private VirtualServer _server1;
@Mock
private VirtualServer _server2;
-
- private final VirtualServerAddress _testAddr = new VirtualServerAddress("test", 123, 0);
-
- private final Random _random = new Random();
+ @Mock
+ private ReceivingMailbox<TransferableBlock> _mailbox1;
+ @Mock
+ private ReceivingMailbox<TransferableBlock> _mailbox2;
@BeforeMethod
public void setUp() {
_mocks = MockitoAnnotations.openMocks(this);
+ when(_mailboxService.getHostname()).thenReturn("localhost");
+ when(_mailboxService.getMailboxPort()).thenReturn(123);
+ when(_server1.getHostname()).thenReturn("localhost");
+ when(_server1.getQueryMailboxPort()).thenReturn(123);
+ when(_server1.getVirtualId()).thenReturn(0);
+ when(_server2.getHostname()).thenReturn("localhost");
+ when(_server2.getQueryMailboxPort()).thenReturn(123);
+ when(_server2.getVirtualId()).thenReturn(1);
}
@AfterMethod
@@ -80,564 +82,319 @@ public class MailboxReceiveOperatorTest {
_mocks.close();
}
- @Test
- public void shouldTimeoutOnExtraLongSleep()
- throws InterruptedException {
- // shorter timeoutMs should result in error.
+ @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Failed to find instance.*")
+ public void shouldThrowSingletonNoMatchMailboxServer() {
+ when(_server1.getQueryMailboxPort()).thenReturn(456);
+ when(_server2.getQueryMailboxPort()).thenReturn(789);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 10L, 10L,
- new HashMap<>(), false);
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, 10L);
- Thread.sleep(200L);
- TransferableBlock mailbox = receiveOp.nextBlock();
- Assert.assertTrue(mailbox.isErrorBlock());
- MetadataBlock errorBlock = (MetadataBlock) mailbox.getDataBlock();
- Assert.assertTrue(errorBlock.getExceptions().containsKey(QueryException.EXECUTION_TIMEOUT_ERROR_CODE));
-
- // 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<>(), false);
- 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<>(), false);
- receiveOp = new MailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, 456, 789, null);
- Thread.sleep(200L);
- mailbox = receiveOp.nextBlock();
- Assert.assertFalse(mailbox.isErrorBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new MailboxReceiveOperator(context, Arrays.asList(_server1, _server2), RelDistribution.Type.SINGLETON, 1, 0);
}
- @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*multiple instance "
- + "found.*")
+ @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Multiple instances.*")
public void shouldThrowReceiveSingletonFromMultiMatchMailboxServer() {
-
- Mockito.when(_mailboxService.getHostname()).thenReturn("singleton");
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(123);
-
- Mockito.when(_server1.getHostname()).thenReturn("singleton");
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(123);
-
- Mockito.when(_server2.getHostname()).thenReturn("singleton");
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- 456, 789, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new MailboxReceiveOperator(context, Arrays.asList(_server1, _server2), RelDistribution.Type.SINGLETON, 1, 0);
}
@Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*RANGE_DISTRIBUTED.*")
public void shouldThrowRangeDistributionNotSupported() {
- Mockito.when(_mailboxService.getHostname()).thenReturn("singleton");
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(123);
-
- Mockito.when(_server1.getHostname()).thenReturn("singleton");
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(123);
-
- Mockito.when(_server2.getHostname()).thenReturn("singleton");
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
- MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.RANGE_DISTRIBUTED, 456, 789, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new MailboxReceiveOperator(context, Collections.emptyList(), RelDistribution.Type.RANGE_DISTRIBUTED, 1, 0);
}
@Test
- public void shouldReceiveSingletonNoMatchMailboxServer() {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = 789;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0);
-
+ public void shouldTimeoutOnExtraLongSleep()
+ throws InterruptedException {
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+
+ // Short timeoutMs should result in timeout
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, 10L, System.currentTimeMillis() + 10L,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, 1, 0)) {
+ Thread.sleep(100L);
+ TransferableBlock mailbox = receiveOp.nextBlock();
+ assertTrue(mailbox.isErrorBlock());
+ MetadataBlock errorBlock = (MetadataBlock) mailbox.getDataBlock();
+ assertTrue(errorBlock.getExceptions().containsKey(QueryException.EXECUTION_TIMEOUT_ERROR_CODE));
+ }
- // Receive end of stream block directly when there is no match.
- Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ // Longer timeout or default timeout (10s) doesn't result in timeout
+ context = new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, 10_000L,
+ System.currentTimeMillis() + 10_000L, Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, 1, 0)) {
+ Thread.sleep(100L);
+ TransferableBlock mailbox = receiveOp.nextBlock();
+ assertFalse(mailbox.isErrorBlock());
+ }
}
@Test
public void shouldReceiveSingletonCloseMailbox() {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(true);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ when(_mailbox1.isClosed()).thenReturn(true);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- 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());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ }
}
@Test
- public void shouldReceiveSingletonNullMailbox()
- throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- // Receive null mailbox during timeout.
- Mockito.when(_mailbox.receive()).thenReturn(null);
+ public void shouldReceiveSingletonNullMailbox() {
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
- // Receive NoOpBlock.
- Assert.assertTrue(receiveOp.nextBlock().isNoOpBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ assertTrue(receiveOp.nextBlock().isNoOpBlock());
+ }
}
@Test
public void shouldReceiveEosDirectlyFromSender()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ when(_mailbox1.receive()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
- // Receive EosBloc.
- Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ }
}
@Test
public void shouldReceiveSingletonMailbox()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox1.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldReceiveSingletonErrorMailbox()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Exception e = new Exception("errorBlock");
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
- Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getErrorTransferableBlock(e));
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenReturn(
+ TransferableBlockUtils.getErrorTransferableBlock(new RuntimeException(errorMessage)));
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- 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();
- Assert.assertTrue(error.getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains("errorBlock"));
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
@Test
public void shouldReceiveMailboxFromTwoServersOneClose()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- 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);
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ when(_mailbox1.isClosed()).thenReturn(true);
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Arrays.asList(_server1, _server2),
+ RelDistribution.Type.HASH_DISTRIBUTED, senderStageId, 0)) {
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldReceiveMailboxFromTwoServersOneNull()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Mockito.when(_mailbox.receive()).thenReturn(null, TransferableBlockUtils.getEndOfStreamTransferableBlock());
-
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ when(_mailbox1.receive()).thenReturn(null, TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- new MailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED,
- stageId, DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Arrays.asList(_server1, _server2),
+ RelDistribution.Type.HASH_DISTRIBUTED, senderStageId, 0)) {
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldReceiveMailboxFromTwoServers()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Object[] expRow1 = new Object[]{1, 1};
- Object[] expRow2 = new Object[]{2, 2};
- Mockito.when(_mailbox.receive())
- .thenReturn(OperatorTestUtil.block(inSchema, expRow1), OperatorTestUtil.block(inSchema, expRow2),
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2);
- Mockito.when(_mailbox2.isClosed()).thenReturn(false);
- Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3));
+ long requestId = 0;
+ int senderStageId = 1;
+ Object[] row1 = new Object[]{1, 1};
+ Object[] row2 = new Object[]{2, 2};
+ Object[] row3 = new Object[]{3, 3};
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ when(_mailbox1.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row1),
+ OperatorTestUtil.block(DATA_SCHEMA, row3), TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row2),
+ TransferableBlockUtils.getEndOfStreamTransferableBlock());
+
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- 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<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow1);
- // Receive second block from first server.
- receivedBlock = receiveOp.nextBlock();
- resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow2);
-
- // Receive from second server.
- receivedBlock = receiveOp.nextBlock();
- resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow3);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Arrays.asList(_server1, _server2),
+ RelDistribution.Type.HASH_DISTRIBUTED, senderStageId, 0)) {
+ // Receive first block from server1
+ assertEquals(receiveOp.nextBlock().getContainer().get(0), row1);
+ // Receive second block from server2
+ assertEquals(receiveOp.nextBlock().getContainer().get(0), row2);
+ // Receive third block from server1
+ assertEquals(receiveOp.nextBlock().getContainer().get(0), row3);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ }
}
@Test
public void shouldGetReceptionReceiveErrorMailbox()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2);
- Mockito.when(_mailbox2.isClosed()).thenReturn(false);
- Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3));
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenReturn(
+ TransferableBlockUtils.getErrorTransferableBlock(new RuntimeException(errorMessage)));
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{3, 3};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
+ TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- 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());
- MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock();
- Assert.assertTrue(error.getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains("mailboxError"));
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
@Test
public void shouldThrowReceiveWhenOneServerReceiveThrowException()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2);
- Mockito.when(_mailbox2.isClosed()).thenReturn(false);
- Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3));
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenThrow(new Exception(errorMessage));
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{3, 3};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
+ TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- MailboxReceiveOperator receiveOp =
- 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");
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (MailboxReceiveOperator receiveOp = new MailboxReceiveOperator(context, Collections.singletonList(_server1),
+ RelDistribution.Type.SINGLETON, senderStageId, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
}
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
index 56ff2b36e3..248b43e89f 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/OpChainTest.java
@@ -77,13 +77,11 @@ public class OpChainTest {
@Mock
private ReceivingMailbox<TransferableBlock> _mailbox;
-
@Mock
private MailboxService<TransferableBlock> _mailboxService2;
@Mock
private ReceivingMailbox<TransferableBlock> _mailbox2;
-
@BeforeMethod
public void setUp() {
_mocks = MockitoAnnotations.openMocks(this);
@@ -198,8 +196,8 @@ public class OpChainTest {
stageMetadata.setServerInstances(ImmutableList.of(_server));
Map<Integer, StageMetadata> stageMetadataMap = Collections.singletonMap(receivedStageId, stageMetadata);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, senderStageId, new VirtualServerAddress(_server), 1000, 1000,
- stageMetadataMap, true);
+ new OpChainExecutionContext(_mailboxService, 1, senderStageId, new VirtualServerAddress(_server), 1000,
+ System.currentTimeMillis() + 1000, stageMetadataMap, true);
Stack<MultiStageOperator> operators =
getFullOpchain(receivedStageId, senderStageId, context, dummyOperatorWaitTime);
@@ -213,11 +211,11 @@ public class OpChainTest {
OpChainExecutionContext secondStageContext =
new OpChainExecutionContext(_mailboxService2, 1, senderStageId + 1, new VirtualServerAddress(_server), 1000,
- 1000, stageMetadataMap, true);
+ System.currentTimeMillis() + 1000, stageMetadataMap, true);
MailboxReceiveOperator secondStageReceiveOp =
new MailboxReceiveOperator(secondStageContext, ImmutableList.of(_server),
- RelDistribution.Type.BROADCAST_DISTRIBUTED, senderStageId, receivedStageId + 1, null);
+ RelDistribution.Type.BROADCAST_DISTRIBUTED, senderStageId, receivedStageId + 1);
Assert.assertTrue(opChain.getStats().getExecutionTime() >= dummyOperatorWaitTime);
int numOperators = operators.size();
@@ -242,8 +240,8 @@ public class OpChainTest {
stageMetadata.setServerInstances(ImmutableList.of(_server));
Map<Integer, StageMetadata> stageMetadataMap = Collections.singletonMap(receivedStageId, stageMetadata);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, senderStageId, new VirtualServerAddress(_server), 1000, 1000,
- stageMetadataMap, false);
+ new OpChainExecutionContext(_mailboxService, 1, senderStageId, new VirtualServerAddress(_server), 1000,
+ System.currentTimeMillis() + 1000, stageMetadataMap, false);
Stack<MultiStageOperator> operators =
getFullOpchain(receivedStageId, senderStageId, context, dummyOperatorWaitTime);
@@ -255,11 +253,11 @@ public class OpChainTest {
OpChainExecutionContext secondStageContext =
new OpChainExecutionContext(_mailboxService2, 1, senderStageId + 1, new VirtualServerAddress(_server), 1000,
- 1000, stageMetadataMap, false);
+ System.currentTimeMillis() + 1000, stageMetadataMap, false);
MailboxReceiveOperator secondStageReceiveOp =
new MailboxReceiveOperator(secondStageContext, ImmutableList.of(_server),
- RelDistribution.Type.BROADCAST_DISTRIBUTED, senderStageId, receivedStageId + 1, null);
+ RelDistribution.Type.BROADCAST_DISTRIBUTED, senderStageId, receivedStageId + 1);
Assert.assertTrue(opChain.getStats().getExecutionTime() >= dummyOperatorWaitTime);
Assert.assertEquals(opChain.getStats().getOperatorStatsMap().size(), 2);
@@ -324,7 +322,6 @@ public class OpChainTest {
return operators;
}
-
static class DummyMultiStageOperator extends MultiStageOperator {
public DummyMultiStageOperator(OpChainExecutionContext context) {
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
index e306875cf3..343904f218 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java
@@ -18,13 +18,9 @@
*/
package org.apache.pinot.query.runtime.operator;
-import com.google.common.collect.ImmutableList;
-import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
-import java.util.HashMap;
import java.util.List;
-import java.util.Random;
import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.pinot.common.datablock.MetadataBlock;
@@ -40,43 +36,50 @@ 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;
-import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT;
import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
public class SortedMailboxReceiveOperatorTest {
-
- private static final int DEFAULT_RECEIVER_STAGE_ID = 10;
+ private static final VirtualServerAddress RECEIVER_ADDRESS = new VirtualServerAddress("localhost", 123, 0);
+ private static final DataSchema DATA_SCHEMA =
+ new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
+ private static final List<RexExpression> COLLATION_KEYS = Collections.singletonList(new RexExpression.InputRef(0));
+ private static final List<RelFieldCollation.Direction> COLLATION_DIRECTIONS =
+ Collections.singletonList(RelFieldCollation.Direction.ASCENDING);
private AutoCloseable _mocks;
-
- @Mock
- private ReceivingMailbox<TransferableBlock> _mailbox;
-
- @Mock
- private ReceivingMailbox<TransferableBlock> _mailbox2;
-
@Mock
private MailboxService<TransferableBlock> _mailboxService;
@Mock
private VirtualServer _server1;
@Mock
private VirtualServer _server2;
-
- private final VirtualServerAddress _testAddr = new VirtualServerAddress("test", 123, 0);
-
- private final Random _random = new Random();
+ @Mock
+ private ReceivingMailbox<TransferableBlock> _mailbox1;
+ @Mock
+ private ReceivingMailbox<TransferableBlock> _mailbox2;
@BeforeMethod
public void setUp() {
_mocks = MockitoAnnotations.openMocks(this);
+ when(_mailboxService.getHostname()).thenReturn("localhost");
+ when(_mailboxService.getMailboxPort()).thenReturn(123);
+ when(_server1.getHostname()).thenReturn("localhost");
+ when(_server1.getQueryMailboxPort()).thenReturn(123);
+ when(_server1.getVirtualId()).thenReturn(0);
+ when(_server2.getHostname()).thenReturn("localhost");
+ when(_server2.getQueryMailboxPort()).thenReturn(123);
+ when(_server2.getVirtualId()).thenReturn(1);
}
@AfterMethod
@@ -85,832 +88,389 @@ public class SortedMailboxReceiveOperatorTest {
_mocks.close();
}
- @Test
- public void shouldTimeoutOnExtraLongSleep()
- throws InterruptedException {
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
- // shorter timeoutMs should result in error.
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
+ @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Failed to find instance.*")
+ public void shouldThrowSingletonNoMatchMailboxServer() {
+ when(_server1.getQueryMailboxPort()).thenReturn(456);
+ when(_server2.getQueryMailboxPort()).thenReturn(789);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, 10L, 10L,
- new HashMap<>(), false);
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON, collationKeys,
- collationDirections, false, true, inSchema, 456, 789, 10L);
- Thread.sleep(200L);
- TransferableBlock mailbox = receiveOp.nextBlock();
- Assert.assertTrue(mailbox.isErrorBlock());
- MetadataBlock errorBlock = (MetadataBlock) mailbox.getDataBlock();
- Assert.assertTrue(errorBlock.getExceptions().containsKey(QueryException.EXECUTION_TIMEOUT_ERROR_CODE));
-
- // 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<>(), false);
- receiveOp = new SortedMailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, 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<>(), false);
- receiveOp = new SortedMailboxReceiveOperator(context, new ArrayList<>(), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, 456, 789, null);
- Thread.sleep(200L);
- mailbox = receiveOp.nextBlock();
- Assert.assertFalse(mailbox.isErrorBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new SortedMailboxReceiveOperator(context, Arrays.asList(_server1, _server2), RelDistribution.Type.SINGLETON,
+ DATA_SCHEMA, COLLATION_KEYS, COLLATION_DIRECTIONS, false, 1, 0);
}
- @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*multiple instance "
- + "found.*")
+ @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Multiple instances.*")
public void shouldThrowReceiveSingletonFromMultiMatchMailboxServer() {
-
- Mockito.when(_mailboxService.getHostname()).thenReturn("singleton");
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(123);
-
- Mockito.when(_server1.getHostname()).thenReturn("singleton");
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(123);
-
- Mockito.when(_server2.getHostname()).thenReturn("singleton");
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, 456, 789, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new SortedMailboxReceiveOperator(context, Arrays.asList(_server1, _server2), RelDistribution.Type.SINGLETON,
+ DATA_SCHEMA, COLLATION_KEYS, COLLATION_DIRECTIONS, false, 1, 0);
}
@Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*RANGE_DISTRIBUTED.*")
public void shouldThrowRangeDistributionNotSupported() {
- Mockito.when(_mailboxService.getHostname()).thenReturn("singleton");
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(123);
-
- Mockito.when(_server1.getHostname()).thenReturn("singleton");
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(123);
-
- Mockito.when(_server2.getHostname()).thenReturn("singleton");
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(123);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
+ OpChainExecutionContext context =
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new SortedMailboxReceiveOperator(context, Collections.emptyList(), RelDistribution.Type.RANGE_DISTRIBUTED,
+ DATA_SCHEMA, COLLATION_KEYS, COLLATION_DIRECTIONS, false, 1, 0);
+ }
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ @Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Collation keys.*")
+ public void shouldThrowOnEmptyCollationKey() {
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, 1, DEFAULT_RECEIVER_STAGE_ID, _testAddr, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
- SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
- ImmutableList.of(_server1, _server2), RelDistribution.Type.RANGE_DISTRIBUTED, collationKeys,
- collationDirections, false, true, inSchema, 456, 789, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, 10L, System.currentTimeMillis() + 10L,
+ Collections.emptyMap(), false);
+ //noinspection resource
+ new SortedMailboxReceiveOperator(context, Collections.singletonList(_server1), RelDistribution.Type.SINGLETON,
+ DATA_SCHEMA, Collections.emptyList(), Collections.emptyList(), false, 1, 0);
}
@Test
- public void shouldReceiveSingletonNoMatchMailboxServer() {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = 789;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
+ public void shouldTimeoutOnExtraLongSleep()
+ throws InterruptedException {
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+
+ // Short timeoutMs should result in timeout
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, 10L, System.currentTimeMillis() + 10L,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ Thread.sleep(100L);
+ TransferableBlock mailbox = receiveOp.nextBlock();
+ assertTrue(mailbox.isErrorBlock());
+ MetadataBlock errorBlock = (MetadataBlock) mailbox.getDataBlock();
+ assertTrue(errorBlock.getExceptions().containsKey(QueryException.EXECUTION_TIMEOUT_ERROR_CODE));
+ }
- // Receive end of stream block directly when there is no match.
- Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ // Longer timeout or default timeout (10s) doesn't result in timeout
+ context = new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, 10_000L,
+ System.currentTimeMillis() + 10_000L, Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ Thread.sleep(100L);
+ TransferableBlock mailbox = receiveOp.nextBlock();
+ assertFalse(mailbox.isErrorBlock());
+ }
}
@Test
public void shouldReceiveSingletonCloseMailbox() {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(true);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ when(_mailbox1.isClosed()).thenReturn(true);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
-
- // Receive end of stream block directly when mailbox is close.
- Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ }
}
@Test
- public void shouldReceiveSingletonNullMailbox()
- throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- // Receive null mailbox during timeout.
- Mockito.when(_mailbox.receive()).thenReturn(null);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ public void shouldReceiveSingletonNullMailbox() {
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
- // Receive NoOpBlock.
- Assert.assertTrue(receiveOp.nextBlock().isNoOpBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ assertTrue(receiveOp.nextBlock().isNoOpBlock());
+ }
}
@Test
public void shouldReceiveEosDirectlyFromSender()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ when(_mailbox1.receive()).thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
- // Receive EosBloc.
- Assert.assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
+ }
}
@Test
public void shouldReceiveSingletonMailbox()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox1.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldReceiveSingletonErrorMailbox()
throws Exception {
- String serverHost = "singleton";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(serverHost);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- int server2port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(serverHost);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2port);
-
- int mailboxPort = server2port;
- Mockito.when(_mailboxService.getHostname()).thenReturn(serverHost);
- Mockito.when(_mailboxService.getMailboxPort()).thenReturn(mailboxPort);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Exception e = new Exception("errorBlock");
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
- Mockito.when(_mailbox.receive()).thenReturn(TransferableBlockUtils.getErrorTransferableBlock(e));
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
- OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2), RelDistribution.Type.SINGLETON,
- collationKeys, collationDirections, false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID,
- null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- Assert.assertTrue(receivedBlock.isErrorBlock());
- MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock();
- Assert.assertTrue(error.getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains("errorBlock"));
- }
-
- @Test(expectedExceptions = IllegalStateException.class,
- expectedExceptionsMessageRegExp = ".*Collation keys should exist.*")
- public void shouldThrowOnEmptyCollationKey()
- throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
-
- OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, true, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
- }
-
- @Test(expectedExceptions = IllegalStateException.class,
- expectedExceptionsMessageRegExp = ".*sorting must be enabled.*")
- public void shouldThrowOnShouldSortOnReceiverFalse()
- throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2"}, new DataSchema.ColumnDataType[]{INT, INT});
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId =
+ new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId), RECEIVER_ADDRESS, RECEIVER_ADDRESS,
+ senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenReturn(
+ TransferableBlockUtils.getErrorTransferableBlock(new RuntimeException(errorMessage)));
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, false, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Collections.singletonList(_server1), RelDistribution.Type.SINGLETON, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
@Test
public void shouldReceiveMailboxFromTwoServersOneClose()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- 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);
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ when(_mailbox1.isClosed()).thenReturn(true);
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, true, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldReceiveMailboxFromTwoServersOneNull()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Mockito.when(_mailbox.receive()).thenReturn(null, TransferableBlockUtils.getEndOfStreamTransferableBlock());
-
- 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),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ when(_mailbox1.receive()).thenReturn(null, TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{1, 1};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, true, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ assertTrue(receiveOp.nextBlock().isNoOpBlock());
+ List<Object[]> actualRows = receiveOp.nextBlock().getContainer();
+ assertEquals(actualRows.size(), 1);
+ assertEquals(actualRows.get(0), row);
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- Assert.assertEquals(resultRows.size(), 1);
- Assert.assertEquals(resultRows.get(0), expRow);
}
@Test
public void shouldGetReceptionReceiveErrorMailbox()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2);
- Mockito.when(_mailbox2.isClosed()).thenReturn(false);
- Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3));
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenReturn(
+ TransferableBlockUtils.getErrorTransferableBlock(new RuntimeException(errorMessage)));
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{3, 3};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
+ TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, true, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
- // Receive error block from first server.
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- Assert.assertTrue(receivedBlock.isErrorBlock());
- MetadataBlock error = (MetadataBlock) receivedBlock.getDataBlock();
- Assert.assertTrue(error.getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains("mailboxError"));
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
@Test
public void shouldThrowReceiveWhenOneServerReceiveThrowException()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId2)).thenReturn(_mailbox2);
- Mockito.when(_mailbox2.isClosed()).thenReturn(false);
- Mockito.when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(inSchema, expRow3));
-
- List<RexExpression> collationKeys = new ArrayList<>();
- List<RelFieldCollation.Direction> collationDirections = new ArrayList<>();
- collationKeys.add(new RexExpression.InputRef(0));
- collationDirections.add(RelFieldCollation.Direction.ASCENDING);
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ String errorMessage = "TEST ERROR";
+ when(_mailbox1.receive()).thenThrow(new Exception(errorMessage));
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row = new Object[]{3, 3};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row),
+ TransferableBlockUtils.getEndOfStreamTransferableBlock());
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp =
- new SortedMailboxReceiveOperator(context, ImmutableList.of(_server1, _server2),
- RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirections, false, true, inSchema, stageId,
- DEFAULT_RECEIVER_STAGE_ID, null);
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- Assert.assertTrue(receivedBlock.isErrorBlock(), "server-1 should have returned an error-block");
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ TransferableBlock block = receiveOp.nextBlock();
+ assertTrue(block.isErrorBlock());
+ assertTrue(block.getDataBlock().getExceptions().get(QueryException.UNKNOWN_ERROR_CODE).contains(errorMessage));
+ }
}
@Test
public void shouldReceiveMailboxFromTwoServersWithCollationKey()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- 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);
- Mockito.when(_mailboxService.getReceivingMailbox(expectedMailboxId1)).thenReturn(_mailbox);
- Mockito.when(_mailbox.isClosed()).thenReturn(false);
- Object[] expRow1 = new Object[]{3, 3};
- Object[] expRow2 = new Object[]{1, 1};
- Mockito.when(_mailbox.receive())
- .thenReturn(OperatorTestUtil.block(inSchema, expRow1), OperatorTestUtil.block(inSchema, expRow2),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
-
- Object[] expRow3 = new Object[]{4, 2};
- Object[] expRow4 = new Object[]{2, 4};
- Object[] expRow5 = new Object[]{-1, 95};
- 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),
- OperatorTestUtil.block(inSchema, expRow4), OperatorTestUtil.block(inSchema, expRow5),
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ Object[] row1 = new Object[]{3, 3};
+ Object[] row2 = new Object[]{1, 1};
+ when(_mailbox1.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row1),
+ OperatorTestUtil.block(DATA_SCHEMA, row2), TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row3 = new Object[]{4, 2};
+ Object[] row4 = new Object[]{2, 4};
+ Object[] row5 = new Object[]{-1, 95};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(DATA_SCHEMA, row3),
+ OperatorTestUtil.block(DATA_SCHEMA, row4), OperatorTestUtil.block(DATA_SCHEMA, row5),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
- // Setup the collation key and direction
- List<RexExpression> collationKeys = new ArrayList<>(Collections.singletonList(new RexExpression.InputRef(0)));
- RelFieldCollation.Direction direction = _random.nextBoolean() ? RelFieldCollation.Direction.ASCENDING
- : RelFieldCollation.Direction.DESCENDING;
- List<RelFieldCollation.Direction> collationDirection = new ArrayList<>(Collections.singletonList(direction));
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
- ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirection,
- false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID, null);
-
- // Receive a set of no-op blocks and skip over them
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
- }
- List<Object[]> resultRows = receivedBlock.getContainer();
- // All blocks should be returned together since ordering was required
- Assert.assertEquals(resultRows.size(), 5);
- if (direction == RelFieldCollation.Direction.ASCENDING) {
- Assert.assertEquals(resultRows.get(0), expRow5);
- Assert.assertEquals(resultRows.get(1), expRow2);
- Assert.assertEquals(resultRows.get(2), expRow4);
- Assert.assertEquals(resultRows.get(3), expRow1);
- Assert.assertEquals(resultRows.get(4), expRow3);
- } else {
- Assert.assertEquals(resultRows.get(0), expRow3);
- Assert.assertEquals(resultRows.get(1), expRow1);
- Assert.assertEquals(resultRows.get(2), expRow4);
- Assert.assertEquals(resultRows.get(3), expRow2);
- Assert.assertEquals(resultRows.get(4), expRow5);
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, DATA_SCHEMA, COLLATION_KEYS,
+ COLLATION_DIRECTIONS, false, 1, 0)) {
+ assertEquals(receiveOp.nextBlock().getContainer(), Arrays.asList(row5, row2, row4, row1, row3));
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
-
- receivedBlock = receiveOp.nextBlock();
- Assert.assertTrue(receivedBlock.isEndOfStreamBlock());
}
@Test
public void shouldReceiveMailboxFromTwoServersWithCollationKeyTwoColumns()
throws Exception {
- String server1Host = "hash1";
- int server1Port = 123;
- Mockito.when(_server1.getHostname()).thenReturn(server1Host);
- Mockito.when(_server1.getQueryMailboxPort()).thenReturn(server1Port);
-
- String server2Host = "hash2";
- int server2Port = 456;
- Mockito.when(_server2.getHostname()).thenReturn(server2Host);
- Mockito.when(_server2.getQueryMailboxPort()).thenReturn(server2Port);
-
- int jobId = 456;
- int stageId = 0;
- int toPort = 8888;
- String toHost = "toHost";
- VirtualServerAddress toAddress = new VirtualServerAddress(toHost, toPort, 0);
-
- DataSchema inSchema = new DataSchema(new String[]{"col1", "col2", "col3"},
- new DataSchema.ColumnDataType[]{INT, INT, STRING});
- 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[]{3, 3, "queen"};
- Object[] expRow2 = new Object[]{1, 1, "pink floyd"};
- Mockito.when(_mailbox.receive())
- .thenReturn(OperatorTestUtil.block(inSchema, expRow1), OperatorTestUtil.block(inSchema, expRow2),
- TransferableBlockUtils.getEndOfStreamTransferableBlock());
-
- Object[] expRow3 = new Object[]{42, 2, "pink floyd"};
- Object[] expRow4 = new Object[]{2, 4, "aerosmith"};
- Object[] expRow5 = new Object[]{-1, 95, "foo fighters"};
- 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),
- OperatorTestUtil.block(inSchema, expRow4), OperatorTestUtil.block(inSchema, expRow5),
+ DataSchema dataSchema =
+ new DataSchema(new String[]{"col1", "col2", "col3"}, new DataSchema.ColumnDataType[]{INT, INT, STRING});
+ List<RexExpression> collationKeys = Arrays.asList(new RexExpression.InputRef(2), new RexExpression.InputRef(0));
+ List<RelFieldCollation.Direction> collationDirection =
+ Arrays.asList(RelFieldCollation.Direction.DESCENDING, RelFieldCollation.Direction.ASCENDING);
+
+ long requestId = 0;
+ int senderStageId = 1;
+ JsonMailboxIdentifier mailboxId1 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 0), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId1)).thenReturn(_mailbox1);
+ Object[] row1 = new Object[]{3, 3, "queen"};
+ Object[] row2 = new Object[]{1, 1, "pink floyd"};
+ when(_mailbox1.receive()).thenReturn(OperatorTestUtil.block(dataSchema, row1),
+ OperatorTestUtil.block(dataSchema, row2), TransferableBlockUtils.getEndOfStreamTransferableBlock());
+ JsonMailboxIdentifier mailboxId2 = new JsonMailboxIdentifier(String.format("%s_%s", requestId, senderStageId),
+ new VirtualServerAddress("localhost", 123, 1), RECEIVER_ADDRESS, senderStageId, 0);
+ when(_mailboxService.getReceivingMailbox(mailboxId2)).thenReturn(_mailbox2);
+ Object[] row3 = new Object[]{4, 2, "pink floyd"};
+ Object[] row4 = new Object[]{2, 4, "aerosmith"};
+ Object[] row5 = new Object[]{-1, 95, "foo fighters"};
+ when(_mailbox2.receive()).thenReturn(OperatorTestUtil.block(dataSchema, row3),
+ OperatorTestUtil.block(dataSchema, row4), OperatorTestUtil.block(dataSchema, row5),
TransferableBlockUtils.getEndOfStreamTransferableBlock());
- // Setup the collation key and direction
- List<RexExpression> collationKeys = new ArrayList<>(Arrays.asList(new RexExpression.InputRef(2),
- new RexExpression.InputRef(0)));
- RelFieldCollation.Direction direction1 = _random.nextBoolean() ? RelFieldCollation.Direction.ASCENDING
- : RelFieldCollation.Direction.DESCENDING;
- RelFieldCollation.Direction direction2 = RelFieldCollation.Direction.ASCENDING;
- List<RelFieldCollation.Direction> collationDirection = new ArrayList<>(Arrays.asList(direction1, direction2));
-
OpChainExecutionContext context =
- new OpChainExecutionContext(_mailboxService, jobId, DEFAULT_RECEIVER_STAGE_ID, toAddress, Long.MAX_VALUE,
- Long.MAX_VALUE, new HashMap<>(), false);
-
- SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
- ImmutableList.of(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, collationKeys, collationDirection,
- false, true, inSchema, stageId, DEFAULT_RECEIVER_STAGE_ID, null);
-
- // Receive a set of no-op blocks and skip over them
- TransferableBlock receivedBlock = receiveOp.nextBlock();
- while (receivedBlock.isNoOpBlock()) {
- receivedBlock = receiveOp.nextBlock();
+ new OpChainExecutionContext(_mailboxService, 0, 0, RECEIVER_ADDRESS, Long.MAX_VALUE, Long.MAX_VALUE,
+ Collections.emptyMap(), false);
+ try (SortedMailboxReceiveOperator receiveOp = new SortedMailboxReceiveOperator(context,
+ Arrays.asList(_server1, _server2), RelDistribution.Type.HASH_DISTRIBUTED, dataSchema, collationKeys,
+ collationDirection, false, 1, 0)) {
+ assertEquals(receiveOp.nextBlock().getContainer(), Arrays.asList(row1, row2, row3, row5, row4));
+ assertTrue(receiveOp.nextBlock().isEndOfStreamBlock());
}
- List<Object[]> resultRows = receivedBlock.getContainer();
- // All blocks should be returned together since ordering was required
- Assert.assertEquals(resultRows.size(), 5);
- if (direction1 == RelFieldCollation.Direction.ASCENDING) {
- Assert.assertEquals(resultRows.get(0), expRow4);
- Assert.assertEquals(resultRows.get(1), expRow5);
- Assert.assertEquals(resultRows.get(2), expRow2);
- Assert.assertEquals(resultRows.get(3), expRow3);
- Assert.assertEquals(resultRows.get(4), expRow1);
- } else {
- Assert.assertEquals(resultRows.get(0), expRow1);
- Assert.assertEquals(resultRows.get(1), expRow2);
- Assert.assertEquals(resultRows.get(2), expRow3);
- Assert.assertEquals(resultRows.get(3), expRow5);
- Assert.assertEquals(resultRows.get(4), expRow4);
- }
-
- receivedBlock = receiveOp.nextBlock();
- Assert.assertTrue(receivedBlock.isEndOfStreamBlock());
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org