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

Example 11 with MessageStreamImpl

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

the class TestOperatorSpecs method testCreateStreamOperator.

@Test
public void testCreateStreamOperator() {
    FlatMapFunction<Object, TestMessageEnvelope> transformFn = m -> new ArrayList<TestMessageEnvelope>() {

        {
            this.add(new TestMessageEnvelope(m.toString(), m.toString(), 12345L));
        }
    };
    MessageStreamImpl<TestMessageEnvelope> mockOutput = mock(MessageStreamImpl.class);
    StreamOperatorSpec<Object, TestMessageEnvelope> streamOp = OperatorSpecs.createStreamOperatorSpec(transformFn, mockOutput, 1);
    assertEquals(streamOp.getTransformFn(), transformFn);
    Object mockInput = mock(Object.class);
    when(mockInput.toString()).thenReturn("test-string-1");
    List<TestMessageEnvelope> outputs = (List<TestMessageEnvelope>) streamOp.getTransformFn().apply(mockInput);
    assertEquals(outputs.size(), 1);
    assertEquals(outputs.get(0).getKey(), "test-string-1");
    assertEquals(outputs.get(0).getMessage().getValue(), "test-string-1");
    assertEquals(outputs.get(0).getMessage().getEventTime(), 12345L);
    assertEquals(streamOp.getNextStream(), mockOutput);
}

Example 12 with MessageStreamImpl

use of org.apache.samza.operators.MessageStreamImpl 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);
}

Example 13 with MessageStreamImpl

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

the class OperatorImplGraph method init.

/**
   * Initialize the DAG of {@link OperatorImpl}s for the input {@link MessageStreamImpl} in the provided
   * {@link StreamGraphImpl}.
   *
   * @param streamGraph  the logical {@link StreamGraphImpl}
   * @param config  the {@link Config} required to instantiate operators
   * @param context  the {@link TaskContext} required to instantiate operators
   */
public void init(StreamGraphImpl streamGraph, Config config, TaskContext context) {
    streamGraph.getInputStreams().forEach((streamSpec, inputStream) -> {
        SystemStream systemStream = new SystemStream(streamSpec.getSystemName(), streamSpec.getPhysicalName());
        this.rootOperators.put(systemStream, this.createOperatorImpls((MessageStreamImpl) inputStream, config, context));
    });
}

Example 14 with MessageStreamImpl

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

the class TestProjectTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalProject mockProject = PowerMockito.mock(LogicalProject.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl testMetricsRegistryImpl = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    List<RelNode> inputs = new ArrayList<>();
    inputs.add(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockProject.getId()).thenReturn(2);
    when(mockProject.getInputs()).thenReturn(inputs);
    when(mockProject.getInput()).thenReturn(mockInput);
    RelDataType mockRowType = mock(RelDataType.class);
    when(mockRowType.getFieldCount()).thenReturn(1);
    when(mockProject.getRowType()).thenReturn(mockRowType);
    RexNode mockRexField = mock(RexNode.class);
    List<Pair<RexNode, String>> namedProjects = new ArrayList<>();
    namedProjects.add(Pair.of(mockRexField, "test_field"));
    when(mockProject.getNamedProjects()).thenReturn(namedProjects);
    StreamApplicationDescriptorImpl mockAppDesc = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockAppDesc, mockInputOp);
    when(mockTranslatorContext.getMessageStream(eq(1))).thenReturn(mockStream);
    doAnswer(this.getRegisterMessageStreamAnswer()).when(mockTranslatorContext).registerMessageStream(eq(2), any(MessageStream.class));
    RexToJavaCompiler mockCompiler = mock(RexToJavaCompiler.class);
    when(mockTranslatorContext.getExpressionCompiler()).thenReturn(mockCompiler);
    Expression mockExpr = mock(Expression.class);
    when(mockCompiler.compile(any(), any())).thenReturn(mockExpr);
    when(mockContext.getContainerContext()).thenReturn(mockContainerContext);
    when(mockContainerContext.getContainerMetricsRegistry()).thenReturn(testMetricsRegistryImpl);
    // Apply translate() method to verify that we are getting the correct map operator constructed
    ProjectTranslator projectTranslator = new ProjectTranslator(1);
    projectTranslator.translate(mockProject, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockProject);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockProject);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec projectSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(projectSpec);
    assertEquals(projectSpec.getOpCode(), OperatorSpec.OpCode.MAP);
    // Verify that the bootstrap() method will establish the context for the map function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    projectSpec.getTransformFn().init(mockContext);
    MapFunction mapFn = (MapFunction) Whitebox.getInternalState(projectSpec, "mapFn");
    assertNotNull(mapFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(mapFn, "translatorContext"));
    assertEquals(mockProject, Whitebox.getInternalState(mapFn, "project"));
    assertEquals(mockExpr, Whitebox.getInternalState(mapFn, "expr"));
    // Verify TestMetricsRegistryImpl works with Project
    assertEquals(1, testMetricsRegistryImpl.getGauges().size());
    assertEquals(2, testMetricsRegistryImpl.getGauges().get(LOGICAL_OP_ID).size());
    assertEquals(1, testMetricsRegistryImpl.getCounters().size());
    assertEquals(2, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling mapFn.apply() to verify the filter function is correctly applied to the input message
    SamzaSqlRelMessage mockInputMsg = new SamzaSqlRelMessage(new ArrayList<>(), new ArrayList<>(), new SamzaSqlRelMsgMetadata(0L, 0L));
    SamzaSqlExecutionContext executionContext = mock(SamzaSqlExecutionContext.class);
    DataContext dataContext = mock(DataContext.class);
    when(mockTranslatorContext.getExecutionContext()).thenReturn(executionContext);
    when(mockTranslatorContext.getDataContext()).thenReturn(dataContext);
    Object[] result = new Object[1];
    final Object mockFieldObj = new Object();
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = mockFieldObj;
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    SamzaSqlRelMessage retMsg = (SamzaSqlRelMessage) mapFn.apply(mockInputMsg);
    assertEquals(retMsg.getSamzaSqlRelRecord().getFieldNames(), Collections.singletonList("test_field"));
    assertEquals(retMsg.getSamzaSqlRelRecord().getFieldValues(), Collections.singletonList(mockFieldObj));
    // Verify mapFn.apply() updates the TestMetricsRegistryImpl metrics
    assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}

Example 15 with MessageStreamImpl

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

the class StreamApplicationDescriptorImpl method getInputStream.

@Override
public <M> MessageStream<M> getInputStream(InputDescriptor<M, ?> inputDescriptor) {
    SystemDescriptor systemDescriptor = inputDescriptor.getSystemDescriptor();
    Optional<StreamExpander> expander = systemDescriptor.getExpander();
    if (expander.isPresent()) {
        return expander.get().apply(this, inputDescriptor);
    }
    // TODO: SAMZA-1841: need to add to the broadcast streams if inputDescriptor is for a broadcast stream
    addInputDescriptor(inputDescriptor);
    String streamId = inputDescriptor.getStreamId();
    Serde serde = inputDescriptor.getSerde();
    KV<Serde, Serde> kvSerdes = getOrCreateStreamSerdes(streamId, serde);
    boolean isKeyed = serde instanceof KVSerde;
    InputTransformer transformer = inputDescriptor.getTransformer().orElse(null);
    InputOperatorSpec inputOperatorSpec = OperatorSpecs.createInputOperatorSpec(streamId, kvSerdes.getKey(), kvSerdes.getValue(), transformer, isKeyed, this.getNextOpId(OpCode.INPUT, null));
    inputOperators.put(streamId, inputOperatorSpec);
    return new MessageStreamImpl(this, inputOperators.get(streamId));
}