Java Code Examples for org.apache.flink.core.memory.MemorySegmentFactory#wrap()

/**
 * This tests the internal capability of using partial {@link ByteBuffer}, making sure the internal methods
 * works when put/get state with a key stored at a none-zero offset of a ByteBuffer.
 */
@Test
public void testPutAndGetNodeWithNoneZeroOffset() {
	final int key = 10;
	final long namespace = 0L;
	final String valueString = "test";
	SkipListKeySerializer<Integer, Long> skipListKeySerializer =
		new SkipListKeySerializer<>(IntSerializer.INSTANCE, LongSerializer.INSTANCE);
	SkipListValueSerializer<String> skipListValueSerializer =
		new SkipListValueSerializer<>(StringSerializer.INSTANCE);
	byte[] keyBytes = skipListKeySerializer.serialize(key, namespace);
	byte[] constructedKeyBytes = new byte[keyBytes.length + 1];
	System.arraycopy(keyBytes, 0, constructedKeyBytes, 1, keyBytes.length);
	MemorySegment keySegment = MemorySegmentFactory.wrap(constructedKeyBytes);
	int keyLen = keyBytes.length;
	byte[] value = skipListValueSerializer.serialize(valueString);
	stateMap.putValue(keySegment, 1, keyLen, value, false);
	String state = stateMap.getNode(keySegment, 1, keyLen);
	assertThat(state, is(valueString));
}
 

/**
 * Test state map iterator illegal next call.
 */
@Test
public void testNamespaceNodeIteratorIllegalNextInvocation() {
	SkipListKeySerializer<Integer, Long> skipListKeySerializer =
		new SkipListKeySerializer<>(IntSerializer.INSTANCE, LongSerializer.INSTANCE);
	byte[] namespaceBytes = skipListKeySerializer.serializeNamespace(namespace);
	MemorySegment namespaceSegment = MemorySegmentFactory.wrap(namespaceBytes);
	Iterator<Long> iterator = stateMap.new NamespaceNodeIterator(namespaceSegment, 0, namespaceBytes.length);
	while (iterator.hasNext()) {
		iterator.next();
	}
	try {
		iterator.next();
		fail("Should have thrown NoSuchElementException.");
	} catch (NoSuchElementException e) {
		// expected
	}
}
 

static void deserializeFromPages(BinaryRowData reuse, ChannelReaderInputView inView,
		BinaryRowDataSerializer buildSideSerializer) throws IOException {
	if (shouldAdvance(
			inView.getCurrentSegmentLimit() - inView.getCurrentPositionInSegment(),
			buildSideSerializer)) {
		inView.advance();
	}
	MemorySegment segment = (reuse.getSegments() != null) ? reuse.getSegments()[0] : null;

	int length = toLength(inView.getCurrentSegment().getLong(inView.getCurrentPositionInSegment()));
	inView.skipBytesToRead(8);

	if (segment == null || segment.size() < length) {
		segment = MemorySegmentFactory.wrap(new byte[length]);
	}
	inView.readFully(segment.getHeapMemory(), 0, length);
	reuse.pointTo(segment, 0, length);
}
 

static void deserializeFromPages(BinaryRow reuse, ChannelReaderInputView inView,
		BinaryRowSerializer buildSideSerializer) throws IOException {
	if (shouldAdvance(
			inView.getCurrentSegmentLimit() - inView.getCurrentPositionInSegment(),
			buildSideSerializer)) {
		inView.advance();
	}
	MemorySegment segment = (reuse.getSegments() != null) ? reuse.getSegments()[0] : null;

	int length = toLength(inView.getCurrentSegment().getLong(inView.getCurrentPositionInSegment()));
	inView.skipBytesToRead(8);

	if (segment == null || segment.size() < length) {
		segment = MemorySegmentFactory.wrap(new byte[length]);
	}
	inView.readFully(segment.getHeapMemory(), 0, length);
	reuse.pointTo(segment, 0, length);
}
 

private void verifyGetKeySpace(KeySpace keySpace, MemorySegment memorySegment, int offset) {
	assertEquals(keySpace.level, SkipListUtils.getLevel(memorySegment, offset));
	assertEquals(keySpace.status, SkipListUtils.getNodeStatus(memorySegment, offset));
	assertEquals(keySpace.keyData.length, SkipListUtils.getKeyLen(memorySegment, offset));
	assertEquals(keySpace.valuePointer, SkipListUtils.getValuePointer(memorySegment, offset));
	assertEquals(keySpace.nextKeyPointer, SkipListUtils.getNextKeyPointer(memorySegment, offset));
	for (int i = 1; i <= keySpace.nextIndexNodes.length; i++) {
		assertEquals(keySpace.nextIndexNodes[i - 1], SkipListUtils.getNextIndexNode(memorySegment, offset, i));
	}
	for (int i = 1; i <= keySpace.prevIndexNodes.length; i++) {
		assertEquals(keySpace.prevIndexNodes[i - 1],
			SkipListUtils.getPrevIndexNode(memorySegment, offset, keySpace.level, i));
	}
	int keyDataOffset = SkipListUtils.getKeyDataOffset(keySpace.level);
	MemorySegment keyDataSegment = MemorySegmentFactory.wrap(keySpace.keyData);
	assertEquals(0, memorySegment.compare(keyDataSegment, offset + keyDataOffset, 0, keySpace.keyData.length));
}
 

public CompressedHeaderlessChannelReaderInputView(
		FileIOChannel.ID id,
		IOManager ioManager,
		BlockCompressionFactory compressionCodecFactory,
		int compressionBlockSize,
		int numBlocks) throws IOException {
	super(0);
	this.numBlocksRemaining = numBlocks;
	this.reader = ioManager.createBufferFileReader(id, this);
	uncompressedBuffer = MemorySegmentFactory.wrap(new byte[compressionBlockSize]);
	decompressor = compressionCodecFactory.getDecompressor();
	cause = new AtomicReference<>();

	BlockCompressor compressor = compressionCodecFactory.getCompressor();
	for (int i = 0; i < 2; i++) {
		MemorySegment segment = MemorySegmentFactory.wrap(new byte[compressor.getMaxCompressedSize(
				compressionBlockSize)]);
		reader.readInto(new NetworkBuffer(segment, this));
	}
}
 

public BinaryRowWriter(BinaryRow row, int initialSize) {
	this.nullBitsSizeInBytes = BinaryRow.calculateBitSetWidthInBytes(row.getArity());
	this.fixedSize = row.getFixedLengthPartSize();
	this.cursor = fixedSize;

	this.segment = MemorySegmentFactory.wrap(new byte[fixedSize + initialSize]);
	this.row = row;
	this.row.pointTo(segment, 0, segment.size());
}
 

@Override
protected BinarySection materialize(TypeSerializer<T> serializer) {
	try {
		byte[] bytes = InstantiationUtil.serializeToByteArray(serializer, javaObject);
		return new BinarySection(new MemorySegment[] {MemorySegmentFactory.wrap(bytes)}, 0, bytes.length);
	} catch (IOException e) {
		throw new RuntimeException(e);
	}
}
 

@Test
public void testBoundaryEquals() {
	BinaryRowData row24 = DataFormatTestUtil.get24BytesBinaryRow();
	BinaryRowData row160 = DataFormatTestUtil.get160BytesBinaryRow();
	BinaryRowData varRow160 = DataFormatTestUtil.getMultiSeg160BytesBinaryRow(row160);
	BinaryRowData varRow160InOne = DataFormatTestUtil.getMultiSeg160BytesInOneSegRow(row160);

	assertEquals(row160, varRow160InOne);
	assertEquals(varRow160, varRow160InOne);
	assertEquals(row160, varRow160);
	assertEquals(varRow160InOne, varRow160);

	assertNotEquals(row24, row160);
	assertNotEquals(row24, varRow160);
	assertNotEquals(row24, varRow160InOne);

	assertTrue(BinarySegmentUtils.equals(row24.getSegments(), 0, row160.getSegments(), 0, 0));
	assertTrue(BinarySegmentUtils.equals(row24.getSegments(), 0, varRow160.getSegments(), 0, 0));

	// test var segs
	MemorySegment[] segments1 = new MemorySegment[2];
	segments1[0] = MemorySegmentFactory.wrap(new byte[32]);
	segments1[1] = MemorySegmentFactory.wrap(new byte[32]);
	MemorySegment[] segments2 = new MemorySegment[3];
	segments2[0] = MemorySegmentFactory.wrap(new byte[16]);
	segments2[1] = MemorySegmentFactory.wrap(new byte[16]);
	segments2[2] = MemorySegmentFactory.wrap(new byte[16]);

	segments1[0].put(9, (byte) 1);
	assertFalse(BinarySegmentUtils.equals(segments1, 0, segments2, 14, 14));
	segments2[1].put(7, (byte) 1);
	assertTrue(BinarySegmentUtils.equals(segments1, 0, segments2, 14, 14));
	assertTrue(BinarySegmentUtils.equals(segments1, 2, segments2, 16, 14));
	assertTrue(BinarySegmentUtils.equals(segments1, 2, segments2, 16, 16));

	segments2[2].put(7, (byte) 1);
	assertTrue(BinarySegmentUtils.equals(segments1, 2, segments2, 32, 14));
}
 

@Test
public void testCopy() {
	// test copy the content of the latter Seg
	MemorySegment[] segments = new MemorySegment[2];
	segments[0] = MemorySegmentFactory.wrap(new byte[]{0, 2, 5});
	segments[1] = MemorySegmentFactory.wrap(new byte[]{6, 12, 15});

	byte[] bytes = BinarySegmentUtils.copyToBytes(segments, 4, 2);
	Assert.assertArrayEquals(new byte[] {12, 15}, bytes);
}
 

@Override
public BinaryGeneric<T> deserialize(DataInputView source) throws IOException {
	int length = source.readInt();
	byte[] bytes = new byte[length];
	source.readFully(bytes);
	return new BinaryGeneric<>(
			new MemorySegment[] {MemorySegmentFactory.wrap(bytes)},
			0, bytes.length, serializer);
}
 

@Override
public void materialize() {
	byte[] bytes = StringUtf8Utils.encodeUTF8(javaObject);
	segments = new MemorySegment[] {MemorySegmentFactory.wrap(bytes)};
	offset = 0;
	sizeInBytes = bytes.length;
}
 

@Override
public RawValueData<T> copy(RawValueData<T> from) {
	BinaryRawValueData<T> rawValue = (BinaryRawValueData<T>) from;
	rawValue.ensureMaterialized(serializer);
	byte[] bytes = BinarySegmentUtils.copyToBytes(
		rawValue.getSegments(),
		rawValue.getOffset(),
		rawValue.getSizeInBytes());
	T newJavaObject = rawValue.getJavaObject() == null ? null : serializer.copy(rawValue.getJavaObject());
	return new BinaryRawValueData<>(
		new MemorySegment[]{MemorySegmentFactory.wrap(bytes)},
		0,
		bytes.length,
		newJavaObject);
}
 

@Test
public void testGenericObject() throws Exception {

	GenericTypeInfo<MyObj> info = new GenericTypeInfo<>(MyObj.class);
	TypeSerializer<MyObj> genericSerializer = info.createSerializer(new ExecutionConfig());

	BinaryRow row = new BinaryRow(4);
	BinaryRowWriter writer = new BinaryRowWriter(row);
	writer.writeInt(0, 0);

	BinaryGeneric<MyObj> myObj1 = new BinaryGeneric<>(new MyObj(0, 1), genericSerializer);
	writer.writeGeneric(1, myObj1);
	BinaryGeneric<MyObj> myObj2 = new BinaryGeneric<>(new MyObj(123, 5.0), genericSerializer);
	myObj2.ensureMaterialized();
	writer.writeGeneric(2, myObj2);
	BinaryGeneric<MyObj> myObj3 = new BinaryGeneric<>(new MyObj(1, 1), genericSerializer);
	writer.writeGeneric(3, myObj3);
	writer.complete();

	assertTestGenericObjectRow(row, genericSerializer);

	// getBytes from var-length memorySegments.
	BinaryRowSerializer serializer = new BinaryRowSerializer(4);
	MemorySegment[] memorySegments = new MemorySegment[3];
	ArrayList<MemorySegment> memorySegmentList = new ArrayList<>();
	for (int i = 0; i < 3; i++) {
		memorySegments[i] = MemorySegmentFactory.wrap(new byte[64]);
		memorySegmentList.add(memorySegments[i]);
	}
	RandomAccessOutputView out = new RandomAccessOutputView(memorySegments, 64);
	serializer.serializeToPages(row, out);

	BinaryRow mapRow = serializer.mapFromPages(new RandomAccessInputView(memorySegmentList, 64));
	assertTestGenericObjectRow(mapRow, genericSerializer);
}
 

@Test
public void testHashAndCopy() throws IOException {
	MemorySegment[] segments = new MemorySegment[3];
	for (int i = 0; i < 3; i++) {
		segments[i] = MemorySegmentFactory.wrap(new byte[64]);
	}
	RandomAccessOutputView out = new RandomAccessOutputView(segments, 64);
	BinaryRowSerializer serializer = new BinaryRowSerializer(2);

	BinaryRow row = new BinaryRow(2);
	BinaryRowWriter writer = new BinaryRowWriter(row);
	writer.writeString(0, BinaryString.fromString("hahahahahahahahahahahahahahahahahahahhahahahahahahahahah"));
	writer.writeString(1, BinaryString.fromString("hahahahahahahahahahahahahahahahahahahhahahahahahahahahaa"));
	writer.complete();
	serializer.serializeToPages(row, out);

	ArrayList<MemorySegment> segmentList = new ArrayList<>(Arrays.asList(segments));
	RandomAccessInputView input = new RandomAccessInputView(segmentList, 64, 64);

	BinaryRow mapRow = serializer.mapFromPages(input);
	assertEquals(row, mapRow);
	assertEquals(row.getString(0), mapRow.getString(0));
	assertEquals(row.getString(1), mapRow.getString(1));
	assertNotEquals(row.getString(0), mapRow.getString(1));

	// test if the hash code before and after serialization are the same
	assertEquals(row.hashCode(), mapRow.hashCode());
	assertEquals(row.getString(0).hashCode(), mapRow.getString(0).hashCode());
	assertEquals(row.getString(1).hashCode(), mapRow.getString(1).hashCode());

	// test if the copy method produce a row with the same contents
	assertEquals(row.copy(), mapRow.copy());
	assertEquals(row.getString(0).copy(), mapRow.getString(0).copy());
	assertEquals(row.getString(1).copy(), mapRow.getString(1).copy());
}
 

@Test
public void testSerializeVariousSize() throws IOException {
	// in this test, we are going to start serializing from the i-th byte (i in 0...`segSize`)
	// and the size of the row we're going to serialize is j bytes
	// (j in `rowFixLength` to the maximum length we can write)

	int segSize = 64;
	int segTotalNumber = 3;

	BinaryRowData row = new BinaryRowData(1);
	BinaryRowWriter writer = new BinaryRowWriter(row);
	Random random = new Random();
	byte[] bytes = new byte[1024];
	random.nextBytes(bytes);
	writer.writeBinary(0, bytes);
	writer.complete();

	MemorySegment[] memorySegments = new MemorySegment[segTotalNumber];
	Map<MemorySegment, Integer> msIndex = new HashMap<>();
	for (int i = 0; i < segTotalNumber; i++) {
		memorySegments[i] = MemorySegmentFactory.wrap(new byte[segSize]);
		msIndex.put(memorySegments[i], i);
	}

	BinaryRowDataSerializer serializer = new BinaryRowDataSerializer(1);

	int rowSizeInt = 4;
	// note that as there is only one field in the row, the fixed-length part is 16 bytes (header + 1 field)
	int rowFixLength = 16;
	for (int i = 0; i < segSize; i++) {
		// this is the maximum row size we can serialize
		// if we are going to serialize from the i-th byte of the input view
		int maxRowSize = (segSize * segTotalNumber) - i - rowSizeInt;
		if (segSize - i < rowFixLength + rowSizeInt) {
			// oops, we can't write the whole fixed-length part in the first segment
			// because the remaining space is too small, so we have to start serializing from the second segment.
			// when serializing, we need to first write the length of the row,
			// then write the fixed-length part of the row.
			maxRowSize -= segSize - i;
		}
		for (int j = rowFixLength; j < maxRowSize; j++) {
			// ok, now we're going to serialize a row of j bytes
			testSerialize(row, memorySegments, msIndex, serializer, i, j);
		}
	}
}
 

public static Buffer toBuffer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	final Buffer buffer = new NetworkBuffer(data, FreeingBufferRecycler.INSTANCE, Buffer.DataType.getDataType(event));
	buffer.setSize(serializedEvent.remaining());

	return buffer;
}
 

public static BufferConsumer toBufferConsumer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	return new BufferConsumer(data, FreeingBufferRecycler.INSTANCE, Buffer.DataType.getDataType(event));
}
 

public static Buffer toBuffer(AbstractEvent event) throws IOException {
	final ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(event);

	MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());

	final Buffer buffer = new NetworkBuffer(data, FreeingBufferRecycler.INSTANCE, false);
	buffer.setSize(serializedEvent.remaining());

	return buffer;
}
 

/**
 * Puts the value data into value space.
 *
 * @param memorySegment memory segment for value space.
 * @param offset offset of value space in memory segment.
 * @param value value data.
 */
public static void putValueData(MemorySegment memorySegment, int offset, byte[] value) {
	MemorySegment valueSegment = MemorySegmentFactory.wrap(value);
	valueSegment.copyTo(0, memorySegment, offset + getValueMetaLen(), value.length);
}