org.apache.flink.streaming.runtime.partitioner.StreamPartitioner Java Examples

/**
 * Tests that the max parallelism is properly set for connected
 * streams.
 */
@Test
public void testMaxParallelismWithConnectedKeyedStream() {
	int maxParallelism = 42;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128);
	DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129);

	env.getConfig().setMaxParallelism(maxParallelism);

	DataStream<Integer> keyedResult = input1
			.connect(input2)
			.keyBy(value -> value, value -> value)
			.map(new NoOpIntCoMap());

	keyedResult.addSink(new DiscardingSink<>());

	StreamGraph graph = env.getStreamGraph();

	StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());

	StreamPartitioner<?> streamPartitioner1 = keyedResultNode.getInEdges().get(0).getPartitioner();
	StreamPartitioner<?> streamPartitioner2 = keyedResultNode.getInEdges().get(1).getPartitioner();
}
 

/**
 * Tests that the max parallelism is properly set for connected
 * streams.
 */
@Test
public void testMaxParallelismWithConnectedKeyedStream() {
	int maxParallelism = 42;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128);
	DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129);

	env.getConfig().setMaxParallelism(maxParallelism);

	DataStream<Integer> keyedResult = input1
			.connect(input2)
			.keyBy(value -> value, value -> value)
			.map(new NoOpIntCoMap());

	keyedResult.addSink(new DiscardingSink<>());

	StreamGraph graph = env.getStreamGraph();

	StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());

	StreamPartitioner<?> streamPartitioner1 = keyedResultNode.getInEdges().get(0).getPartitioner();
	StreamPartitioner<?> streamPartitioner2 = keyedResultNode.getInEdges().get(1).getPartitioner();
}
 

public StreamEdge(StreamNode sourceVertex, StreamNode targetVertex, int typeNumber,
		List<String> selectedNames, StreamPartitioner<?> outputPartitioner, OutputTag outputTag,
		ShuffleMode shuffleMode) {
	this.sourceId = sourceVertex.getId();
	this.targetId = targetVertex.getId();
	this.typeNumber = typeNumber;
	this.selectedNames = selectedNames;
	this.outputPartitioner = outputPartitioner;
	this.outputTag = outputTag;
	this.sourceOperatorName = sourceVertex.getOperatorName();
	this.targetOperatorName = targetVertex.getOperatorName();
	this.shuffleMode = checkNotNull(shuffleMode);

	this.edgeId = sourceVertex + "_" + targetVertex + "_" + typeNumber + "_" + selectedNames
			+ "_" + outputPartitioner;
}
 

public StreamEdge(StreamNode sourceVertex, StreamNode targetVertex, int typeNumber,
		List<String> selectedNames, StreamPartitioner<?> outputPartitioner, OutputTag outputTag,
		ShuffleMode shuffleMode) {
	this.sourceId = sourceVertex.getId();
	this.targetId = targetVertex.getId();
	this.typeNumber = typeNumber;
	this.selectedNames = selectedNames;
	this.outputPartitioner = outputPartitioner;
	this.outputTag = outputTag;
	this.sourceOperatorName = sourceVertex.getOperatorName();
	this.targetOperatorName = targetVertex.getOperatorName();
	this.shuffleMode = checkNotNull(shuffleMode);

	this.edgeId = sourceVertex + "_" + targetVertex + "_" + typeNumber + "_" + selectedNames
			+ "_" + outputPartitioner;
}
 

/**
 * Tests that the max parallelism is properly set for connected
 * streams.
 */
@Test
public void testMaxParallelismWithConnectedKeyedStream() {
	int maxParallelism = 42;

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128);
	DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129);

	env.getConfig().setMaxParallelism(maxParallelism);

	DataStream<Integer> keyedResult = input1
			.connect(input2)
			.keyBy(value -> value, value -> value)
			.map(new NoOpIntCoMap());

	keyedResult.addSink(new DiscardingSink<>());

	StreamGraph graph = env.getStreamGraph();

	StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());

	StreamPartitioner<?> streamPartitioner1 = keyedResultNode.getInEdges().get(0).getPartitioner();
	StreamPartitioner<?> streamPartitioner2 = keyedResultNode.getInEdges().get(1).getPartitioner();
}
 

private static <OUT> RecordWriter<SerializationDelegate<StreamRecord<OUT>>> createRecordWriter(
		StreamEdge edge,
		int outputIndex,
		Environment environment,
		String taskName,
		long bufferTimeout) {
	@SuppressWarnings("unchecked")
	StreamPartitioner<OUT> outputPartitioner = (StreamPartitioner<OUT>) edge.getPartitioner();

	LOG.debug("Using partitioner {} for output {} of task {}", outputPartitioner, outputIndex, taskName);

	ResultPartitionWriter bufferWriter = environment.getWriter(outputIndex);

	// we initialize the partitioner here with the number of key groups (aka max. parallelism)
	if (outputPartitioner instanceof ConfigurableStreamPartitioner) {
		int numKeyGroups = bufferWriter.getNumTargetKeyGroups();
		if (0 < numKeyGroups) {
			((ConfigurableStreamPartitioner) outputPartitioner).configure(numKeyGroups);
		}
	}

	RecordWriter<SerializationDelegate<StreamRecord<OUT>>> output =
		RecordWriter.createRecordWriter(bufferWriter, outputPartitioner, bufferTimeout, taskName);
	output.setMetricGroup(environment.getMetricGroup().getIOMetricGroup());
	return output;
}
 

private ResultPartitionType determineResultPartitionType(StreamPartitioner<?> partitioner) {
	switch (streamGraph.getGlobalDataExchangeMode()) {
		case ALL_EDGES_BLOCKING:
			return ResultPartitionType.BLOCKING;
		case FORWARD_EDGES_PIPELINED:
			if (partitioner instanceof ForwardPartitioner) {
				return ResultPartitionType.PIPELINED_BOUNDED;
			} else {
				return ResultPartitionType.BLOCKING;
			}
		case POINTWISE_EDGES_PIPELINED:
			if (isPointwisePartitioner(partitioner)) {
				return ResultPartitionType.PIPELINED_BOUNDED;
			} else {
				return ResultPartitionType.BLOCKING;
			}
		case ALL_EDGES_PIPELINED:
			return ResultPartitionType.PIPELINED_BOUNDED;
		default:
			throw new RuntimeException("Unrecognized global data exchange mode " + streamGraph.getGlobalDataExchangeMode());
	}
}
 

private void connect(Integer headOfChain, StreamEdge edge) {

		physicalEdgesInOrder.add(edge);

		Integer downStreamvertexID = edge.getTargetId();

		JobVertex headVertex = jobVertices.get(headOfChain);
		JobVertex downStreamVertex = jobVertices.get(downStreamvertexID);

		StreamConfig downStreamConfig = new StreamConfig(downStreamVertex.getConfiguration());

		downStreamConfig.setNumberOfInputs(downStreamConfig.getNumberOfInputs() + 1);

		StreamPartitioner<?> partitioner = edge.getPartitioner();
		JobEdge jobEdge;
		if (partitioner instanceof ForwardPartitioner || partitioner instanceof RescalePartitioner) {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
				headVertex,
				DistributionPattern.POINTWISE,
				ResultPartitionType.PIPELINED_BOUNDED);
		} else {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
					headVertex,
					DistributionPattern.ALL_TO_ALL,
					ResultPartitionType.PIPELINED_BOUNDED);
		}
		// set strategy name so that web interface can show it.
		jobEdge.setShipStrategyName(partitioner.toString());

		if (LOG.isDebugEnabled()) {
			LOG.debug("CONNECTED: {} - {} -> {}", partitioner.getClass().getSimpleName(),
					headOfChain, downStreamvertexID);
		}
	}
 

private static <OUT> RecordWriter<SerializationDelegate<StreamRecord<OUT>>> createRecordWriter(
		StreamEdge edge,
		int outputIndex,
		Environment environment,
		String taskName,
		long bufferTimeout) {
	@SuppressWarnings("unchecked")
	StreamPartitioner<OUT> outputPartitioner = (StreamPartitioner<OUT>) edge.getPartitioner();

	LOG.debug("Using partitioner {} for output {} of task {}", outputPartitioner, outputIndex, taskName);

	ResultPartitionWriter bufferWriter = environment.getWriter(outputIndex);

	// we initialize the partitioner here with the number of key groups (aka max. parallelism)
	if (outputPartitioner instanceof ConfigurableStreamPartitioner) {
		int numKeyGroups = bufferWriter.getNumTargetKeyGroups();
		if (0 < numKeyGroups) {
			((ConfigurableStreamPartitioner) outputPartitioner).configure(numKeyGroups);
		}
	}

	RecordWriter<SerializationDelegate<StreamRecord<OUT>>> output = new RecordWriterBuilder<SerializationDelegate<StreamRecord<OUT>>>()
		.setChannelSelector(outputPartitioner)
		.setTimeout(bufferTimeout)
		.setTaskName(taskName)
		.build(bufferWriter);
	output.setMetricGroup(environment.getMetricGroup().getIOMetricGroup());
	return output;
}
 

private static <OUT> RecordWriter<SerializationDelegate<StreamRecord<OUT>>> createRecordWriter(
		StreamEdge edge,
		int outputIndex,
		Environment environment,
		String taskName,
		long bufferTimeout) {
	@SuppressWarnings("unchecked")
	StreamPartitioner<OUT> outputPartitioner = (StreamPartitioner<OUT>) edge.getPartitioner();

	LOG.debug("Using partitioner {} for output {} of task {}", outputPartitioner, outputIndex, taskName);

	ResultPartitionWriter bufferWriter = environment.getWriter(outputIndex);

	// we initialize the partitioner here with the number of key groups (aka max. parallelism)
	if (outputPartitioner instanceof ConfigurableStreamPartitioner) {
		int numKeyGroups = bufferWriter.getNumTargetKeyGroups();
		if (0 < numKeyGroups) {
			((ConfigurableStreamPartitioner) outputPartitioner).configure(numKeyGroups);
		}
	}

	RecordWriter<SerializationDelegate<StreamRecord<OUT>>> output = new RecordWriterBuilder()
		.setChannelSelector(outputPartitioner)
		.setTimeout(bufferTimeout)
		.setTaskName(taskName)
		.build(bufferWriter);
	output.setMetricGroup(environment.getMetricGroup().getIOMetricGroup());
	return output;
}
 

/**
 * Adds a new virtual node that is used to connect a downstream vertex to an input with a
 * certain partitioning.
 *
 * <p>When adding an edge from the virtual node to a downstream node the connection will be made
 * to the original node, but with the partitioning given here.
 *
 * @param originalId ID of the node that should be connected to.
 * @param virtualId ID of the virtual node.
 * @param partitioner The partitioner
 */
public void addVirtualPartitionNode(
		Integer originalId,
		Integer virtualId,
		StreamPartitioner<?> partitioner,
		ShuffleMode shuffleMode) {

	if (virtualPartitionNodes.containsKey(virtualId)) {
		throw new IllegalStateException("Already has virtual partition node with id " + virtualId);
	}

	virtualPartitionNodes.put(virtualId, new Tuple3<>(originalId, partitioner, shuffleMode));
}
 

public StreamEdge(StreamNode sourceVertex, StreamNode targetVertex, int typeNumber,
		List<String> selectedNames, StreamPartitioner<?> outputPartitioner, OutputTag outputTag) {
	this.sourceId = sourceVertex.getId();
	this.targetId = targetVertex.getId();
	this.typeNumber = typeNumber;
	this.selectedNames = selectedNames;
	this.outputPartitioner = outputPartitioner;
	this.outputTag = outputTag;
	this.sourceOperatorName = sourceVertex.getOperatorName();
	this.targetOperatorName = targetVertex.getOperatorName();

	this.edgeId = sourceVertex + "_" + targetVertex + "_" + typeNumber + "_" + selectedNames
			+ "_" + outputPartitioner;
}
 

/**
 * Adds a new virtual node that is used to connect a downstream vertex to an input with a
 * certain partitioning.
 *
 * <p>When adding an edge from the virtual node to a downstream node the connection will be made
 * to the original node, but with the partitioning given here.
 *
 * @param originalId ID of the node that should be connected to.
 * @param virtualId ID of the virtual node.
 * @param partitioner The partitioner
 */
public void addVirtualPartitionNode(
		Integer originalId,
		Integer virtualId,
		StreamPartitioner<?> partitioner,
		ShuffleMode shuffleMode) {

	if (virtualPartitionNodes.containsKey(virtualId)) {
		throw new IllegalStateException("Already has virtual partition node with id " + virtualId);
	}

	virtualPartitionNodes.put(virtualId, new Tuple3<>(originalId, partitioner, shuffleMode));
}
 

public StreamEdge(StreamNode sourceVertex, StreamNode targetVertex, int typeNumber,
		List<String> selectedNames, StreamPartitioner<?> outputPartitioner, OutputTag outputTag) {
	this(sourceVertex,
			targetVertex,
			typeNumber,
			selectedNames,
			outputPartitioner,
			outputTag,
			ShuffleMode.UNDEFINED);
}
 

/**
 * Creates a new {@code PartitionTransformation} from the given input and
 * {@link StreamPartitioner}.
 *
 * @param input The input {@code Transformation}
 * @param partitioner The {@code StreamPartitioner}
 * @param shuffleMode The {@code ShuffleMode}
 */
public PartitionTransformation(
		Transformation<T> input,
		StreamPartitioner<T> partitioner,
		ShuffleMode shuffleMode) {
	super("Partition", input.getOutputType(), input.getParallelism());
	this.input = input;
	this.partitioner = partitioner;
	this.shuffleMode = checkNotNull(shuffleMode);
}
 

/**
 * Creates a new {@code PartitionTransformation} from the given input and
 * {@link StreamPartitioner}.
 *
 * @param input The input {@code Transformation}
 * @param partitioner The {@code StreamPartitioner}
 * @param shuffleMode The {@code ShuffleMode}
 */
public PartitionTransformation(
		Transformation<T> input,
		StreamPartitioner<T> partitioner,
		ShuffleMode shuffleMode) {
	super("Partition", input.getOutputType(), input.getParallelism());
	this.input = input;
	this.partitioner = partitioner;
	this.shuffleMode = checkNotNull(shuffleMode);
}
 

public StreamEdge(StreamNode sourceVertex, StreamNode targetVertex, int typeNumber,
		List<String> selectedNames, StreamPartitioner<?> outputPartitioner, OutputTag outputTag) {
	this(sourceVertex,
			targetVertex,
			typeNumber,
			selectedNames,
			outputPartitioner,
			outputTag,
			ShuffleMode.UNDEFINED);
}
 

private void connect(Integer headOfChain, StreamEdge edge) {

		physicalEdgesInOrder.add(edge);

		Integer downStreamvertexID = edge.getTargetId();

		JobVertex headVertex = jobVertices.get(headOfChain);
		JobVertex downStreamVertex = jobVertices.get(downStreamvertexID);

		StreamConfig downStreamConfig = new StreamConfig(downStreamVertex.getConfiguration());

		downStreamConfig.setNumberOfInputs(downStreamConfig.getNumberOfInputs() + 1);

		StreamPartitioner<?> partitioner = edge.getPartitioner();

		ResultPartitionType resultPartitionType;
		switch (edge.getShuffleMode()) {
			case PIPELINED:
				resultPartitionType = ResultPartitionType.PIPELINED_BOUNDED;
				break;
			case BATCH:
				resultPartitionType = ResultPartitionType.BLOCKING;
				break;
			case UNDEFINED:
				resultPartitionType = streamGraph.isBlockingConnectionsBetweenChains() ?
						ResultPartitionType.BLOCKING : ResultPartitionType.PIPELINED_BOUNDED;
				break;
			default:
				throw new UnsupportedOperationException("Data exchange mode " +
					edge.getShuffleMode() + " is not supported yet.");
		}

		JobEdge jobEdge;
		if (partitioner instanceof ForwardPartitioner || partitioner instanceof RescalePartitioner) {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
				headVertex,
				DistributionPattern.POINTWISE,
				resultPartitionType);
		} else {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
					headVertex,
					DistributionPattern.ALL_TO_ALL,
					resultPartitionType);
		}
		// set strategy name so that web interface can show it.
		jobEdge.setShipStrategyName(partitioner.toString());

		if (LOG.isDebugEnabled()) {
			LOG.debug("CONNECTED: {} - {} -> {}", partitioner.getClass().getSimpleName(),
					headOfChain, downStreamvertexID);
		}
	}
 

@Override
public StreamPartitioner<Tuple2<StreamRoute, Object>> copy() {
    return new DynamicPartitioner();
}
 

@Override
public StreamPartitioner<RowData> copy() {
	return this;
}
 

@Override
protected DataStream<T> setConnectionType(StreamPartitioner<T> partitioner) {
	throw new UnsupportedOperationException("Cannot override partitioning for KeyedStream.");
}
 

/**
 * Tests that the KeyGroupStreamPartitioner are properly set up with the correct value of
 * maximum parallelism.
 */
@Test
public void testSetupOfKeyGroupPartitioner() {
	int maxParallelism = 42;
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.getConfig().setMaxParallelism(maxParallelism);

	DataStream<Integer> source = env.fromElements(1, 2, 3);

	DataStream<Integer> keyedResult = source.keyBy(value -> value).map(new NoOpIntMap());

	keyedResult.addSink(new DiscardingSink<>());

	StreamGraph graph = env.getStreamGraph();

	StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());

	StreamPartitioner<?> streamPartitioner = keyedResultNode.getInEdges().get(0).getPartitioner();
}
 

@Test
public void testChannelSelectors() {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	DataStreamSource<Long> src = env.generateSequence(0, 0);

	DataStream<Long> broadcast = src.broadcast();
	DataStreamSink<Long> broadcastSink = broadcast.print();
	StreamPartitioner<?> broadcastPartitioner =
			getStreamGraph(env).getStreamEdges(src.getId(),
					broadcastSink.getTransformation().getId()).get(0).getPartitioner();
	assertTrue(broadcastPartitioner instanceof BroadcastPartitioner);

	DataStream<Long> shuffle = src.shuffle();
	DataStreamSink<Long> shuffleSink = shuffle.print();
	StreamPartitioner<?> shufflePartitioner =
			getStreamGraph(env).getStreamEdges(src.getId(),
					shuffleSink.getTransformation().getId()).get(0).getPartitioner();
	assertTrue(shufflePartitioner instanceof ShufflePartitioner);

	DataStream<Long> forward = src.forward();
	DataStreamSink<Long> forwardSink = forward.print();
	StreamPartitioner<?> forwardPartitioner =
			getStreamGraph(env).getStreamEdges(src.getId(),
					forwardSink.getTransformation().getId()).get(0).getPartitioner();
	assertTrue(forwardPartitioner instanceof ForwardPartitioner);

	DataStream<Long> rebalance = src.rebalance();
	DataStreamSink<Long> rebalanceSink = rebalance.print();
	StreamPartitioner<?> rebalancePartitioner =
			getStreamGraph(env).getStreamEdges(src.getId(),
					rebalanceSink.getTransformation().getId()).get(0).getPartitioner();
	assertTrue(rebalancePartitioner instanceof RebalancePartitioner);

	DataStream<Long> global = src.global();
	DataStreamSink<Long> globalSink = global.print();
	StreamPartitioner<?> globalPartitioner =
			getStreamGraph(env).getStreamEdges(src.getId(),
					globalSink.getTransformation().getId()).get(0).getPartitioner();
	assertTrue(globalPartitioner instanceof GlobalPartitioner);
}
 

public StreamPartitioner<?> getPartitioner() {
	return outputPartitioner;
}
 

public void setPartitioner(StreamPartitioner<?> partitioner) {
	this.outputPartitioner = partitioner;
}
 

private void connect(Integer headOfChain, StreamEdge edge) {

		physicalEdgesInOrder.add(edge);

		Integer downStreamVertexID = edge.getTargetId();

		JobVertex headVertex = jobVertices.get(headOfChain);
		JobVertex downStreamVertex = jobVertices.get(downStreamVertexID);

		StreamConfig downStreamConfig = new StreamConfig(downStreamVertex.getConfiguration());

		downStreamConfig.setNumberOfInputs(downStreamConfig.getNumberOfInputs() + 1);

		StreamPartitioner<?> partitioner = edge.getPartitioner();

		ResultPartitionType resultPartitionType;
		switch (edge.getShuffleMode()) {
			case PIPELINED:
				resultPartitionType = ResultPartitionType.PIPELINED_BOUNDED;
				break;
			case BATCH:
				resultPartitionType = ResultPartitionType.BLOCKING;
				break;
			case UNDEFINED:
				resultPartitionType = determineResultPartitionType(partitioner);
				break;
			default:
				throw new UnsupportedOperationException("Data exchange mode " +
					edge.getShuffleMode() + " is not supported yet.");
		}

		JobEdge jobEdge;
		if (isPointwisePartitioner(partitioner)) {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
				headVertex,
				DistributionPattern.POINTWISE,
				resultPartitionType);
		} else {
			jobEdge = downStreamVertex.connectNewDataSetAsInput(
					headVertex,
					DistributionPattern.ALL_TO_ALL,
					resultPartitionType);
		}
		// set strategy name so that web interface can show it.
		jobEdge.setShipStrategyName(partitioner.toString());

		if (LOG.isDebugEnabled()) {
			LOG.debug("CONNECTED: {} - {} -> {}", partitioner.getClass().getSimpleName(),
					headOfChain, downStreamVertexID);
		}
	}
 

private static boolean isPointwisePartitioner(StreamPartitioner<?> partitioner) {
	return partitioner instanceof ForwardPartitioner || partitioner instanceof RescalePartitioner;
}
 

@Override
protected DataStream<T> setConnectionType(StreamPartitioner<T> partitioner) {
	throw new UnsupportedOperationException("Cannot override partitioning for KeyedStream.");
}
 

@Override
protected DataStream<T> setConnectionType(StreamPartitioner<T> partitioner) {
	throw new UnsupportedOperationException("Cannot override partitioning for KeyedStream.");
}
 

public StreamPartitioner<?> getPartitioner() {
	return outputPartitioner;
}