Search in sources :

Example 1 with FilterFunction

use of org.apache.samza.operators.functions.FilterFunction in project samza by apache.

the class TestOperatorSpec method testStreamOperatorSpecWithFilter.

@Test
public void testStreamOperatorSpecWithFilter() {
    FilterFunction<TestMessageEnvelope> filterFn = m -> m.getKey().equals("key1");
    StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope> streamOperatorSpec = OperatorSpecs.createFilterOperatorSpec(filterFn, "op0");
    StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope> cloneOperatorSpec = (StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope>) OperatorSpecTestUtils.copyOpSpec(streamOperatorSpec);
    assertNotEquals(streamOperatorSpec, cloneOperatorSpec);
    assertTrue(streamOperatorSpec.isClone(cloneOperatorSpec));
    FilterFunction userFn = (FilterFunction) Whitebox.getInternalState(streamOperatorSpec, "filterFn");
    assertEquals(userFn, filterFn);
    assertNotEquals(streamOperatorSpec.getTransformFn(), cloneOperatorSpec.getTransformFn());
    FilterFunction clonedUserFn = (FilterFunction) Whitebox.getInternalState(cloneOperatorSpec, "filterFn");
    assertTrue(cloneOperatorSpec.getTransformFn() instanceof FlatMapFunction);
    assertTrue(clonedUserFn instanceof FilterFunction);
    assertNotEquals(userFn, clonedUserFn);
    assertNull(streamOperatorSpec.getWatermarkFn());
    assertNull(cloneOperatorSpec.getWatermarkFn());
    assertNull(streamOperatorSpec.getScheduledFn());
    assertNull(cloneOperatorSpec.getScheduledFn());
}
Also used : TestMessageEnvelope(org.apache.samza.operators.data.TestMessageEnvelope) FilterFunction(org.apache.samza.operators.functions.FilterFunction) ScheduledFunction(org.apache.samza.operators.functions.ScheduledFunction) StreamTableJoinFunction(org.apache.samza.operators.functions.StreamTableJoinFunction) Collection(java.util.Collection) Scheduler(org.apache.samza.operators.Scheduler) Serde(org.apache.samza.serializers.Serde) Test(org.junit.Test) JoinFunction(org.apache.samza.operators.functions.JoinFunction) FlatMapFunction(org.apache.samza.operators.functions.FlatMapFunction) ArrayList(java.util.ArrayList) MapFunction(org.apache.samza.operators.functions.MapFunction) WatermarkFunction(org.apache.samza.operators.functions.WatermarkFunction) StringSerde(org.apache.samza.serializers.StringSerde) List(java.util.List) Whitebox(org.mockito.internal.util.reflection.Whitebox) TestOutputMessageEnvelope(org.apache.samza.operators.data.TestOutputMessageEnvelope) SinkFunction(org.apache.samza.operators.functions.SinkFunction) JsonSerdeV2(org.apache.samza.serializers.JsonSerdeV2) Assert(org.junit.Assert) KV(org.apache.samza.operators.KV) NoOpSerde(org.apache.samza.serializers.NoOpSerde) FilterFunction(org.apache.samza.operators.functions.FilterFunction) TestMessageEnvelope(org.apache.samza.operators.data.TestMessageEnvelope) FlatMapFunction(org.apache.samza.operators.functions.FlatMapFunction) TestOutputMessageEnvelope(org.apache.samza.operators.data.TestOutputMessageEnvelope) Test(org.junit.Test)

Example 2 with FilterFunction

use of org.apache.samza.operators.functions.FilterFunction in project samza by apache.

the class TestFilterTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalFilter mockFilter = PowerMockito.mock(LogicalFilter.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl metricsRegistry = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    when(mockFilter.getInput()).thenReturn(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockFilter.getId()).thenReturn(2);
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockGraph, 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(metricsRegistry);
    // Apply translate() method to verify that we are getting the correct filter operator constructed
    FilterTranslator filterTranslator = new FilterTranslator(1);
    filterTranslator.translate(mockFilter, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockFilter);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockFilter);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec filterSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(filterSpec);
    assertEquals(filterSpec.getOpCode(), OperatorSpec.OpCode.FILTER);
    // Verify that the describe() method will establish the context for the filter function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    filterSpec.getTransformFn().init(mockContext);
    FilterFunction filterFn = (FilterFunction) Whitebox.getInternalState(filterSpec, "filterFn");
    assertNotNull(filterFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(filterFn, "translatorContext"));
    assertEquals(mockFilter, Whitebox.getInternalState(filterFn, "filter"));
    assertEquals(mockExpr, Whitebox.getInternalState(filterFn, "expr"));
    // Verify MetricsRegistry works with Project
    assertEquals(1, metricsRegistry.getGauges().size());
    assertTrue(metricsRegistry.getGauges().get(LOGICAL_OP_ID).size() > 0);
    assertEquals(1, metricsRegistry.getCounters().size());
    assertEquals(3, metricsRegistry.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling filterFn.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];
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(true);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertTrue(filterFn.apply(mockInputMsg));
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(false);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertFalse(filterFn.apply(mockInputMsg));
    // Verify filterFn.apply() updates the MetricsRegistry metrics
    assertEquals(2, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}
Also used : MessageStreamImpl(org.apache.samza.operators.MessageStreamImpl) FilterFunction(org.apache.samza.operators.functions.FilterFunction) HashMap(java.util.HashMap) ContainerContext(org.apache.samza.context.ContainerContext) StreamOperatorSpec(org.apache.samza.operators.spec.StreamOperatorSpec) DataContext(org.apache.calcite.DataContext) StreamApplicationDescriptorImpl(org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl) SamzaSqlApplicationContext(org.apache.samza.sql.runner.SamzaSqlApplicationContext) MessageStream(org.apache.samza.operators.MessageStream) ContainerContext(org.apache.samza.context.ContainerContext) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) DataContext(org.apache.calcite.DataContext) Context(org.apache.samza.context.Context) SamzaSqlApplicationContext(org.apache.samza.sql.runner.SamzaSqlApplicationContext) SamzaSqlRelMsgMetadata(org.apache.samza.sql.data.SamzaSqlRelMsgMetadata) LogicalFilter(org.apache.calcite.rel.logical.LogicalFilter) RexToJavaCompiler(org.apache.samza.sql.data.RexToJavaCompiler) TestMetricsRegistryImpl(org.apache.samza.sql.util.TestMetricsRegistryImpl) RelNode(org.apache.calcite.rel.RelNode) Expression(org.apache.samza.sql.data.Expression) SamzaSqlExecutionContext(org.apache.samza.sql.data.SamzaSqlExecutionContext) SamzaSqlRelMessage(org.apache.samza.sql.data.SamzaSqlRelMessage) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Aggregations

FilterFunction (org.apache.samza.operators.functions.FilterFunction)2 Test (org.junit.Test)2 ArrayList (java.util.ArrayList)1 Collection (java.util.Collection)1 HashMap (java.util.HashMap)1 List (java.util.List)1 DataContext (org.apache.calcite.DataContext)1 RelNode (org.apache.calcite.rel.RelNode)1 LogicalFilter (org.apache.calcite.rel.logical.LogicalFilter)1 StreamApplicationDescriptorImpl (org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl)1 ContainerContext (org.apache.samza.context.ContainerContext)1 Context (org.apache.samza.context.Context)1 KV (org.apache.samza.operators.KV)1 MessageStream (org.apache.samza.operators.MessageStream)1 MessageStreamImpl (org.apache.samza.operators.MessageStreamImpl)1 Scheduler (org.apache.samza.operators.Scheduler)1 TestMessageEnvelope (org.apache.samza.operators.data.TestMessageEnvelope)1 TestOutputMessageEnvelope (org.apache.samza.operators.data.TestOutputMessageEnvelope)1 FlatMapFunction (org.apache.samza.operators.functions.FlatMapFunction)1 JoinFunction (org.apache.samza.operators.functions.JoinFunction)1