org.apache.flink.runtime.plugable.SerializationDelegate Java Examples

@SuppressWarnings("unchecked")
public RecordWriterOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		TypeSerializer<OUT> outSerializer,
		OutputTag outputTag,
		StreamStatusProvider streamStatusProvider) {

	checkNotNull(recordWriter);
	this.outputTag = outputTag;
	// generic hack: cast the writer to generic Object type so we can use it
	// with multiplexed records and watermarks
	this.recordWriter = (RecordWriter<SerializationDelegate<StreamElement>>)
			(RecordWriter<?>) recordWriter;

	TypeSerializer<StreamElement> outRecordSerializer =
			new StreamElementSerializer<>(outSerializer);

	if (outSerializer != null) {
		serializationDelegate = new SerializationDelegate<StreamElement>(outRecordSerializer);
	}

	this.streamStatusProvider = checkNotNull(streamStatusProvider);
}
 

@VisibleForTesting
public static <OUT> List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> createRecordWriters(
		StreamConfig configuration,
		Environment environment) {
	List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> recordWriters = new ArrayList<>();
	List<StreamEdge> outEdgesInOrder = configuration.getOutEdgesInOrder(environment.getUserClassLoader());
	Map<Integer, StreamConfig> chainedConfigs = configuration.getTransitiveChainedTaskConfigsWithSelf(environment.getUserClassLoader());

	for (int i = 0; i < outEdgesInOrder.size(); i++) {
		StreamEdge edge = outEdgesInOrder.get(i);
		recordWriters.add(
			createRecordWriter(
				edge,
				i,
				environment,
				environment.getTaskInfo().getTaskName(),
				chainedConfigs.get(edge.getSourceId()).getBufferTimeout()));
	}
	return recordWriters;
}
 

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;
}
 

@Test
public void testPartitionHash() {
	// Test for IntValue
	verifyPartitionHashSelectedChannels(50000, 100, RecordType.INTEGER);
	// Test for StringValue
	verifyPartitionHashSelectedChannels(10000, 100, RecordType.STRING);

	// Test hash corner cases
	final TestIntComparator testIntComp = new TestIntComparator();
	final ChannelSelector<SerializationDelegate<Integer>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, testIntComp, 100);
	final SerializationDelegate<Integer> serializationDelegate = new SerializationDelegate<>(new IntSerializer());

	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MIN_VALUE, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, -1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 0, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MAX_VALUE, 100);
}
 

@SuppressWarnings("unchecked")
public RecordWriterOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		TypeSerializer<OUT> outSerializer,
		OutputTag outputTag,
		StreamStatusProvider streamStatusProvider) {

	checkNotNull(recordWriter);
	this.outputTag = outputTag;
	// generic hack: cast the writer to generic Object type so we can use it
	// with multiplexed records and watermarks
	this.recordWriter = (RecordWriter<SerializationDelegate<StreamElement>>)
			(RecordWriter<?>) recordWriter;

	TypeSerializer<StreamElement> outRecordSerializer =
			new StreamElementSerializer<>(outSerializer);

	if (outSerializer != null) {
		serializationDelegate = new SerializationDelegate<StreamElement>(outRecordSerializer);
	}

	this.streamStatusProvider = checkNotNull(streamStatusProvider);
}
 

private RecordWriterOutput<OUT> createStreamOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		StreamEdge edge,
		StreamConfig upStreamConfig,
		Environment taskEnvironment) {
	OutputTag sideOutputTag = edge.getOutputTag(); // OutputTag, return null if not sideOutput

	TypeSerializer outSerializer = null;

	if (edge.getOutputTag() != null) {
		// side output
		outSerializer = upStreamConfig.getTypeSerializerSideOut(
				edge.getOutputTag(), taskEnvironment.getUserClassLoader());
	} else {
		// main output
		outSerializer = upStreamConfig.getTypeSerializerOut(taskEnvironment.getUserClassLoader());
	}

	return new RecordWriterOutput<>(recordWriter, outSerializer, sideOutputTag, this);
}
 

private boolean verifyWrongPartitionHashKey(int position, int fieldNum) {
	final TypeComparator<Record> comparator = new RecordComparatorFactory(
		new int[] {position}, new Class[] {IntValue.class}).createComparator();
	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, comparator, 100);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	Record record = new Record(2);
	record.setField(fieldNum, new IntValue(1));
	delegate.setInstance(record);

	try {
		selector.selectChannel(delegate);
	} catch (NullKeyFieldException re) {
		Assert.assertEquals(position, re.getFieldNumber());
		return true;
	}
	return false;
}
 

@SuppressWarnings("unchecked")
public RecordWriterOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		TypeSerializer<OUT> outSerializer,
		OutputTag outputTag,
		StreamStatusProvider streamStatusProvider) {

	checkNotNull(recordWriter);
	this.outputTag = outputTag;
	// generic hack: cast the writer to generic Object type so we can use it
	// with multiplexed records and watermarks
	this.recordWriter = (RecordWriter<SerializationDelegate<StreamElement>>)
			(RecordWriter<?>) recordWriter;

	TypeSerializer<StreamElement> outRecordSerializer =
			new StreamElementSerializer<>(outSerializer);

	if (outSerializer != null) {
		serializationDelegate = new SerializationDelegate<StreamElement>(outRecordSerializer);
	}

	this.streamStatusProvider = checkNotNull(streamStatusProvider);
}
 

private RecordWriterOutput<OUT> createStreamOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		StreamEdge edge,
		StreamConfig upStreamConfig,
		Environment taskEnvironment) {
	OutputTag sideOutputTag = edge.getOutputTag(); // OutputTag, return null if not sideOutput

	TypeSerializer outSerializer = null;

	if (edge.getOutputTag() != null) {
		// side output
		outSerializer = upStreamConfig.getTypeSerializerSideOut(
				edge.getOutputTag(), taskEnvironment.getUserClassLoader());
	} else {
		// main output
		outSerializer = upStreamConfig.getTypeSerializerOut(taskEnvironment.getUserClassLoader());
	}

	return new RecordWriterOutput<>(recordWriter, outSerializer, sideOutputTag, this);
}
 

private static <OUT> List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> createRecordWriters(
		StreamConfig configuration,
		Environment environment) {
	List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> recordWriters = new ArrayList<>();
	List<StreamEdge> outEdgesInOrder = configuration.getOutEdgesInOrder(environment.getUserClassLoader());
	Map<Integer, StreamConfig> chainedConfigs = configuration.getTransitiveChainedTaskConfigsWithSelf(environment.getUserClassLoader());

	for (int i = 0; i < outEdgesInOrder.size(); i++) {
		StreamEdge edge = outEdgesInOrder.get(i);
		recordWriters.add(
			createRecordWriter(
				edge,
				i,
				environment,
				environment.getTaskInfo().getTaskName(),
				chainedConfigs.get(edge.getSourceId()).getBufferTimeout()));
	}
	return recordWriters;
}
 

@Test
public void testPartitionHash() {
	// Test for IntValue
	verifyPartitionHashSelectedChannels(50000, 100, RecordType.INTEGER);
	// Test for StringValue
	verifyPartitionHashSelectedChannels(10000, 100, RecordType.STRING);

	// Test hash corner cases
	final TestIntComparator testIntComp = new TestIntComparator();
	final ChannelSelector<SerializationDelegate<Integer>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, testIntComp, 100);
	final SerializationDelegate<Integer> serializationDelegate = new SerializationDelegate<>(new IntSerializer());

	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MIN_VALUE, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, -1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 0, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MAX_VALUE, 100);
}
 

@Override
public final int selectChannel(SerializationDelegate<T> record) {
	switch (strategy) {
	case FORWARD:
		return forward();
	case PARTITION_RANDOM:
	case PARTITION_FORCED_REBALANCE:
		return robin(numberOfChannels);
	case PARTITION_HASH:
		return hashPartitionDefault(record.getInstance(), numberOfChannels);
	case PARTITION_CUSTOM:
		return customPartition(record.getInstance(), numberOfChannels);
	case PARTITION_RANGE:
		return rangePartition(record.getInstance(), numberOfChannels);
	default:
		throw new UnsupportedOperationException("Unsupported distribution strategy: " + strategy.name());
	}
}
 

@Override
public final int selectChannel(SerializationDelegate<T> record) {
	switch (strategy) {
	case FORWARD:
		return forward();
	case PARTITION_RANDOM:
	case PARTITION_FORCED_REBALANCE:
		return robin(numberOfChannels);
	case PARTITION_HASH:
		return hashPartitionDefault(record.getInstance(), numberOfChannels);
	case PARTITION_CUSTOM:
		return customPartition(record.getInstance(), numberOfChannels);
	case PARTITION_RANGE:
		return rangePartition(record.getInstance(), numberOfChannels);
	default:
		throw new UnsupportedOperationException("Unsupported distribution strategy: " + strategy.name());
	}
}
 

@Test
public void testPartitionHash() {
	// Test for IntValue
	verifyPartitionHashSelectedChannels(50000, 100, RecordType.INTEGER);
	// Test for StringValue
	verifyPartitionHashSelectedChannels(10000, 100, RecordType.STRING);

	// Test hash corner cases
	final TestIntComparator testIntComp = new TestIntComparator();
	final ChannelSelector<SerializationDelegate<Integer>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, testIntComp, 100);
	final SerializationDelegate<Integer> serializationDelegate = new SerializationDelegate<>(new IntSerializer());

	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MIN_VALUE, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, -1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 0, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, 1, 100);
	assertPartitionHashSelectedChannels(selector, serializationDelegate, Integer.MAX_VALUE, 100);
}
 

private RecordWriterOutput<OUT> createStreamOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		StreamEdge edge,
		StreamConfig upStreamConfig,
		Environment taskEnvironment) {
	OutputTag sideOutputTag = edge.getOutputTag(); // OutputTag, return null if not sideOutput

	TypeSerializer outSerializer = null;

	if (edge.getOutputTag() != null) {
		// side output
		outSerializer = upStreamConfig.getTypeSerializerSideOut(
				edge.getOutputTag(), taskEnvironment.getUserClassLoader());
	} else {
		// main output
		outSerializer = upStreamConfig.getTypeSerializerOut(taskEnvironment.getUserClassLoader());
	}

	return new RecordWriterOutput<>(recordWriter, outSerializer, sideOutputTag, this);
}
 

@VisibleForTesting
public static <OUT> List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> createRecordWriters(
		StreamConfig configuration,
		Environment environment) {
	List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> recordWriters = new ArrayList<>();
	List<StreamEdge> outEdgesInOrder = configuration.getOutEdgesInOrder(environment.getUserClassLoader());
	Map<Integer, StreamConfig> chainedConfigs = configuration.getTransitiveChainedTaskConfigsWithSelf(environment.getUserClassLoader());

	for (int i = 0; i < outEdgesInOrder.size(); i++) {
		StreamEdge edge = outEdgesInOrder.get(i);
		recordWriters.add(
			createRecordWriter(
				edge,
				i,
				environment,
				environment.getTaskInfo().getTaskName(),
				chainedConfigs.get(edge.getSourceId()).getBufferTimeout()));
	}
	return recordWriters;
}
 

private boolean verifyWrongPartitionHashKey(int position, int fieldNum) {
	final TypeComparator<Record> comparator = new RecordComparatorFactory(
		new int[] {position}, new Class[] {IntValue.class}).createComparator();
	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, comparator, 100);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	Record record = new Record(2);
	record.setField(fieldNum, new IntValue(1));
	delegate.setInstance(record);

	try {
		selector.selectChannel(delegate);
	} catch (NullKeyFieldException re) {
		Assert.assertEquals(position, re.getFieldNumber());
		return true;
	}
	return false;
}
 

private boolean verifyWrongPartitionHashKey(int position, int fieldNum) {
	final TypeComparator<Record> comparator = new RecordComparatorFactory(
		new int[] {position}, new Class[] {IntValue.class}).createComparator();
	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, comparator, 100);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	Record record = new Record(2);
	record.setField(fieldNum, new IntValue(1));
	delegate.setInstance(record);

	try {
		selector.selectChannel(delegate);
	} catch (NullKeyFieldException re) {
		Assert.assertEquals(position, re.getFieldNumber());
		return true;
	}
	return false;
}
 

@Override
public final int selectChannel(SerializationDelegate<T> record) {
	switch (strategy) {
	case FORWARD:
		return forward();
	case PARTITION_RANDOM:
	case PARTITION_FORCED_REBALANCE:
		return robin(numberOfChannels);
	case PARTITION_HASH:
		return hashPartitionDefault(record.getInstance(), numberOfChannels);
	case PARTITION_CUSTOM:
		return customPartition(record.getInstance(), numberOfChannels);
	case PARTITION_RANGE:
		return rangePartition(record.getInstance(), numberOfChannels);
	default:
		throw new UnsupportedOperationException("Unsupported distribution strategy: " + strategy.name());
	}
}
 

@Test
public void testMultiKeys() {
	final int numberOfChannels = 100;
	final int numRecords = 5000;
	final TypeComparator<Record> multiComp = new RecordComparatorFactory(
		new int[] {0,1, 3}, new Class[] {IntValue.class, StringValue.class, DoubleValue.class}).createComparator();

	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, multiComp, numberOfChannels);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	int[] hits = new int[numberOfChannels];
	for (int i = 0; i < numRecords; i++) {
		Record record = new Record(4);
		record.setField(0, new IntValue(i));
		record.setField(1, new StringValue("AB" + i + "CD" + i));
		record.setField(3, new DoubleValue(i * 3.141d));
		delegate.setInstance(record);

		int channel = selector.selectChannel(delegate);
		hits[channel]++;
	}

	int totalHitCount = 0;
	for (int hit : hits) {
		assertTrue(hit > 0);
		totalHitCount += hit;
	}
	assertTrue(totalHitCount == numRecords);
}
 

private int[] getSelectedChannelsHitCount(
		ShipStrategyType shipStrategyType,
		int numRecords,
		int numberOfChannels,
		Enum recordType) {
	final TypeComparator<Record> comparator = new RecordComparatorFactory(
		new int[] {0}, new Class[] {recordType == RecordType.INTEGER ? IntValue.class : StringValue.class}).createComparator();
	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(shipStrategyType, comparator, numberOfChannels);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	return getSelectedChannelsHitCount(selector, delegate, recordType, numRecords, numberOfChannels);
}
 

@Test
public void testMultiKeys() {
	final int numberOfChannels = 100;
	final int numRecords = 5000;
	final TypeComparator<Record> multiComp = new RecordComparatorFactory(
		new int[] {0,1, 3}, new Class[] {IntValue.class, StringValue.class, DoubleValue.class}).createComparator();

	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, multiComp, numberOfChannels);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	int[] hits = new int[numberOfChannels];
	for (int i = 0; i < numRecords; i++) {
		Record record = new Record(4);
		record.setField(0, new IntValue(i));
		record.setField(1, new StringValue("AB" + i + "CD" + i));
		record.setField(3, new DoubleValue(i * 3.141d));
		delegate.setInstance(record);

		int channel = selector.selectChannel(delegate);
		hits[channel]++;
	}

	int totalHitCount = 0;
	for (int hit : hits) {
		assertTrue(hit > 0);
		totalHitCount += hit;
	}
	assertTrue(totalHitCount == numRecords);
}
 

private int[] getSelectedChannelsHitCount(
		ShipStrategyType shipStrategyType,
		int numRecords,
		int numberOfChannels,
		Enum recordType) {
	final TypeComparator<Record> comparator = new RecordComparatorFactory(
		new int[] {0}, new Class[] {recordType == RecordType.INTEGER ? IntValue.class : StringValue.class}).createComparator();
	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(shipStrategyType, comparator, numberOfChannels);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	return getSelectedChannelsHitCount(selector, delegate, recordType, numRecords, numberOfChannels);
}
 

private void assertPartitionHashSelectedChannels(
		ChannelSelector selector,
		SerializationDelegate<Integer> serializationDelegate,
		int record,
		int numberOfChannels) {
	serializationDelegate.setInstance(record);
	int selectedChannel = selector.selectChannel(serializationDelegate);

	assertTrue(selectedChannel >= 0 && selectedChannel <= numberOfChannels - 1);
}
 

private void serializeRecord(long value, BufferBuilder bufferBuilder) throws IOException {
	RecordSerializer<SerializationDelegate<StreamElement>> serializer = new SpanningRecordSerializer<>();
	SerializationDelegate<StreamElement> serializationDelegate =
		new SerializationDelegate<>(
			new StreamElementSerializer<>(LongSerializer.INSTANCE));
	serializationDelegate.setInstance(new StreamRecord<>(value));
	serializer.serializeRecord(serializationDelegate);

	assertFalse(serializer.copyToBufferBuilder(bufferBuilder).isFullBuffer());
}
 

private void serializeRecord(long value, BufferBuilder bufferBuilder) throws IOException {
	RecordSerializer<SerializationDelegate<StreamElement>> serializer = new SpanningRecordSerializer<>();
	SerializationDelegate<StreamElement> serializationDelegate =
		new SerializationDelegate<>(
			new StreamElementSerializer<>(LongSerializer.INSTANCE));
	serializationDelegate.setInstance(new StreamRecord<>(value));
	serializer.serializeRecord(serializationDelegate);

	assertFalse(serializer.copyToBufferBuilder(bufferBuilder).isFullBuffer());
}
 

@Override
public int selectChannel(SerializationDelegate<StreamRecord<Tuple2<StreamRoute, Object>>> streamRecordSerializationDelegate) {
    Tuple2<StreamRoute, Object> value = streamRecordSerializationDelegate.getInstance().getValue();
    if (value.f0.getPartitionKey() == -1) {
        // random partition
        return random.nextInt(numberOfChannels);
    } else {
        return partitioner.partition(value.f0.getPartitionKey(), numberOfChannels);
    }
}
 

@Override
public int selectChannel(SerializationDelegate<StreamRecord<T>> record) {
	if (++nextChannelToSendTo >= numberOfChannels) {
		nextChannelToSendTo = 0;
	}
	return nextChannelToSendTo;
}
 

@Test
public void testMultiKeys() {
	final int numberOfChannels = 100;
	final int numRecords = 5000;
	final TypeComparator<Record> multiComp = new RecordComparatorFactory(
		new int[] {0,1, 3}, new Class[] {IntValue.class, StringValue.class, DoubleValue.class}).createComparator();

	final ChannelSelector<SerializationDelegate<Record>> selector = createChannelSelector(
		ShipStrategyType.PARTITION_HASH, multiComp, numberOfChannels);
	final SerializationDelegate<Record> delegate = new SerializationDelegate<>(new RecordSerializerFactory().getSerializer());

	int[] hits = new int[numberOfChannels];
	for (int i = 0; i < numRecords; i++) {
		Record record = new Record(4);
		record.setField(0, new IntValue(i));
		record.setField(1, new StringValue("AB" + i + "CD" + i));
		record.setField(3, new DoubleValue(i * 3.141d));
		delegate.setInstance(record);

		int channel = selector.selectChannel(delegate);
		hits[channel]++;
	}

	int totalHitCount = 0;
	for (int hit : hits) {
		assertTrue(hit > 0);
		totalHitCount += hit;
	}
	assertTrue(totalHitCount == numRecords);
}
 

@Override
public int selectChannel(SerializationDelegate<StreamRecord<T>> record) {
	K key;
	try {
		key = keySelector.getKey(record.getInstance().getValue());
	} catch (Exception e) {
		throw new RuntimeException("Could not extract key from " + record.getInstance().getValue(), e);
	}
	return KeyGroupRangeAssignment.assignKeyToParallelOperator(key, maxParallelism, numberOfChannels);
}