Examples with IntermediateStreamInternalImpl org.apache.samza.operators.stream.IntermediateStreamInternalImpl used on opensource projects

Example 1 with IntermediateStreamInternalImpl

use of org.apache.samza.operators.stream.IntermediateStreamInternalImpl in project samza by apache.

the class TestStreamGraphImpl method testGetIntermediateStream.

@Test
public void testGetIntermediateStream() {
    ApplicationRunner mockRunner = mock(ApplicationRunner.class);
    Config mockConfig = mock(Config.class);
    StreamSpec testStreamSpec = new StreamSpec("myJob-i001-test-stream-1", "physical-stream-1", "test-system");
    when(mockRunner.getStreamSpec("myJob-i001-test-stream-1")).thenReturn(testStreamSpec);
    when(mockConfig.get(JobConfig.JOB_NAME())).thenReturn("myJob");
    when(mockConfig.get(JobConfig.JOB_ID(), "1")).thenReturn("i001");
    class MyMessageType extends MessageType {

        public final String outputId;

        public MyMessageType(String value, long eventTime, String outputId) {
            super(value, eventTime);
            this.outputId = outputId;
        }
    }
    StreamGraphImpl graph = new StreamGraphImpl(mockRunner, mockConfig);
    Function<TestMessageEnvelope, String> xKeyExtractor = x -> x.getKey();
    Function<TestMessageEnvelope, MyMessageType> xMsgExtractor = x -> new MyMessageType(x.getMessage().getValue(), x.getMessage().getEventTime(), "test-output-id-1");
    BiFunction<String, MessageType, TestInputMessageEnvelope> xMsgBuilder = (k, v) -> new TestInputMessageEnvelope(k, v.getValue(), v.getEventTime(), "input-id-1");
    MessageStream<TestMessageEnvelope> mIntermediateStream = graph.getIntermediateStream("test-stream-1", xKeyExtractor, xMsgExtractor, xMsgBuilder);
    assertEquals(graph.getOutputStreams().get(testStreamSpec), mIntermediateStream);
    assertTrue(mIntermediateStream instanceof IntermediateStreamInternalImpl);
    assertEquals(((IntermediateStreamInternalImpl) mIntermediateStream).getKeyExtractor(), xKeyExtractor);
    assertEquals(((IntermediateStreamInternalImpl) mIntermediateStream).getMsgExtractor(), xMsgExtractor);
    assertEquals(((IntermediateStreamInternalImpl) mIntermediateStream).getMsgBuilder(), xMsgBuilder);
    TestMessageEnvelope xInputMsg = new TestMessageEnvelope("test-key-1", "test-msg-1", 33333L);
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getKeyExtractor().apply(xInputMsg), "test-key-1");
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getMsgExtractor().apply(xInputMsg).getValue(), "test-msg-1");
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getMsgExtractor().apply(xInputMsg).getEventTime(), 33333L);
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getMsgBuilder().apply("test-key-1", new MyMessageType("test-msg-1", 33333L, "test-output-id-1")).getKey(), "test-key-1");
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getMsgBuilder().apply("test-key-1", new MyMessageType("test-msg-1", 33333L, "test-output-id-1")).getMessage().getValue(), "test-msg-1");
    assertEquals(((IntermediateStreamInternalImpl<String, MessageType, TestMessageEnvelope>) mIntermediateStream).getMsgBuilder().apply("test-key-1", new MyMessageType("test-msg-1", 33333L, "test-output-id-1")).getMessage().getEventTime(), 33333L);
}

Example 2 with IntermediateStreamInternalImpl

use of org.apache.samza.operators.stream.IntermediateStreamInternalImpl in project samza by apache.

the class StreamGraphImpl method getIntermediateStream.

/**
   * Internal helper for {@link MessageStreamImpl} to add an intermediate {@link MessageStream} to the graph.
   * An intermediate {@link MessageStream} is both an output and an input stream.
   *
   * @param streamName the name of the stream to be created. Will be prefixed with job name and id to generate the
   *                   logical streamId.
   * @param keyExtractor the {@link Function} to extract the outgoing key from the intermediate message
   * @param msgExtractor the {@link Function} to extract the outgoing message from the intermediate message
   * @param msgBuilder the {@link BiFunction} to convert the incoming key and message to a message
   *                   in the intermediate {@link MessageStream}
   * @param <K> the type of key in the intermediate message
   * @param <V> the type of message in the intermediate message
   * @param <M> the type of messages in the intermediate {@link MessageStream}
   * @return  the intermediate {@link MessageStreamImpl}
   */
<K, V, M> MessageStreamImpl<M> getIntermediateStream(String streamName, Function<? super M, ? extends K> keyExtractor, Function<? super M, ? extends V> msgExtractor, BiFunction<? super K, ? super V, ? extends M> msgBuilder) {
    String streamId = String.format("%s-%s-%s", config.get(JobConfig.JOB_NAME()), config.get(JobConfig.JOB_ID(), "1"), streamName);
    if (msgBuilder == null) {
        throw new IllegalArgumentException("msgBuilder cannot be null for an intermediate stream");
    }
    if (keyExtractor == null) {
        throw new IllegalArgumentException("keyExtractor can't be null for an output stream.");
    }
    if (msgExtractor == null) {
        throw new IllegalArgumentException("msgExtractor can't be null for an output stream.");
    }
    StreamSpec streamSpec = runner.getStreamSpec(streamId);
    IntermediateStreamInternalImpl<K, V, M> intStream = (IntermediateStreamInternalImpl<K, V, M>) inStreams.computeIfAbsent(streamSpec, k -> new IntermediateStreamInternalImpl<>(this, streamSpec, (Function<M, K>) keyExtractor, (Function<M, V>) msgExtractor, (BiFunction<K, V, M>) msgBuilder));
    outStreams.putIfAbsent(streamSpec, intStream);
    return intStream;
}