You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@samza.apache.org by xi...@apache.org on 2017/12/04 23:46:14 UTC
[2/3] samza git commit: Updated Serde related documentation and error
messages for High Level API
Updated Serde related documentation and error messages for High Level API
Updated and clarified the documentation and error messages related to Serdes for Input/Output/PartitionBy streams.
Author: Prateek Maheshwari <pm...@linkedin.com>
Reviewers: Jagadish Venkatraman <vj...@gmail.com>
Closes #376 from prateekm/documentation-cleanup
Project: http://git-wip-us.apache.org/repos/asf/samza/repo
Commit: http://git-wip-us.apache.org/repos/asf/samza/commit/32d2af5d
Tree: http://git-wip-us.apache.org/repos/asf/samza/tree/32d2af5d
Diff: http://git-wip-us.apache.org/repos/asf/samza/diff/32d2af5d
Branch: refs/heads/0.14.0
Commit: 32d2af5dd10b2c2c1cf3f6e683967bafd5285cdf
Parents: 9c9bbc4
Author: Prateek Maheshwari <pm...@linkedin.com>
Authored: Mon Dec 4 14:21:21 2017 -0800
Committer: Prateek Maheshwari <pm...@linkedin.com>
Committed: Mon Dec 4 14:21:21 2017 -0800
----------------------------------------------------------------------
.../apache/samza/operators/MessageStream.java | 18 ++++++++---
.../org/apache/samza/operators/StreamGraph.java | 33 ++++++++++++++++----
.../org/apache/samza/serializers/KVSerde.java | 8 +++--
.../samza/operators/MessageStreamImpl.java | 4 +++
.../apache/samza/operators/StreamGraphImpl.java | 6 ++--
.../samza/operators/impl/InputOperatorImpl.java | 2 +-
.../operators/impl/OutputOperatorImpl.java | 2 +-
.../operators/impl/PartitionByOperatorImpl.java | 4 ---
.../samza/operators/spec/InputOperatorSpec.java | 10 +++---
.../samza/operators/spec/OutputStreamImpl.java | 10 +++---
.../stream/IntermediateMessageStreamImpl.java | 13 ++++++++
.../samza/operators/TestMessageStreamImpl.java | 2 ++
12 files changed, 80 insertions(+), 32 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-api/src/main/java/org/apache/samza/operators/MessageStream.java
----------------------------------------------------------------------
diff --git a/samza-api/src/main/java/org/apache/samza/operators/MessageStream.java b/samza-api/src/main/java/org/apache/samza/operators/MessageStream.java
index dcce7c8..acb2f33 100644
--- a/samza-api/src/main/java/org/apache/samza/operators/MessageStream.java
+++ b/samza-api/src/main/java/org/apache/samza/operators/MessageStream.java
@@ -93,6 +93,9 @@ public interface MessageStream<M> {
/**
* Allows sending messages in this {@link MessageStream} to an {@link OutputStream}.
+ * <p>
+ * When sending messages to an {@code OutputStream<KV<K, V>>}, messages are partitioned using their serialized key.
+ * When sending messages to any other {@code OutputStream<M>}, messages are partitioned using a null partition key.
*
* @param outputStream the output stream to send messages to
*/
@@ -185,8 +188,9 @@ public interface MessageStream<M> {
* input to the job.
* <p>
* Uses the provided {@link KVSerde} for serialization of keys and values. If the provided {@code serde} is null,
- * uses the default serde provided via {@link StreamGraph#setDefaultSerde}, which must be a KVSerde.
- * If no default serde has been provided <b>before</b> calling this method, no-op serdes are used for keys and values.
+ * uses the default serde provided via {@link StreamGraph#setDefaultSerde}, which must be a KVSerde. If the default
+ * serde is not a {@link KVSerde}, a runtime exception will be thrown. If no default serde has been provided
+ * <b>before</b> calling this method, a {@code KVSerde<NoOpSerde, NoOpSerde>} is used.
* <p>
* The number of partitions for this intermediate stream is determined as follows:
* If the stream is an eventual input to a {@link #join}, and the number of partitions for the other stream is known,
@@ -200,8 +204,11 @@ public interface MessageStream<M> {
* for any state stores and streams created by this operator (the full ID also contains the job name, job id and
* operator type). If the application logic is changed, this ID must be reused in the new operator to retain
* state from the previous version, and changed for the new operator to discard the state from the previous version.
+ * <p>
+ * Unlike {@link #sendTo}, messages with a null key are all sent to partition 0.
*
- * @param keyExtractor the {@link Function} to extract the message and partition key from the input message
+ * @param keyExtractor the {@link Function} to extract the message and partition key from the input message.
+ * Messages with a null key are all sent to partition 0.
* @param valueExtractor the {@link Function} to extract the value from the input message
* @param serde the {@link KVSerde} to use for (de)serializing the key and value.
* @param id the unique id of this operator in this application
@@ -212,9 +219,12 @@ public interface MessageStream<M> {
<K, V> MessageStream<KV<K, V>> partitionBy(Function<? super M, ? extends K> keyExtractor,
Function<? super M, ? extends V> valueExtractor, KVSerde<K, V> serde, String id);
-
/**
* Same as calling {@link #partitionBy(Function, Function, KVSerde, String)} with a null KVSerde.
+ * <p>
+ * Uses the default serde provided via {@link StreamGraph#setDefaultSerde}, which must be a KVSerde. If the default
+ * serde is not a {@link KVSerde}, a runtime exception will be thrown. If no default serde has been provided
+ * <b>before</b> calling this method, a {@code KVSerde<NoOpSerde, NoOpSerde>} is used.
*
* @param keyExtractor the {@link Function} to extract the message and partition key from the input message
* @param valueExtractor the {@link Function} to extract the value from the input message
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-api/src/main/java/org/apache/samza/operators/StreamGraph.java
----------------------------------------------------------------------
diff --git a/samza-api/src/main/java/org/apache/samza/operators/StreamGraph.java b/samza-api/src/main/java/org/apache/samza/operators/StreamGraph.java
index 17223b1..4930631 100644
--- a/samza-api/src/main/java/org/apache/samza/operators/StreamGraph.java
+++ b/samza-api/src/main/java/org/apache/samza/operators/StreamGraph.java
@@ -32,8 +32,10 @@ public interface StreamGraph {
* Sets the default {@link Serde} to use for (de)serializing messages.
* <p>.
* If the default serde is set, it must be set <b>before</b> creating any input or output streams.
- * If no explicit or default serdes are provided, a NoOpSerde is used for keys and values. This means that any
- * streams created without explicit or default serdes should be cast to MessageStream<KV<Object, Object>>.
+ * <p>
+ * If no explicit or default serdes are provided, a {@code KVSerde<NoOpSerde, NoOpSerde>} is used. This means that
+ * any streams created without explicit or default serdes should be cast to {@code MessageStream<KV<Object, Object>>}.
+ * <p>
* Providing an incompatible message type for the input/output streams that use the default serde will result in
* {@link ClassCastException}s at runtime.
*
@@ -44,6 +46,13 @@ public interface StreamGraph {
/**
* Gets the input {@link MessageStream} corresponding to the {@code streamId}.
* <p>
+ * An input {@code MessageStream<KV<K, V>}, which can be obtained by calling this method with a {@code KVSerde<K, V>},
+ * can receive messages of type {@code KV<K, V>}. An input {@code MessageStream<M>} with any other {@code Serde<M>}
+ * can receive messages of type M - the key in the incoming message is ignored.
+ * <p>
+ * A {@code KVSerde<NoOpSerde, NoOpSerde>} or {@code NoOpSerde} may be used if the {@code SystemConsumer}
+ * deserializes the incoming messages itself, and no further deserialization is required from the framework.
+ * <p>
* Multiple invocations of this method with the same {@code streamId} will throw an {@link IllegalStateException}.
*
* @param streamId the unique ID for the stream
@@ -58,9 +67,10 @@ public interface StreamGraph {
* Same as {@link #getInputStream(String, Serde)}, but uses the default {@link Serde} provided via
* {@link #setDefaultSerde(Serde)} for deserializing input messages.
* <p>
- * If no default serde has been provided <b>before</b> calling this method, a no-op serde is used.
- * Providing a message type {@code M} that is incompatible with the default Serde will result in
+ * If no default serde has been provided <b>before</b> calling this method, a {@code KVSerde<NoOpSerde, NoOpSerde>}
+ * is used. Providing a message type {@code M} that is incompatible with the default Serde will result in
* {@link ClassCastException}s at runtime.
+ * <p>
* Multiple invocations of this method with the same {@code streamId} will throw an {@link IllegalStateException}.
*
* @param streamId the unique ID for the stream
@@ -73,6 +83,16 @@ public interface StreamGraph {
/**
* Gets the {@link OutputStream} corresponding to the {@code streamId}.
* <p>
+ * An {@code OutputStream<KV<K, V>>}, which can be obtained by calling this method with a {@code KVSerde<K, V>},
+ * can send messages of type {@code KV<K, V>}. An {@code OutputStream<M>} with any other {@code Serde<M>} can
+ * send messages of type M without a key.
+ * <p>
+ * A {@code KVSerde<NoOpSerde, NoOpSerde>} or {@code NoOpSerde} may be used if the {@code SystemProducer}
+ * serializes the outgoing messages itself, and no prior serialization is required from the framework.
+ * <p>
+ * When sending messages to an {@code OutputStream<KV<K, V>>}, messages are partitioned using their serialized key.
+ * When sending messages to any other {@code OutputStream<M>}, messages are partitioned using a null partition key.
+ * <p>
* Multiple invocations of this method with the same {@code streamId} will throw an {@link IllegalStateException}.
*
* @param streamId the unique ID for the stream
@@ -87,9 +107,10 @@ public interface StreamGraph {
* Same as {@link #getOutputStream(String, Serde)}, but uses the default {@link Serde} provided via
* {@link #setDefaultSerde(Serde)} for serializing output messages.
* <p>
- * If no default serde has been provided <b>before</b> calling this method, a no-op serde is used.
- * Providing a message type {@code M} that is incompatible with the default Serde will result in
+ * If no default serde has been provided <b>before</b> calling this method, a {@code KVSerde<NoOpSerde, NoOpSerde>}
+ * is used. Providing a message type {@code M} that is incompatible with the default Serde will result in
* {@link ClassCastException}s at runtime.
+ * <p>
* Multiple invocations of this method with the same {@code streamId} will throw an {@link IllegalStateException}.
*
* @param streamId the unique ID for the stream
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-api/src/main/java/org/apache/samza/serializers/KVSerde.java
----------------------------------------------------------------------
diff --git a/samza-api/src/main/java/org/apache/samza/serializers/KVSerde.java b/samza-api/src/main/java/org/apache/samza/serializers/KVSerde.java
index 10f224b..42a527d 100644
--- a/samza-api/src/main/java/org/apache/samza/serializers/KVSerde.java
+++ b/samza-api/src/main/java/org/apache/samza/serializers/KVSerde.java
@@ -25,8 +25,12 @@ import java.nio.ByteBuffer;
/**
- * A marker serde class to indicate that messages are keyed and should be deserialized as K-V pairs. This class is
- * intended for use cases where a single Serde parameter or configuration is required.
+ * A serializer that serves one of the following two purposes:
+ * <ol>
+ * <li>When used for input/output streams, serves as an indication to the framework that messages are keyed and
+ * should be deserialized as K-V pairs.</li>
+ * <li>When used for a key or value of type KV, serializes the key and value pair using the provided serdes.</li>
+ * </ol>
*
* @param <K> type of the key in the message
* @param <V> type of the value in the message
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java b/samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java
index e6e711c..3f4e40d 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java
@@ -152,6 +152,10 @@ public class MessageStreamImpl<M> implements MessageStream<M> {
Function<? super M, ? extends V> valueExtractor, KVSerde<K, V> serde, String userDefinedId) {
String opId = this.graph.getNextOpId(OpCode.PARTITION_BY, userDefinedId);
IntermediateMessageStreamImpl<KV<K, V>> intermediateStream = this.graph.getIntermediateStream(opId, serde);
+ if (!intermediateStream.isKeyed()) {
+ // this can only happen when the default serde partitionBy variant is being used
+ throw new SamzaException("partitionBy can not be used with a default serde that is not a KVSerde.");
+ }
PartitionByOperatorSpec<M, K, V> partitionByOperatorSpec =
OperatorSpecs.createPartitionByOperatorSpec(
intermediateStream.getOutputStream(), keyExtractor, valueExtractor, opId);
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/StreamGraphImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/StreamGraphImpl.java b/samza-core/src/main/java/org/apache/samza/operators/StreamGraphImpl.java
index 5323c32..d014cb9 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/StreamGraphImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/StreamGraphImpl.java
@@ -108,8 +108,7 @@ public class StreamGraphImpl implements StreamGraph {
InputOperatorSpec inputOperatorSpec =
OperatorSpecs.createInputOperatorSpec(streamSpec, kvSerdes.getKey(), kvSerdes.getValue(),
isKeyed, this.getNextOpId(OpCode.INPUT, null));
- inputOperators.put(streamSpec,
- inputOperatorSpec);
+ inputOperators.put(streamSpec, inputOperatorSpec);
return new MessageStreamImpl<>(this, inputOperators.get(streamSpec));
}
@@ -137,8 +136,7 @@ public class StreamGraphImpl implements StreamGraph {
}
boolean isKeyed = serde instanceof KVSerde;
- outputStreams.put(streamSpec,
- new OutputStreamImpl<>(streamSpec, kvSerdes.getKey(), kvSerdes.getValue(), isKeyed));
+ outputStreams.put(streamSpec, new OutputStreamImpl<>(streamSpec, kvSerdes.getKey(), kvSerdes.getValue(), isKeyed));
return outputStreams.get(streamSpec);
}
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java
index 9cc5370..ef3c322 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java
@@ -50,7 +50,7 @@ public final class InputOperatorImpl<K, V> extends OperatorImpl<KV<K, V>, Object
@Override
public Collection<Object> handleMessage(KV<K, V> pair, MessageCollector collector, TaskCoordinator coordinator) {
- Object message = this.inputOpSpec.isKeyedInput() ? pair : pair.getValue();
+ Object message = this.inputOpSpec.isKeyed() ? pair : pair.getValue();
return Collections.singletonList(message);
}
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/impl/OutputOperatorImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/OutputOperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/OutputOperatorImpl.java
index 7b7e49c..27bef87 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/OutputOperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/OutputOperatorImpl.java
@@ -57,7 +57,7 @@ class OutputOperatorImpl<M> extends OperatorImpl<M, Void> {
public Collection<Void> handleMessage(M message, MessageCollector collector,
TaskCoordinator coordinator) {
Object key, value;
- if (outputStream.isKeyedOutput()) {
+ if (outputStream.isKeyed()) {
key = ((KV) message).getKey();
value = ((KV) message).getValue();
} else {
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java b/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
index b3fb4b2..9fc1e7c 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/impl/PartitionByOperatorImpl.java
@@ -18,7 +18,6 @@
*/
package org.apache.samza.operators.impl;
-import org.apache.samza.SamzaException;
import org.apache.samza.config.Config;
import org.apache.samza.container.TaskContextImpl;
import org.apache.samza.operators.KV;
@@ -55,9 +54,6 @@ class PartitionByOperatorImpl<M, K, V> extends OperatorImpl<M, Void> {
PartitionByOperatorImpl(PartitionByOperatorSpec<M, K, V> partitionByOpSpec, Config config, TaskContext context) {
this.partitionByOpSpec = partitionByOpSpec;
OutputStreamImpl<KV<K, V>> outputStream = partitionByOpSpec.getOutputStream();
- if (!outputStream.isKeyedOutput()) {
- throw new SamzaException("Output stream for repartitioning must be a keyed stream.");
- }
this.systemStream = new SystemStream(
outputStream.getStreamSpec().getSystemName(),
outputStream.getStreamSpec().getPhysicalName());
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/spec/InputOperatorSpec.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/spec/InputOperatorSpec.java b/samza-core/src/main/java/org/apache/samza/operators/spec/InputOperatorSpec.java
index 4c0687a..2ad0597 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/spec/InputOperatorSpec.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/spec/InputOperatorSpec.java
@@ -35,15 +35,15 @@ public class InputOperatorSpec<K, V> extends OperatorSpec<KV<K, V>, Object> { //
private final StreamSpec streamSpec;
private final Serde<K> keySerde;
private final Serde<V> valueSerde;
- private final boolean isKeyedInput;
+ private final boolean isKeyed;
public InputOperatorSpec(StreamSpec streamSpec,
- Serde<K> keySerde, Serde<V> valueSerde, boolean isKeyedInput, String opId) {
+ Serde<K> keySerde, Serde<V> valueSerde, boolean isKeyed, String opId) {
super(OpCode.INPUT, opId);
this.streamSpec = streamSpec;
this.keySerde = keySerde;
this.valueSerde = valueSerde;
- this.isKeyedInput = isKeyedInput;
+ this.isKeyed = isKeyed;
}
public StreamSpec getStreamSpec() {
@@ -58,8 +58,8 @@ public class InputOperatorSpec<K, V> extends OperatorSpec<KV<K, V>, Object> { //
return valueSerde;
}
- public boolean isKeyedInput() {
- return isKeyedInput;
+ public boolean isKeyed() {
+ return isKeyed;
}
@Override
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/spec/OutputStreamImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/spec/OutputStreamImpl.java b/samza-core/src/main/java/org/apache/samza/operators/spec/OutputStreamImpl.java
index a793e0c..e439c4e 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/spec/OutputStreamImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/spec/OutputStreamImpl.java
@@ -28,14 +28,14 @@ public class OutputStreamImpl<M> implements OutputStream<M> {
private final StreamSpec streamSpec;
private final Serde keySerde;
private final Serde valueSerde;
- private final boolean isKeyedOutput;
+ private final boolean isKeyed;
public OutputStreamImpl(StreamSpec streamSpec,
- Serde keySerde, Serde valueSerde, boolean isKeyedOutput) {
+ Serde keySerde, Serde valueSerde, boolean isKeyed) {
this.streamSpec = streamSpec;
this.keySerde = keySerde;
this.valueSerde = valueSerde;
- this.isKeyedOutput = isKeyedOutput;
+ this.isKeyed = isKeyed;
}
public StreamSpec getStreamSpec() {
@@ -50,7 +50,7 @@ public class OutputStreamImpl<M> implements OutputStream<M> {
return valueSerde;
}
- public boolean isKeyedOutput() {
- return isKeyedOutput;
+ public boolean isKeyed() {
+ return isKeyed;
}
}
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/main/java/org/apache/samza/operators/stream/IntermediateMessageStreamImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/main/java/org/apache/samza/operators/stream/IntermediateMessageStreamImpl.java b/samza-core/src/main/java/org/apache/samza/operators/stream/IntermediateMessageStreamImpl.java
index 279cdd4..5eeca99 100644
--- a/samza-core/src/main/java/org/apache/samza/operators/stream/IntermediateMessageStreamImpl.java
+++ b/samza-core/src/main/java/org/apache/samza/operators/stream/IntermediateMessageStreamImpl.java
@@ -25,6 +25,8 @@ import org.apache.samza.operators.spec.InputOperatorSpec;
import org.apache.samza.operators.spec.OperatorSpec;
import org.apache.samza.operators.spec.OutputStreamImpl;
import org.apache.samza.system.StreamSpec;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* An intermediate stream is both an input and an output stream (e.g. a repartitioned stream).
@@ -39,12 +41,19 @@ import org.apache.samza.system.StreamSpec;
*/
public class IntermediateMessageStreamImpl<M> extends MessageStreamImpl<M> implements OutputStream<M> {
+ private static final Logger LOGGER = LoggerFactory.getLogger(IntermediateMessageStreamImpl.class);
private final OutputStreamImpl<M> outputStream;
+ private final boolean isKeyed;
public IntermediateMessageStreamImpl(StreamGraphImpl graph, InputOperatorSpec<?, M> inputOperatorSpec,
OutputStreamImpl<M> outputStream) {
super(graph, (OperatorSpec<?, M>) inputOperatorSpec);
this.outputStream = outputStream;
+ if (inputOperatorSpec.isKeyed() != outputStream.isKeyed()) {
+ LOGGER.error("Input and output streams for intermediate stream {} aren't keyed consistently. Input: {}, Output: {}",
+ new Object[]{inputOperatorSpec.getStreamSpec().getId(), inputOperatorSpec.isKeyed(), outputStream.isKeyed()});
+ }
+ this.isKeyed = inputOperatorSpec.isKeyed() && outputStream.isKeyed();
}
public StreamSpec getStreamSpec() {
@@ -54,4 +63,8 @@ public class IntermediateMessageStreamImpl<M> extends MessageStreamImpl<M> imple
public OutputStreamImpl<M> getOutputStream() {
return this.outputStream;
}
+
+ public boolean isKeyed() {
+ return isKeyed;
+ }
}
http://git-wip-us.apache.org/repos/asf/samza/blob/32d2af5d/samza-core/src/test/java/org/apache/samza/operators/TestMessageStreamImpl.java
----------------------------------------------------------------------
diff --git a/samza-core/src/test/java/org/apache/samza/operators/TestMessageStreamImpl.java b/samza-core/src/test/java/org/apache/samza/operators/TestMessageStreamImpl.java
index 5b3c3a0..d97d494 100644
--- a/samza-core/src/test/java/org/apache/samza/operators/TestMessageStreamImpl.java
+++ b/samza-core/src/test/java/org/apache/samza/operators/TestMessageStreamImpl.java
@@ -207,6 +207,7 @@ public class TestMessageStreamImpl {
.thenReturn(mockIntermediateStream);
when(mockIntermediateStream.getOutputStream())
.thenReturn(mockOutputStreamImpl);
+ when(mockIntermediateStream.isKeyed()).thenReturn(true);
MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec);
Function mockKeyFunction = mock(Function.class);
@@ -236,6 +237,7 @@ public class TestMessageStreamImpl {
.thenReturn(mockIntermediateStream);
when(mockIntermediateStream.getOutputStream())
.thenReturn(mockOutputStreamImpl);
+ when(mockIntermediateStream.isKeyed()).thenReturn(true);
MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec);
Function mockKeyFunction = mock(Function.class);