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

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

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

private static void copyMultiSegmentsToUnsafe(
		MemorySegment[] segments,
		int offset,
		Object target,
		int pointer,
		int numBytes) {
	int remainSize = numBytes;
	for (MemorySegment segment : segments) {
		int remain = segment.size() - offset;
		if (remain > 0) {
			int nCopy = Math.min(remain, remainSize);
			segment.copyToUnsafe(offset, target, numBytes - remainSize + pointer, 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 <T> void transform(
	K key,
	N namespace,
	T value,
	StateTransformationFunction<S, T> transformation) throws Exception {
	updateStat();
	MemorySegment keySegment = getKeySegment(key, namespace);
	int keyLen = keySegment.size();

	S oldState = getNode(keySegment, 0, keyLen);
	S newState = transformation.apply(oldState, value);
	byte[] stateBytes = skipListValueSerializer.serialize(newState);
	putValue(keySegment, 0, keyLen, stateBytes, false);
}
 

/**
 * Creates an new ChannelWriterOutputView that writes to the given channel and buffers data
 * in the given memory segments. If the given memory segments are null, the writer takes its buffers
 * directly from the return queue of the writer. Note that this variant locks if no buffers are contained
 * in the return queue.
 * 
 * @param writer The writer to write to.
 * @param memory The memory used to buffer data, or null, to utilize solely the return queue.
 * @param segmentSize The size of the memory segments.
 */
public ChannelWriterOutputView(BlockChannelWriter<MemorySegment> writer, List<MemorySegment> memory, int segmentSize) {
	super(segmentSize, HEADER_LENGTH);
	
	if (writer == null) {
		throw new NullPointerException();
	}
	
	this.writer = writer;
	
	if (memory == null) {
		this.numSegments = 0;
	} else {
		this.numSegments = memory.size();
		// load the segments into the queue
		final LinkedBlockingQueue<MemorySegment> queue = writer.getReturnQueue();
		for (int i = memory.size() - 1; i >= 0; --i) {
			final MemorySegment seg = memory.get(i);
			if (seg.size() != segmentSize) {
				throw new IllegalArgumentException("The supplied memory segments are not of the specified size.");
			}
			queue.add(seg);
		}
	}
	
	// get the first segment
	try {
		advance();
	}
	catch (IOException ioex) {
		throw new RuntimeException("BUG: IOException occurred while getting first block for ChannelWriterOutputView.", ioex);
	}
}
 

@Override
public Stream<K> getKeys(N namespace) {
	updateStat();
	MemorySegment namespaceSegment = skipListKeySerializer.serializeNamespaceToSegment(namespace);
	Iterator<Long> nodeIter = new NamespaceNodeIterator(namespaceSegment, 0, namespaceSegment.size());
	return StreamSupport.stream(Spliterators.spliteratorUnknownSize(nodeIter, 0), false)
		.map(this::helpGetKey);
}
 

public HeaderlessChannelWriterOutputView(
		BlockChannelWriter<MemorySegment> writer,
		List<MemorySegment> memory,
		int segmentSize) {
	super(segmentSize, 0);

	if (writer == null) {
		throw new NullPointerException();
	}

	this.writer = writer;

	Preconditions.checkNotNull(memory);

	// load the segments into the queue
	final LinkedBlockingQueue<MemorySegment> queue = writer.getReturnQueue();
	for (int i = memory.size() - 1; i >= 0; --i) {
		final MemorySegment seg = memory.get(i);
		if (seg.size() != segmentSize) {
			throw new IllegalArgumentException("This segment are not of the specified size.");
		}
		queue.add(seg);
	}

	// get the first segment
	try {
		advance();
	} catch (IOException ioex) {
		throw new RuntimeException(ioex);
	}
}
 

@Test
public void channelReadWriteOneSegment() {
	final int NUM_IOS = 1111;
	
	try {
		final FileIOChannel.ID channelID = this.ioManager.createChannel();
		final BlockChannelWriter<MemorySegment> writer = this.ioManager.createBlockChannelWriter(channelID);
		
		MemorySegment memSeg = MemorySegmentFactory.allocateUnpooledSegment(32 * 1024);
		
		for (int i = 0; i < NUM_IOS; i++) {
			for (int pos = 0; pos < memSeg.size(); pos += 4) {
				memSeg.putInt(pos, i);
			}
			
			writer.writeBlock(memSeg);
			memSeg = writer.getNextReturnedBlock();
		}
		
		writer.close();
		
		final BlockChannelReader<MemorySegment> reader = this.ioManager.createBlockChannelReader(channelID);
		for (int i = 0; i < NUM_IOS; i++) {
			reader.readBlock(memSeg);
			memSeg = reader.getNextReturnedBlock();
			
			for (int pos = 0; pos < memSeg.size(); pos += 4) {
				if (memSeg.getInt(pos) != i) {
					fail("Read memory segment contains invalid data.");
				}
			}
		}
		
		reader.closeAndDelete();
	}
	catch (Exception ex) {
		ex.printStackTrace();
		fail("Test encountered an exception: " + ex.getMessage());
	}
}
 

@Override
public S putAndGetOld(K key, N namespace, S state) {
	updateStat();
	MemorySegment keySegment = getKeySegment(key, namespace);
	int keyLen = keySegment.size();
	byte[] value = skipListValueSerializer.serialize(state);

	return putValue(keySegment, 0, keyLen, value, true);
}
 

private BuildSideBuffer(MemorySegment initialSegment, MemorySegmentSource memSource) {
	super(initialSegment, initialSegment.size(), 0);
	
	this.targetList = new ArrayList<MemorySegment>();
	this.memSource = memSource;
	this.sizeBits = MathUtils.log2strict(initialSegment.size());
}
 

/**
 * Creates an new ChannelWriterOutputView that writes to the given channel and buffers data
 * in the given memory segments. If the given memory segments are null, the writer takes its buffers
 * directly from the return queue of the writer. Note that this variant locks if no buffers are contained
 * in the return queue.
 * 
 * @param writer The writer to write to.
 * @param memory The memory used to buffer data, or null, to utilize solely the return queue.
 * @param segmentSize The size of the memory segments.
 */
public ChannelWriterOutputView(BlockChannelWriter<MemorySegment> writer, List<MemorySegment> memory, int segmentSize) {
	super(segmentSize, HEADER_LENGTH);
	
	if (writer == null) {
		throw new NullPointerException();
	}
	
	this.writer = writer;
	
	if (memory == null) {
		this.numSegments = 0;
	} else {
		this.numSegments = memory.size();
		// load the segments into the queue
		final LinkedBlockingQueue<MemorySegment> queue = writer.getReturnQueue();
		for (int i = memory.size() - 1; i >= 0; --i) {
			final MemorySegment seg = memory.get(i);
			if (seg.size() != segmentSize) {
				throw new IllegalArgumentException("The supplied memory segments are not of the specified size.");
			}
			queue.add(seg);
		}
	}
	
	// get the first segment
	try {
		advance();
	}
	catch (IOException ioex) {
		throw new RuntimeException("BUG: IOException occurred while getting first block for ChannelWriterOutputView.", ioex);
	}
}
 

private BuildSideBuffer(MemorySegment initialSegment, MemorySegmentSource memSource) {
	super(initialSegment, initialSegment.size(), 0);

	this.memSource = memSource;
	this.sizeBits = MathUtils.log2strict(initialSegment.size());
	this.targetList = new ArrayList<>();
	this.buildStageSegments = new ArrayList<>();
	this.buildStageSegments.add(initialSegment);
	this.buildStageInputView = new RandomAccessInputView(
			buildStageSegments, initialSegment.size());
}
 

/**
 * Constructs {@link BufferConsumer} instance with static content.
 */
public BufferConsumer(MemorySegment memorySegment, BufferRecycler recycler, Buffer.DataType dataType) {
	this(memorySegment, recycler, memorySegment.size(), dataType);
}
 

@Override
protected int getLimitForSegment(MemorySegment segment) {
	return numBlocksRemaining > 0 ? segment.size() : sizeOfLastBlock;
}
 

private void write(ChannelStateCheckpointWriter writer, InputChannelInfo channelInfo, byte[] data) throws Exception {
	MemorySegment segment = wrap(data);
	NetworkBuffer buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, Buffer.DataType.DATA_BUFFER, segment.size());
	writer.writeInput(channelInfo, buffer);
}
 

@Test
public void testBucketsNotFulfillSegment() throws Exception {
	final int NUM_KEYS = 10000;
	final int BUILD_VALS_PER_KEY = 3;
	final int PROBE_VALS_PER_KEY = 10;

	// create a build input that gives 30000 pairs with 3 values sharing the same key
	MutableObjectIterator<IntPair> buildInput = new UniformIntPairGenerator(NUM_KEYS, BUILD_VALS_PER_KEY, false);

	// create a probe input that gives 100000 pairs with 10 values sharing a key
	MutableObjectIterator<IntPair> probeInput = new UniformIntPairGenerator(NUM_KEYS, PROBE_VALS_PER_KEY, true);

	// allocate the memory for the HashTable
	List<MemorySegment> memSegments;
	try {
		// 33 is minimum number of pages required to perform hash join this inputs
		memSegments = this.memManager.allocatePages(MEM_OWNER, 33);
	}
	catch (MemoryAllocationException maex) {
		fail("Memory for the Join could not be provided.");
		return;
	}

	// For FLINK-2545, the buckets data may not fulfill it's buffer, for example, the buffer may contains 256 buckets,
	// while hash table only assign 250 bucket on it. The unused buffer bytes may contains arbitrary data, which may
	// influence hash table if forget to skip it. To mock this, put the invalid bucket data(partition=1, inMemory=true, count=-1)
	// at the end of buffer.
	for (MemorySegment segment : memSegments) {
		int newBucketOffset = segment.size() - 128;
		// initialize the header fields
		segment.put(newBucketOffset + 0, (byte)0);
		segment.put(newBucketOffset + 1, (byte)0);
		segment.putShort(newBucketOffset + 2, (short) -1);
		segment.putLong(newBucketOffset + 4, ~0x0L);
	}

	// ----------------------------------------------------------------------------------------

	final MutableHashTable<IntPair, IntPair> join = new MutableHashTable<IntPair, IntPair>(
		this.pairBuildSideAccesssor, this.pairProbeSideAccesssor,
		this.pairBuildSideComparator, this.pairProbeSideComparator, this.pairComparator,
		memSegments, ioManager);
	join.open(buildInput, probeInput);

	final IntPair recordReuse = new IntPair();
	int numRecordsInJoinResult = 0;

	while (join.nextRecord()) {
		MutableObjectIterator<IntPair> buildSide = join.getBuildSideIterator();
		while (buildSide.next(recordReuse) != null) {
			numRecordsInJoinResult++;
		}
	}
	Assert.assertEquals("Wrong number of records in join result.", NUM_KEYS * BUILD_VALS_PER_KEY * PROBE_VALS_PER_KEY, numRecordsInJoinResult);

	join.close();
	this.memManager.release(join.getFreedMemory());
}
 

@Override
protected int getLimitForSegment(MemorySegment segment) {
	return this.numBlocksRemaining > 0 ? segment.size() : this.lastBlockBytes;
}
 

private BuildSideBuffer(MemorySegment segment) {
	super(segment, segment.size(), 0);
	this.targetList = new ArrayList<>();
}
 

@Test
public void testBucketsNotFulfillSegment() throws Exception {
	final int numKeys = 10000;
	final int buildValsPerKey = 3;
	final int probeValsPerKey = 10;

	// create a build input that gives 30000 pairs with 3 values sharing the same key
	MutableObjectIterator<BinaryRowData> buildInput = new UniformBinaryRowGenerator(numKeys, buildValsPerKey, false);

	// create a probe input that gives 100000 pairs with 10 values sharing a key
	MutableObjectIterator<BinaryRowData> probeInput = new UniformBinaryRowGenerator(numKeys, probeValsPerKey, true);

	// allocate the memory for the HashTable
	MemoryManager memManager = MemoryManagerBuilder.newBuilder().setMemorySize(35 * PAGE_SIZE).build();
	// ----------------------------------------------------------------------------------------

	final BinaryHashTable table = new BinaryHashTable(conf, new Object(),
			this.buildSideSerializer, this.probeSideSerializer,
			new MyProjection(), new MyProjection(),
			memManager, 35 * PAGE_SIZE, ioManager, 24, 200000,
			true, HashJoinType.INNER, null, false, new boolean[]{true}, false);

	// For FLINK-2545, the buckets data may not fulfill it's buffer, for example, the buffer may contains 256 buckets,
	// while hash table only assign 250 bucket on it. The unused buffer bytes may contains arbitrary data, which may
	// influence hash table if forget to skip it. To mock this, put the invalid bucket data(partition=1, inMemory=true, count=-1)
	// at the end of buffer.
	int totalPages = table.getInternalPool().freePages();
	for (int i = 0; i < totalPages; i++) {
		MemorySegment segment = table.getInternalPool().nextSegment();
		int newBucketOffset = segment.size() - 128;
		// initialize the header fields
		segment.put(newBucketOffset, (byte) 0);
		segment.put(newBucketOffset + 1, (byte) 0);
		segment.putShort(newBucketOffset + 2, (short) -1);
		segment.putLong(newBucketOffset + 4, ~0x0L);
		table.returnPage(segment);
	}

	int numRecordsInJoinResult = join(table, buildInput, probeInput);

	Assert.assertEquals("Wrong number of records in join result.", numKeys * buildValsPerKey * probeValsPerKey, numRecordsInJoinResult);

	table.close();
	table.free();
}
 

/**
 * Creates a new buffer instance backed by the given <tt>memorySegment</tt> with <tt>0</tt> for
 * the <tt>readerIndex</tt> and <tt>size</tt> as <tt>writerIndex</tt>.
 *
 * @param memorySegment
 * 		backing memory segment (defines {@link #maxCapacity})
 * @param recycler
 * 		will be called to recycle this buffer once the reference count is <tt>0</tt>
 * @param isBuffer
 * 		whether this buffer represents a buffer (<tt>true</tt>) or an event (<tt>false</tt>)
 * @param size
 * 		current size of data in the buffer, i.e. the writer index to set
 */
public NetworkBuffer(MemorySegment memorySegment, BufferRecycler recycler, boolean isBuffer, int size) {
	super(memorySegment.size());
	this.memorySegment = checkNotNull(memorySegment);
	this.recycler = checkNotNull(recycler);
	this.isBuffer = isBuffer;
	this.currentSize = memorySegment.size();
	setSize(size);
}