Examples with StreamGraphImpl org.apache.samza.operators.StreamGraphImpl used on opensource projects

Example 6 with StreamGraphImpl

use of org.apache.samza.operators.StreamGraphImpl in project samza by apache.

the class TestOperatorImpls method testJoinChain.

@Test
public void testJoinChain() throws IllegalAccessException, InvocationTargetException {
    // test creation of join chain
    StreamGraphImpl mockGraph = mock(StreamGraphImpl.class);
    MessageStreamImpl<TestMessageEnvelope> input1 = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    MessageStreamImpl<TestMessageEnvelope> input2 = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    TaskContext mockContext = mock(TaskContext.class);
    when(mockContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config mockConfig = mock(Config.class);
    input1.join(input2, new JoinFunction<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope>() {

        @Override
        public TestOutputMessageEnvelope apply(TestMessageEnvelope m1, TestMessageEnvelope m2) {
            return new TestOutputMessageEnvelope(m1.getKey(), m1.getMessage().getValue().length() + m2.getMessage().getValue().length());
        }

        @Override
        public String getFirstKey(TestMessageEnvelope message) {
            return message.getKey();
        }

        @Override
        public String getSecondKey(TestMessageEnvelope message) {
            return message.getKey();
        }
    }, Duration.ofMinutes(1)).map(m -> m);
    OperatorImplGraph opGraph = new OperatorImplGraph();
    // now, we create chained operators from each input sources
    RootOperatorImpl chain1 = (RootOperatorImpl) createOpsMethod.invoke(opGraph, input1, mockConfig, mockContext);
    RootOperatorImpl chain2 = (RootOperatorImpl) createOpsMethod.invoke(opGraph, input2, mockConfig, mockContext);
    // check that those two chains will merge at map operator
    // first branch of the join
    Set<OperatorImpl> subsSet = (Set<OperatorImpl>) nextOperatorsField.get(chain1);
    assertEquals(subsSet.size(), 1);
    OperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> joinOp1 = subsSet.iterator().next();
    Set<OperatorImpl> subsOps = (Set<OperatorImpl>) nextOperatorsField.get(joinOp1);
    assertEquals(subsOps.size(), 1);
    // the map operator consumes the common join output, where two branches merge
    OperatorImpl mapImpl = subsOps.iterator().next();
    // second branch of the join
    subsSet = (Set<OperatorImpl>) nextOperatorsField.get(chain2);
    assertEquals(subsSet.size(), 1);
    OperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> joinOp2 = subsSet.iterator().next();
    assertNotSame(joinOp1, joinOp2);
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(joinOp2);
    assertEquals(subsOps.size(), 1);
    // make sure that the map operator is the same
    assertEquals(mapImpl, subsOps.iterator().next());
}

Example 7 with StreamGraphImpl

use of org.apache.samza.operators.StreamGraphImpl in project samza by apache.

the class TestOperatorSpecs method testCreateMergeOperator.

@Test
public void testCreateMergeOperator() {
    StreamGraphImpl mockGraph = mock(StreamGraphImpl.class);
    MessageStreamImpl<TestMessageEnvelope> output = TestMessageStreamImplUtil.<TestMessageEnvelope>getMessageStreamImpl(mockGraph);
    StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope> mergeOp = OperatorSpecs.createMergeOperatorSpec(output, 1);
    Function<TestMessageEnvelope, Collection<TestMessageEnvelope>> mergeFn = t -> new ArrayList<TestMessageEnvelope>() {

        {
            this.add(t);
        }
    };
    TestMessageEnvelope t = mock(TestMessageEnvelope.class);
    assertEquals(mergeOp.getTransformFn().apply(t), mergeFn.apply(t));
    assertEquals(mergeOp.getNextStream(), output);
}

Example 8 with StreamGraphImpl

use of org.apache.samza.operators.StreamGraphImpl in project samza by apache.

the class TestOperatorImpls method testBroadcastChain.

@Test
public void testBroadcastChain() throws IllegalAccessException, InvocationTargetException {
    // test creation of broadcast chain
    StreamGraphImpl mockGraph = mock(StreamGraphImpl.class);
    MessageStreamImpl<TestMessageEnvelope> testInput = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    TaskContext mockContext = mock(TaskContext.class);
    when(mockContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config mockConfig = mock(Config.class);
    testInput.filter(m -> m.getMessage().getEventTime() > 123456L).flatMap(m -> new ArrayList() {

        {
            this.add(m);
            this.add(m);
        }
    });
    testInput.filter(m -> m.getMessage().getEventTime() < 123456L).map(m -> m);
    OperatorImplGraph opGraph = new OperatorImplGraph();
    RootOperatorImpl operatorChain = (RootOperatorImpl) createOpsMethod.invoke(opGraph, testInput, mockConfig, mockContext);
    Set<OperatorImpl> subsSet = (Set<OperatorImpl>) nextOperatorsField.get(operatorChain);
    assertEquals(subsSet.size(), 2);
    Iterator<OperatorImpl> iter = subsSet.iterator();
    // check the first branch w/ flatMap
    OperatorImpl<TestMessageEnvelope, TestMessageEnvelope> opImpl = iter.next();
    Set<OperatorImpl> subsOps = (Set<OperatorImpl>) nextOperatorsField.get(opImpl);
    assertEquals(subsOps.size(), 1);
    OperatorImpl flatMapImpl = subsOps.iterator().next();
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(flatMapImpl);
    assertEquals(subsOps.size(), 0);
    // check the second branch w/ map
    opImpl = iter.next();
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(opImpl);
    assertEquals(subsOps.size(), 1);
    OperatorImpl mapImpl = subsOps.iterator().next();
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(mapImpl);
    assertEquals(subsOps.size(), 0);
}

Example 9 with StreamGraphImpl

use of org.apache.samza.operators.StreamGraphImpl in project samza by apache.

the class TestOperatorImpls method testLinearChain.

@Test
public void testLinearChain() throws IllegalAccessException, InvocationTargetException {
    // test creation of linear chain
    StreamGraphImpl mockGraph = mock(StreamGraphImpl.class);
    MessageStreamImpl<TestMessageEnvelope> testInput = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    TaskContext mockContext = mock(TaskContext.class);
    when(mockContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config mockConfig = mock(Config.class);
    testInput.map(m -> m).window(Windows.keyedSessionWindow(TestMessageEnvelope::getKey, Duration.ofMinutes(10)));
    OperatorImplGraph opGraph = new OperatorImplGraph();
    RootOperatorImpl operatorChain = (RootOperatorImpl) createOpsMethod.invoke(opGraph, testInput, mockConfig, mockContext);
    Set<OperatorImpl> subsSet = (Set<OperatorImpl>) nextOperatorsField.get(operatorChain);
    assertEquals(subsSet.size(), 1);
    OperatorImpl<TestMessageEnvelope, TestMessageEnvelope> firstOpImpl = subsSet.iterator().next();
    Set<OperatorImpl> subsOps = (Set<OperatorImpl>) nextOperatorsField.get(firstOpImpl);
    assertEquals(subsOps.size(), 1);
    OperatorImpl wndOpImpl = subsOps.iterator().next();
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(wndOpImpl);
    assertEquals(subsOps.size(), 0);
}

Example 10 with StreamGraphImpl

use of org.apache.samza.operators.StreamGraphImpl in project samza by apache.

the class TestExecutionPlanner method createStreamGraphWithJoin.

private StreamGraphImpl createStreamGraphWithJoin() {
    /**
     * the graph looks like the following. number of partitions in parentheses. quotes indicate expected value.
     *
     *                               input1 (64) -> map -> join -> output1 (8)
     *                                                       |
     *          input2 (16) -> partitionBy ("64") -> filter -|
     *                                                       |
     * input3 (32) -> filter -> partitionBy ("64") -> map -> join -> output2 (16)
     *
     */
    StreamGraphImpl streamGraph = new StreamGraphImpl(runner, config);
    BiFunction msgBuilder = mock(BiFunction.class);
    MessageStream m1 = streamGraph.getInputStream("input1", msgBuilder).map(m -> m);
    MessageStream m2 = streamGraph.getInputStream("input2", msgBuilder).partitionBy(m -> "haha").filter(m -> true);
    MessageStream m3 = streamGraph.getInputStream("input3", msgBuilder).filter(m -> true).partitionBy(m -> "hehe").map(m -> m);
    Function mockFn = mock(Function.class);
    OutputStream<Object, Object, Object> output1 = streamGraph.getOutputStream("output1", mockFn, mockFn);
    OutputStream<Object, Object, Object> output2 = streamGraph.getOutputStream("output2", mockFn, mockFn);
    m1.join(m2, mock(JoinFunction.class), Duration.ofHours(2)).sendTo(output1);
    m3.join(m2, mock(JoinFunction.class), Duration.ofHours(1)).sendTo(output2);
    return streamGraph;
}