Examples with OperatorSpec org.apache.samza.operators.spec.OperatorSpec used on opensource projects

Example 16 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec in project samza by apache.

the class TestMessageStreamImpl method testWindowWithRelaxedTypes.

@Test
public void testWindowWithRelaxedTypes() throws Exception {
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec mockOpSpec = mock(OperatorSpec.class);
    MessageStream<TestInputMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec);
    MapFunction<TestMessageEnvelope, String> keyExtractor = m -> m.getKey();
    FoldLeftFunction<TestMessageEnvelope, Integer> aggregator = (m, c) -> c + 1;
    SupplierFunction<Integer> initialValue = () -> 0;
    // should compile since TestMessageEnvelope (input for functions) is base class of TestInputMessageEnvelope (M)
    Window<TestInputMessageEnvelope, String, Integer> window = Windows.keyedTumblingWindow(keyExtractor, Duration.ofHours(1), initialValue, aggregator, null, mock(Serde.class));
    MessageStream<WindowPane<String, Integer>> windowedStream = inputStream.window(window, "w1");
    ArgumentCaptor<OperatorSpec> registeredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec).registerNextOperatorSpec(registeredOpCaptor.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec = registeredOpCaptor.getValue();
    assertTrue(registeredOpSpec instanceof WindowOperatorSpec);
    assertEquals(OpCode.WINDOW, registeredOpSpec.getOpCode());
    assertEquals(window, ((WindowOperatorSpec) registeredOpSpec).getWindow());
}

Example 17 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec in project samza by apache.

the class TestMessageStreamImpl method testJoin.

@Test
public void testJoin() {
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec leftInputOpSpec = mock(OperatorSpec.class);
    MessageStreamImpl<TestMessageEnvelope> source1 = new MessageStreamImpl<>(mockGraph, leftInputOpSpec);
    OperatorSpec rightInputOpSpec = mock(OperatorSpec.class);
    MessageStreamImpl<TestMessageEnvelope> source2 = new MessageStreamImpl<>(mockGraph, rightInputOpSpec);
    JoinFunction<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope> mockJoinFn = mock(JoinFunction.class);
    Duration joinTtl = Duration.ofMinutes(1);
    source1.join(source2, mockJoinFn, mock(Serde.class), mock(Serde.class), mock(Serde.class), joinTtl, "j1");
    ArgumentCaptor<OperatorSpec> leftRegisteredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(leftInputOpSpec).registerNextOperatorSpec(leftRegisteredOpCaptor.capture());
    OperatorSpec<?, TestMessageEnvelope> leftRegisteredOpSpec = leftRegisteredOpCaptor.getValue();
    ArgumentCaptor<OperatorSpec> rightRegisteredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(rightInputOpSpec).registerNextOperatorSpec(rightRegisteredOpCaptor.capture());
    OperatorSpec<?, TestMessageEnvelope> rightRegisteredOpSpec = rightRegisteredOpCaptor.getValue();
    assertEquals(leftRegisteredOpSpec, rightRegisteredOpSpec);
    assertEquals(OpCode.JOIN, leftRegisteredOpSpec.getOpCode());
    assertTrue(leftRegisteredOpSpec instanceof JoinOperatorSpec);
    assertEquals(mockJoinFn, ((JoinOperatorSpec) leftRegisteredOpSpec).getJoinFn());
    assertEquals(joinTtl.toMillis(), ((JoinOperatorSpec) leftRegisteredOpSpec).getTtlMs());
    assertEquals(leftInputOpSpec, ((JoinOperatorSpec) leftRegisteredOpSpec).getLeftInputOpSpec());
    assertEquals(rightInputOpSpec, ((JoinOperatorSpec) leftRegisteredOpSpec).getRightInputOpSpec());
}

Example 18 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec in project samza by apache.

the class TestMessageStreamImpl method validateMergeOperator.

private void validateMergeOperator(MessageStream<TestMessageEnvelope> mergeSource, MessageStream<TestMessageEnvelope> mergeOutput) {
    Collection<OperatorSpec> subs = ((MessageStreamImpl<TestMessageEnvelope>) mergeSource).getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> mergeOp = subs.iterator().next();
    assertTrue(mergeOp instanceof StreamOperatorSpec);
    assertEquals(((StreamOperatorSpec) mergeOp).getNextStream(), mergeOutput);
    TestMessageEnvelope mockMsg = mock(TestMessageEnvelope.class);
    Collection<TestMessageEnvelope> outputs = ((StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope>) mergeOp).getTransformFn().apply(mockMsg);
    assertEquals(outputs.size(), 1);
    assertEquals(outputs.iterator().next(), mockMsg);
}

Example 19 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec in project samza by apache.

the class StreamGraphImpl method getAllOperatorSpecs.

/**
   * Get all {@link OperatorSpec}s available in this {@link StreamGraphImpl}
   *
   * @return  a set of all available {@link OperatorSpec}s
   */
public Collection<OperatorSpec> getAllOperatorSpecs() {
    Collection<InputStreamInternal> inputStreams = inStreams.values();
    Set<OperatorSpec> operatorSpecs = new HashSet<>();
    for (InputStreamInternal stream : inputStreams) {
        doGetOperatorSpecs((MessageStreamImpl) stream, operatorSpecs);
    }
    return operatorSpecs;
}

Example 20 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec in project samza by apache.

the class OperatorImplGraph method createAndRegisterOperatorImpl.

/**
   * Helper method to recursively traverse the {@link OperatorSpec} DAG and instantiate and link the corresponding
   * {@link OperatorImpl}s.
   *
   * @param operatorSpec  the operatorSpec to create the {@link OperatorImpl} for
   * @param config  the {@link Config} required to instantiate operators
   * @param context  the {@link TaskContext} required to instantiate operators
   * @param <M>  type of input message
   * @return  the operator implementation for the operatorSpec
   */
private <M> OperatorImpl<M, ?> createAndRegisterOperatorImpl(OperatorSpec operatorSpec, Config config, TaskContext context) {
    if (!operatorImpls.containsKey(operatorSpec)) {
        OperatorImpl<M, ?> operatorImpl = createOperatorImpl(operatorSpec, config, context);
        if (operatorImpls.putIfAbsent(operatorSpec, operatorImpl) == null) {
            // this is the first time we've added the operatorImpl corresponding to the operatorSpec,
            // so traverse and initialize and register the rest of the DAG.
            // initialize the corresponding operator function
            operatorImpl.init(config, context);
            MessageStreamImpl nextStream = operatorSpec.getNextStream();
            if (nextStream != null) {
                Collection<OperatorSpec> registeredSpecs = nextStream.getRegisteredOperatorSpecs();
                registeredSpecs.forEach(registeredSpec -> {
                    OperatorImpl subImpl = createAndRegisterOperatorImpl(registeredSpec, config, context);
                    operatorImpl.registerNextOperator(subImpl);
                });
            }
            return operatorImpl;
        }
    }
    // and registered, so we do not need to traverse the DAG further.
    return operatorImpls.get(operatorSpec);
}