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

Example 1 with PartitionByOperatorSpec

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

the class TestMessageStreamImpl method testPartitionBy.

@Test
public void testPartitionBy() throws IOException {
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec mockOpSpec = mock(OperatorSpec.class);
    String mockOpName = "mockName";
    when(mockGraph.getNextOpId(anyObject(), anyObject())).thenReturn(mockOpName);
    OutputStreamImpl mockOutputStreamImpl = mock(OutputStreamImpl.class);
    KVSerde mockKVSerde = mock(KVSerde.class);
    IntermediateMessageStreamImpl mockIntermediateStream = mock(IntermediateMessageStreamImpl.class);
    when(mockGraph.getIntermediateStream(eq(mockOpName), eq(mockKVSerde), eq(false))).thenReturn(mockIntermediateStream);
    when(mockIntermediateStream.getOutputStream()).thenReturn(mockOutputStreamImpl);
    when(mockIntermediateStream.isKeyed()).thenReturn(true);
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec);
    MapFunction mockKeyFunction = mock(MapFunction.class);
    MapFunction mockValueFunction = mock(MapFunction.class);
    inputStream.partitionBy(mockKeyFunction, mockValueFunction, mockKVSerde, "p1");
    ArgumentCaptor<OperatorSpec> registeredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec).registerNextOperatorSpec(registeredOpCaptor.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec = registeredOpCaptor.getValue();
    assertTrue(registeredOpSpec instanceof PartitionByOperatorSpec);
    assertEquals(OpCode.PARTITION_BY, registeredOpSpec.getOpCode());
    assertEquals(mockOutputStreamImpl, ((PartitionByOperatorSpec) registeredOpSpec).getOutputStream());
    assertEquals(mockKeyFunction, ((PartitionByOperatorSpec) registeredOpSpec).getKeyFunction());
    assertEquals(mockValueFunction, ((PartitionByOperatorSpec) registeredOpSpec).getValueFunction());
}

Example 2 with PartitionByOperatorSpec

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

the class OperatorImplGraph method computeOutputToInput.

private static void computeOutputToInput(SystemStream input, OperatorSpec opSpec, Multimap<SystemStream, SystemStream> outputToInputStreams, StreamConfig streamConfig) {
    if (opSpec instanceof PartitionByOperatorSpec) {
        PartitionByOperatorSpec spec = (PartitionByOperatorSpec) opSpec;
        SystemStream systemStream = streamConfig.streamIdToSystemStream(spec.getOutputStream().getStreamId());
        outputToInputStreams.put(systemStream, input);
    } else if (opSpec instanceof BroadcastOperatorSpec) {
        BroadcastOperatorSpec spec = (BroadcastOperatorSpec) opSpec;
        SystemStream systemStream = streamConfig.streamIdToSystemStream(spec.getOutputStream().getStreamId());
        outputToInputStreams.put(systemStream, input);
    } else {
        Collection<OperatorSpec> nextOperators = opSpec.getRegisteredOperatorSpecs();
        nextOperators.forEach(spec -> computeOutputToInput(input, spec, outputToInputStreams, streamConfig));
    }
}

Example 3 with PartitionByOperatorSpec

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

the class OperatorImplGraph method createOperatorImpl.

/**
 * Creates a new {@link OperatorImpl} instance for the provided {@link OperatorSpec}.
 *
 * @param prevOperatorSpec the original {@link OperatorSpec} that produces output for {@code operatorSpec} from {@link OperatorSpecGraph}
 * @param operatorSpec  the original {@link OperatorSpec} from {@link OperatorSpecGraph}
 * @param context  the {@link Context} required to instantiate operators
 * @return  the {@link OperatorImpl} implementation instance
 */
OperatorImpl createOperatorImpl(OperatorSpec prevOperatorSpec, OperatorSpec operatorSpec, Context context) {
    Config config = context.getJobContext().getConfig();
    StreamConfig streamConfig = new StreamConfig(config);
    if (operatorSpec instanceof InputOperatorSpec) {
        return new InputOperatorImpl((InputOperatorSpec) operatorSpec);
    } else if (operatorSpec instanceof StreamOperatorSpec) {
        return new FlatmapOperatorImpl((StreamOperatorSpec) operatorSpec);
    } else if (operatorSpec instanceof SinkOperatorSpec) {
        return new SinkOperatorImpl((SinkOperatorSpec) operatorSpec);
    } else if (operatorSpec instanceof OutputOperatorSpec) {
        String streamId = ((OutputOperatorSpec) operatorSpec).getOutputStream().getStreamId();
        SystemStream systemStream = streamConfig.streamIdToSystemStream(streamId);
        return new OutputOperatorImpl((OutputOperatorSpec) operatorSpec, systemStream);
    } else if (operatorSpec instanceof PartitionByOperatorSpec) {
        String streamId = ((PartitionByOperatorSpec) operatorSpec).getOutputStream().getStreamId();
        SystemStream systemStream = streamConfig.streamIdToSystemStream(streamId);
        return new PartitionByOperatorImpl((PartitionByOperatorSpec) operatorSpec, systemStream, internalTaskContext);
    } else if (operatorSpec instanceof WindowOperatorSpec) {
        return new WindowOperatorImpl((WindowOperatorSpec) operatorSpec, clock);
    } else if (operatorSpec instanceof JoinOperatorSpec) {
        return getOrCreatePartialJoinOpImpls((JoinOperatorSpec) operatorSpec, prevOperatorSpec.equals(((JoinOperatorSpec) operatorSpec).getLeftInputOpSpec()), clock);
    } else if (operatorSpec instanceof StreamTableJoinOperatorSpec) {
        return new StreamTableJoinOperatorImpl((StreamTableJoinOperatorSpec) operatorSpec, context);
    } else if (operatorSpec instanceof SendToTableOperatorSpec) {
        return new SendToTableOperatorImpl((SendToTableOperatorSpec) operatorSpec, context);
    } else if (operatorSpec instanceof SendToTableWithUpdateOperatorSpec) {
        return new SendToTableWithUpdateOperatorImpl((SendToTableWithUpdateOperatorSpec) operatorSpec, context);
    } else if (operatorSpec instanceof BroadcastOperatorSpec) {
        String streamId = ((BroadcastOperatorSpec) operatorSpec).getOutputStream().getStreamId();
        SystemStream systemStream = streamConfig.streamIdToSystemStream(streamId);
        return new BroadcastOperatorImpl((BroadcastOperatorSpec) operatorSpec, systemStream, context);
    } else if (operatorSpec instanceof AsyncFlatMapOperatorSpec) {
        return new AsyncFlatmapOperatorImpl((AsyncFlatMapOperatorSpec) operatorSpec);
    }
    throw new IllegalArgumentException(String.format("Unsupported OperatorSpec: %s", operatorSpec.getClass().getName()));
}

Example 4 with PartitionByOperatorSpec

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

the class JobGraphJsonGenerator method operatorToMap.

/**
 * Format the operator properties into a map
 * @param spec a {@link OperatorSpec} instance
 * @return map of the operator properties
 */
@VisibleForTesting
Map<String, Object> operatorToMap(OperatorSpec spec) {
    Map<String, Object> map = new HashMap<>();
    map.put("opCode", spec.getOpCode().name());
    map.put("opId", spec.getOpId());
    map.put("sourceLocation", spec.getSourceLocation());
    Collection<OperatorSpec> nextOperators = spec.getRegisteredOperatorSpecs();
    map.put("nextOperatorIds", nextOperators.stream().map(OperatorSpec::getOpId).collect(Collectors.toSet()));
    if (spec instanceof OutputOperatorSpec) {
        OutputStreamImpl outputStream = ((OutputOperatorSpec) spec).getOutputStream();
        map.put("outputStreamId", outputStream.getStreamId());
    } else if (spec instanceof PartitionByOperatorSpec) {
        OutputStreamImpl outputStream = ((PartitionByOperatorSpec) spec).getOutputStream();
        map.put("outputStreamId", outputStream.getStreamId());
    }
    if (spec instanceof StreamTableJoinOperatorSpec) {
        String tableId = ((StreamTableJoinOperatorSpec) spec).getTableId();
        map.put("tableId", tableId);
    }
    if (spec instanceof SendToTableOperatorSpec) {
        String tableId = ((SendToTableOperatorSpec) spec).getTableId();
        map.put("tableId", tableId);
    }
    if (spec instanceof JoinOperatorSpec) {
        map.put("ttlMs", ((JoinOperatorSpec) spec).getTtlMs());
    }
    return map;
}