Java Code Examples for org.apache.flink.core.memory.MemorySegment#put()

@Override
public void putNormalizedKey(KV<K, V> record, MemorySegment target, int offset, int numBytes) {
	buffer1.reset();
	try {
		keyCoder.encode(record.getKey(), buffer1, Coder.Context.NESTED);
	} catch (IOException e) {
		throw new RuntimeException("Could not serializer " + record + " using coder " + coder + ": " + e);
	}
	final byte[] data = buffer1.getBuffer();
	final int limit = offset + numBytes;

	int numBytesPut = Math.min(numBytes, buffer1.size());

	target.put(offset, data, 0, numBytesPut);

	offset += numBytesPut;

	while (offset < limit) {
		target.put(offset++, (byte) 0);
	}
}
 

private static void copyMultiSegmentsFromBytes(
		MemorySegment[] segments, int offset, byte[] bytes, int bytesOffset, int numBytes) {
	int remainSize = numBytes;
	for (MemorySegment segment : segments) {
		int remain = segment.size() - offset;
		if (remain > 0) {
			int nCopy = Math.min(remain, remainSize);
			segment.put(offset, bytes, numBytes - remainSize + bytesOffset, nCopy);
			remainSize -= nCopy;
			// next new segment.
			offset = 0;
			if (remainSize == 0) {
				return;
			}
		} else {
			// remain is negative, let's advance to next segment
			// now the offset = offset - segmentSize (-remain)
			offset = -remain;
		}
	}
}
 

@Override
public void putNormalizedKey(Character value, MemorySegment target, int offset, int numBytes) {
	// note that the char is an unsigned data type in java and consequently needs
	// no code that transforms the signed representation to an offset representation
	// that is equivalent to unsigned, when compared byte by byte
	if (numBytes == 2) {
		// default case, full normalized key
		target.put(offset,     (byte) ((value >>> 8) & 0xff));
		target.put(offset + 1, (byte) ((value      ) & 0xff));
	}
	else if (numBytes <= 0) {
	}
	else if (numBytes == 1) {
		target.put(offset,     (byte) ((value >>> 8) & 0xff));
	}
	else {
		target.put(offset,     (byte) ((value >>> 8) & 0xff));
		target.put(offset + 1, (byte) ((value      ) & 0xff));
		for (int i = 2; i < numBytes; i++) {
			target.put(offset + i, (byte) 0);
		}
	}
}
 

private static void setLongSlowly(
		MemorySegment[] segments, int segSize, int segNum, int segOffset, long value) {
	MemorySegment segment = segments[segNum];
	for (int i = 0; i < 8; i++) {
		if (segOffset == segSize) {
			segment = segments[++segNum];
			segOffset = 0;
		}
		long unsignedByte;
		if (LITTLE_ENDIAN) {
			unsignedByte = value >> (i * 8);
		} else {
			unsignedByte = value >> ((7 - i) * 8);
		}
		segment.put(segOffset, (byte) unsignedByte);
		segOffset++;
	}
}
 

private static void setIntSlowly(
		MemorySegment[] segments, int segSize, int segNum, int segOffset, int value) {
	MemorySegment segment = segments[segNum];
	for (int i = 0; i < 4; i++) {
		if (segOffset == segSize) {
			segment = segments[++segNum];
			segOffset = 0;
		}
		int unsignedByte;
		if (LITTLE_ENDIAN) {
			unsignedByte = value >> (i * 8);
		} else {
			unsignedByte = value >> ((3 - i) * 8);
		}
		segment.put(segOffset, (byte) unsignedByte);
		segOffset++;
	}
}
 

public static void putUnsignedIntegerNormalizedKey(int value, MemorySegment target, int offset, int numBytes) {
	if (numBytes == 4) {
		// default case, full normalized key
		target.putIntBigEndian(offset, value);
	} else if (numBytes > 0) {
		if (numBytes < 4) {
			for (int i = 0; numBytes > 0; numBytes--, i++) {
				target.put(offset + i, (byte) (value >>> ((3 - i) << 3)));
			}
		} else {
			target.putIntBigEndian(offset, value);
			for (int i = 4; i < numBytes; i++) {
				target.put(offset + i, (byte) 0);
			}
		}
	}
}
 

private MemorySegment copyDataBuffer() throws IOException {
	int leftOverSize = leftOverLimit - leftOverStart;
	int unconsumedSize = LENGTH_BYTES + accumulatedRecordBytes + leftOverSize;
	DataOutputSerializer serializer = new DataOutputSerializer(unconsumedSize);
	serializer.writeInt(recordLength);
	serializer.write(buffer, 0, accumulatedRecordBytes);
	if (leftOverData != null) {
		serializer.write(leftOverData, leftOverStart, leftOverSize);
	}
	MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(unconsumedSize);
	segment.put(0, serializer.getSharedBuffer(), 0, segment.size());
	return segment;
}
 

protected void initTable(int numBuckets, byte numPartitions) {
	final int bucketsPerSegment = this.bucketsPerSegmentMask + 1;
	final int numSegs = (numBuckets >>> this.bucketsPerSegmentBits) + ( (numBuckets & this.bucketsPerSegmentMask) == 0 ? 0 : 1);
	final MemorySegment[] table = new MemorySegment[numSegs];
	
	ensureNumBuffersReturned(numSegs);
	
	// go over all segments that are part of the table
	for (int i = 0, bucket = 0; i < numSegs && bucket < numBuckets; i++) {
		final MemorySegment seg = getNextBuffer();
		
		// go over all buckets in the segment
		for (int k = 0; k < bucketsPerSegment && bucket < numBuckets; k++, bucket++) {
			final int bucketOffset = k * HASH_BUCKET_SIZE;	
			
			// compute the partition that the bucket corresponds to
			final byte partition = assignPartition(bucket, numPartitions);
			
			// initialize the header fields
			seg.put(bucketOffset + HEADER_PARTITION_OFFSET, partition);
			seg.put(bucketOffset + HEADER_STATUS_OFFSET, BUCKET_STATUS_IN_MEMORY);
			seg.putShort(bucketOffset + HEADER_COUNT_OFFSET, (short) 0);
			seg.putLong(bucketOffset + HEADER_FORWARD_OFFSET, BUCKET_FORWARD_POINTER_NOT_SET);
			seg.putShort(bucketOffset + HEADER_PROBED_FLAGS_OFFSET, (short) 0);
		}
		
		table[i] = seg;
	}
	this.buckets = table;
	this.numBuckets = numBuckets;
	
	if (useBloomFilters) {
		initBloomFilter(numBuckets);
	}
}
 

@Override
public void putNormalizedKey(String record, MemorySegment target, int offset, int len) {
	final int limit = offset + len;
	final int end = record.length();
	int pos = 0;
	
	while (pos < end && offset < limit) {
		char c = record.charAt(pos++);
		if (c < HIGH_BIT) {
			target.put(offset++, (byte) c);
		}
		else if (c < HIGH_BIT2) {
			target.put(offset++, (byte) ((c >>> 7) | HIGH_BIT));
			if (offset < limit) {
				target.put(offset++, (byte) c);
			}
		}
		else {
			target.put(offset++, (byte) ((c >>> 10) | HIGH_BIT2_MASK));
			if (offset < limit) {
				target.put(offset++, (byte) (c >>> 2));
			}
			if (offset < limit) {
				target.put(offset++, (byte) c);
			}
		}
	}
	while (offset < limit) {
		target.put(offset++, (byte) 0);
	}
}
 

public static BufferOrEvent generateRandomBuffer(int size, int channelIndex) {
	MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	for (int i = 0; i < size; i++) {
		seg.put(i, (byte) i);
	}

	Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);
	return new BufferOrEvent(buf, channelIndex);
}
 

public static void putBooleanNormalizedKey(boolean value, MemorySegment target, int offset, int numBytes) {
	if (numBytes > 0) {
		target.put(offset, (byte) (value ? 1 : 0));

		for (offset = offset + 1; numBytes > 1; numBytes--) {
			target.put(offset++, (byte) 0);
		}
	}
}
 

/**
 * unset bit from segments.
 *
 * @param segments target segments.
 * @param baseOffset bits base offset.
 * @param index bit index from base offset.
 */
public static void bitUnSet(MemorySegment[] segments, int baseOffset, int index) {
	if (segments.length == 1) {
		MemorySegment segment = segments[0];
		int offset = baseOffset + byteIndex(index);
		byte current = segment.get(offset);
		current &= ~(1 << (index & BIT_BYTE_INDEX_MASK));
		segment.put(offset, current);
	} else {
		bitUnSetMultiSegments(segments, baseOffset, index);
	}
}
 

private static BufferOrEvent createBuffer(int channel, int size) {
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	memory.put(0, bytes);

	Buffer buf = new NetworkBuffer(memory, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);

	// retain an additional time so it does not get disposed after being read by the input gate
	buf.retainBuffer();

	return new BufferOrEvent(buf, channel);
}
 

/**
 * unset bit from segments.
 *
 * @param segments target segments.
 * @param baseOffset bits base offset.
 * @param index bit index from base offset.
 */
public static void bitUnSet(MemorySegment[] segments, int baseOffset, int index) {
	if (segments.length == 1) {
		MemorySegment segment = segments[0];
		int offset = baseOffset + byteIndex(index);
		byte current = segment.get(offset);
		current &= ~(1 << (index & BIT_BYTE_INDEX_MASK));
		segment.put(offset, current);
	} else {
		bitUnSetMultiSegments(segments, baseOffset, index);
	}
}
 

public static void minNormalizedKey(MemorySegment target, int offset, int numBytes) {
	//write min value.
	for (int i = 0; i < numBytes; i++) {
		target.put(offset + i, (byte) 0);
	}
}
 

@Override
public void copyNormalizedKey(MemorySegment target, int offset, int len) {
	for (int i = offset; i < offset + len; i++) {
		target.put(i, (byte) 0);
	}
}
 

@Override
public void copyNormalizedKey(MemorySegment target, int offset, int len) {
	for (int i = offset; i < offset + len; i++) {
		target.put(i, (byte) 0);
	}
}
 

@Override
public BufferOrEvent getNext() throws IOException {
	if (buffer.remaining() < HEADER_LENGTH) {
		buffer.compact();

		while (buffer.position() < HEADER_LENGTH) {
			if (fileChannel.read(buffer) == -1) {
				if (buffer.position() == 0) {
					// no trailing data
					return null;
				} else {
					throw new IOException("Found trailing incomplete buffer or event");
				}
			}
		}

		buffer.flip();
	}

	final int channel = buffer.getInt();
	final int length = buffer.getInt();
	final boolean isBuffer = buffer.get() == 0;

	if (isBuffer) {
		// deserialize buffer
		if (length > pageSize) {
			throw new IOException(String.format(
					"Spilled buffer (%d bytes) is larger than page size of (%d bytes)", length, pageSize));
		}

		MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(pageSize);

		int segPos = 0;
		int bytesRemaining = length;

		while (true) {
			int toCopy = Math.min(buffer.remaining(), bytesRemaining);
			if (toCopy > 0) {
				seg.put(segPos, buffer, toCopy);
				segPos += toCopy;
				bytesRemaining -= toCopy;
			}

			if (bytesRemaining == 0) {
				break;
			}
			else {
				buffer.clear();
				if (fileChannel.read(buffer) == -1) {
					throw new IOException("Found trailing incomplete buffer");
				}
				buffer.flip();
			}
		}

		Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
		buf.setSize(length);

		return new BufferOrEvent(buf, channel, true);
	}
	else {
		// deserialize event
		if (length > buffer.capacity() - HEADER_LENGTH) {
			throw new IOException("Event is too large");
		}

		if (buffer.remaining() < length) {
			buffer.compact();

			while (buffer.position() < length) {
				if (fileChannel.read(buffer) == -1) {
					throw new IOException("Found trailing incomplete event");
				}
			}

			buffer.flip();
		}

		int oldLimit = buffer.limit();
		buffer.limit(buffer.position() + length);
		AbstractEvent evt = EventSerializer.fromSerializedEvent(buffer, getClass().getClassLoader());
		buffer.limit(oldLimit);

		return new BufferOrEvent(evt, channel, true, length);
	}
}
 

/**
 * unset bit.
 *
 * @param segment target segment.
 * @param baseOffset bits base offset.
 * @param index bit index from base offset.
 */
public static void bitUnSet(MemorySegment segment, int baseOffset, int index) {
	int offset = baseOffset + byteIndex(index);
	byte current = segment.get(offset);
	current &= ~(1 << (index & BIT_BYTE_INDEX_MASK));
	segment.put(offset, current);
}
 

/**
 * set bit.
 *
 * @param segment target segment.
 * @param baseOffset bits base offset.
 * @param index bit index from base offset.
 */
public static void bitSet(MemorySegment segment, int baseOffset, int index) {
	int offset = baseOffset + byteIndex(index);
	byte current = segment.get(offset);
	current |= (1 << (index & BIT_BYTE_INDEX_MASK));
	segment.put(offset, current);
}