org.apache.flink.util.CloseableIterator Java Examples

@Override
public void spillInflightBuffers(long checkpointId, ChannelStateWriter channelStateWriter) throws IOException {
	synchronized (receivedBuffers) {
		checkState(checkpointId > lastRequestedCheckpointId, "Need to request the next checkpointId");

		final List<Buffer> inflightBuffers = new ArrayList<>(receivedBuffers.size());
		for (Buffer buffer : receivedBuffers) {
			CheckpointBarrier checkpointBarrier = parseCheckpointBarrierOrNull(buffer);
			if (checkpointBarrier != null && checkpointBarrier.getId() >= checkpointId) {
				break;
			}
			if (buffer.isBuffer()) {
				inflightBuffers.add(buffer.retainBuffer());
			}
		}

		lastRequestedCheckpointId = checkpointId;

		channelStateWriter.addInputData(
			checkpointId,
			channelInfo,
			ChannelStateWriter.SEQUENCE_NUMBER_UNKNOWN,
			CloseableIterator.fromList(inflightBuffers, Buffer::recycleBuffer));
	}
}
 

static ChannelStateWriteRequest buildWriteRequest(
		long checkpointId,
		String name,
		CloseableIterator<Buffer> iterator,
		BiConsumerWithException<ChannelStateCheckpointWriter, Buffer, Exception> bufferConsumer) {
	return new CheckpointInProgressRequest(
		name,
		checkpointId,
		writer -> {
			while (iterator.hasNext()) {
				Buffer buffer = iterator.next();
				try {
					checkArgument(buffer.isBuffer());
				} catch (Exception e) {
					buffer.recycleBuffer();
					throw e;
				}
				bufferConsumer.accept(writer, buffer);
			}
		},
		throwable -> iterator.close(),
		false);
}
 

@Test
public void testIterator() throws Exception {
	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	// test empty iterator
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		Assert.assertFalse(iterator.hasNext());
		try {
			iterator.next();
			Assert.fail();
		} catch (NoSuchElementException ignore) {
		}
	}

	// iterate some data
	final int testSize = 10;
	HashSet<TestElement> checkSet = new HashSet<>(testSize);
	insertRandomElements(priorityQueue, checkSet, testSize);
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			Assert.assertTrue(checkSet.remove(iterator.next()));
		}
		Assert.assertTrue(checkSet.isEmpty());
	}
}
 

private byte[] toByteArray(CloseableIterator<Buffer> unconsumed) {
	final List<Buffer> buffers = new ArrayList<>();
	try {
		unconsumed.forEachRemaining(buffers::add);
		byte[] result = new byte[buffers.stream().mapToInt(Buffer::readableBytes).sum()];
		int offset = 0;
		for (Buffer buffer : buffers) {
			int len = buffer.readableBytes();
			buffer.getNioBuffer(0, len).get(result, offset, len);
			offset += len;
		}
		return result;
	} finally {
		buffers.forEach(Buffer::recycleBuffer);
	}
}
 

@Test
public void testLargeUnconsumedSegment() throws Exception {
	int recordLen = 100;
	int firstChunk = (int) (recordLen * .9);
	int spillingThreshold = (int) (firstChunk * .9);

	byte[] record1 = recordBytes(recordLen);
	byte[] record2 = recordBytes(recordLen * 2);

	SpanningWrapper spanningWrapper = new SpanningWrapper(new String[]{folder.newFolder().getAbsolutePath()}, spillingThreshold, recordLen);
	spanningWrapper.transferFrom(wrapNonSpanning(record1, firstChunk), recordLen);
	spanningWrapper.addNextChunkFromMemorySegment(wrap(record1), firstChunk, recordLen - firstChunk + LENGTH_BYTES);
	spanningWrapper.addNextChunkFromMemorySegment(wrap(record2), 0, record2.length);

	CloseableIterator<Buffer> unconsumedSegment = spanningWrapper.getUnconsumedSegment();

	spanningWrapper.getInputView().readFully(new byte[recordLen], 0, recordLen); // read out from file
	spanningWrapper.transferLeftOverTo(new NonSpanningWrapper()); // clear any leftover
	spanningWrapper.transferFrom(wrapNonSpanning(recordBytes(recordLen), recordLen), recordLen); // overwrite with new data

	assertArrayEquals(concat(record1, record2), toByteArray(unconsumedSegment));
}
 

@Test
public void testIterator() throws Exception {
	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	// test empty iterator
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		Assert.assertFalse(iterator.hasNext());
		try {
			iterator.next();
			Assert.fail();
		} catch (NoSuchElementException ignore) {
		}
	}

	// iterate some data
	final int testSize = 10;
	HashSet<TestElement> checkSet = new HashSet<>(testSize);
	insertRandomElements(priorityQueue, checkSet, testSize);
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			Assert.assertTrue(checkSet.remove(iterator.next()));
		}
		Assert.assertTrue(checkSet.isEmpty());
	}
}
 

@Test
public void testIterator() throws Exception {
	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	// test empty iterator
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		Assert.assertFalse(iterator.hasNext());
		try {
			iterator.next();
			Assert.fail();
		} catch (NoSuchElementException ignore) {
		}
	}

	// iterate some data
	final int testSize = 10;
	HashSet<TestElement> checkSet = new HashSet<>(testSize);
	insertRandomElements(priorityQueue, checkSet, testSize);
	try (CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			Assert.assertTrue(checkSet.remove(iterator.next()));
		}
		Assert.assertTrue(checkSet.isEmpty());
	}
}
 

@SuppressWarnings("unchecked")
private CloseableIterator<Buffer> createSpilledDataIterator() throws IOException {
	if (spillingChannel != null && spillingChannel.isOpen()) {
		spillingChannel.force(false);
	}
	return CloseableIterator.flatten(
		toSingleBufferIterator(wrapInt(recordLength)),
		new FileBasedBufferIterator(spillFile, min(accumulatedRecordBytes, recordLength), fileBufferSize),
		leftOverData == null ? empty() : toSingleBufferIterator(wrapCopy(leftOverData.getArray(), leftOverStart, leftOverLimit))
	);
}
 

private void foreachTimer(BiConsumerWithException<N, Long, Exception> consumer, KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) throws Exception {
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			keyContext.setCurrentKey(timer.getKey());
			consumer.accept(timer.getNamespace(), timer.getTimestamp());
		}
	}
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}
 

public SelectResultProvider getSelectResultProvider() {
	return new SelectResultProvider() {

		@Override
		public void setJobClient(JobClient jobClient) {
			iterator.setJobClient(jobClient);
		}

		@Override
		public CloseableIterator<Row> getResultIterator() {
			return new RowIteratorWrapper(iterator);
		}
	};
}
 

private TableResultImpl(
		@Nullable JobClient jobClient,
		TableSchema tableSchema,
		ResultKind resultKind,
		CloseableIterator<Row> data,
		PrintStyle printStyle) {
	this.jobClient = jobClient;
	this.tableSchema = Preconditions.checkNotNull(tableSchema, "tableSchema should not be null");
	this.resultKind = Preconditions.checkNotNull(resultKind, "resultKind should not be null");
	this.data = Preconditions.checkNotNull(data, "data should not be null");
	this.printStyle = Preconditions.checkNotNull(printStyle, "printStyle should not be null");
}
 

CloseableIterator<Buffer> getUnconsumedSegment() {
	if (!hasRemaining()) {
		return CloseableIterator.empty();
	}
	MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(remaining());
	this.segment.copyTo(position, segment, 0, remaining());
	return singleBufferIterator(segment);
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}
 

CloseableIterator<Buffer> getUnconsumedSegment() throws IOException {
	if (isReadingLength()) {
		return singleBufferIterator(wrapCopy(lengthBuffer.array(), 0, lengthBuffer.position()));
	} else if (isAboveSpillingThreshold()) {
		return createSpilledDataIterator();
	} else if (recordLength == -1) {
		return empty(); // no remaining partial length or data
	} else {
		return singleBufferIterator(copyDataBuffer());
	}
}
 

public SelectResultProvider getSelectResultProvider() {
	return new SelectResultProvider() {
		@Override
		public void setJobClient(JobClient jobClient) {
			iterator.setJobClient(jobClient);
		}

		@Override
		public CloseableIterator<Row> getResultIterator() {
			return new RowIteratorWrapper(iterator);
		}
	};
}
 

@Nonnull
@Override
public Set<T> getSubsetForKeyGroup(int keyGroupId) {
	HashSet<T> result = new HashSet<>();
	PQ partitionQueue = keyGroupedHeaps[globalKeyGroupToLocalIndex(keyGroupId)];
	try (CloseableIterator<T> iterator = partitionQueue.iterator()) {
		while (iterator.hasNext()) {
			result.add(iterator.next());
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception while iterating key group.", e);
	}
	return result;
}
 

@Override
public void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, CloseableIterator<Buffer> iterator) {
	LOG.debug(
		"{} adding input data, checkpoint {}, channel: {}, startSeqNum: {}",
		taskName,
		checkpointId,
		info,
		startSeqNum);
	enqueue(write(checkpointId, info, iterator), false);
}
 

@Override
public void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, CloseableIterator<Buffer> iterator) {
	checkCheckpointId(checkpointId);
	iterator.forEachRemaining(b -> addedInput.put(info, b));
	try {
		iterator.close();
	} catch (Exception e) {
		rethrow(e);
	}
}
 

private void foreachTimer(BiConsumerWithException<N, Long, Exception> consumer, KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) throws Exception {
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			keyContext.setCurrentKey(timer.getKey());
			consumer.accept(timer.getNamespace(), timer.getTimestamp());
		}
	}
}
 

@Nonnull
@Override
public Set<T> getSubsetForKeyGroup(int keyGroupId) {
	HashSet<T> result = new HashSet<>();
	PQ partitionQueue = keyGroupedHeaps[globalKeyGroupToLocalIndex(keyGroupId)];
	try (CloseableIterator<T> iterator = partitionQueue.iterator()) {
		while (iterator.hasNext()) {
			result.add(iterator.next());
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception while iterating key group.", e);
	}
	return result;
}
 

private static void assertUnconsumedBuffer(ByteArrayOutputStream expected, CloseableIterator<Buffer> actual) throws Exception {
	if (!actual.hasNext()) {
		Assert.assertEquals(expected.size(), 0);
	}

	ByteBuffer expectedByteBuffer = ByteBuffer.wrap(expected.toByteArray());
	ByteBuffer actualByteBuffer = actual.next().getNioBufferReadable();
	Assert.assertEquals(expectedByteBuffer, actualByteBuffer);
	actual.close();
}
 

@Test
public void testBulkAddRestoredElements() throws Exception {
	final int testSize = 10;
	HashSet<TestElement> elementSet = new HashSet<>(testSize);
	for (int i = 0; i < testSize; ++i) {
		elementSet.add(new TestElement(i, i));
	}

	List<TestElement> twoTimesElementSet = new ArrayList<>(elementSet.size() * 2);

	for (TestElement testElement : elementSet) {
		twoTimesElementSet.add(testElement.deepCopy());
		twoTimesElementSet.add(testElement.deepCopy());
	}

	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	priorityQueue.addAll(twoTimesElementSet);
	priorityQueue.addAll(elementSet);

	final int expectedSize = testSetSemanticsAgainstDuplicateElements() ? elementSet.size() : 3 * elementSet.size();

	Assert.assertEquals(expectedSize, priorityQueue.size());
	try (final CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			if (testSetSemanticsAgainstDuplicateElements()) {
				Assert.assertTrue(elementSet.remove(iterator.next()));
			} else {
				Assert.assertTrue(elementSet.contains(iterator.next()));
			}
		}
	}
	if (testSetSemanticsAgainstDuplicateElements()) {
		Assert.assertTrue(elementSet.isEmpty());
	}
}
 

private int countTimersInNamespaceInternal(N namespace, InternalPriorityQueue<TimerHeapInternalTimer<K, N>> queue) {
	int count = 0;
	try (final CloseableIterator<TimerHeapInternalTimer<K, N>> iterator = queue.iterator()) {
		while (iterator.hasNext()) {
			final TimerHeapInternalTimer<K, N> timer = iterator.next();
			if (timer.getNamespace().equals(namespace)) {
				count++;
			}
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception when closing iterator.", e);
	}
	return count;
}
 

private static List<Row> materialize(TableResult results) {
    try (CloseableIterator<Row> resultIterator = results.collect()) {
        return StreamSupport
                .stream(Spliterators.spliteratorUnknownSize(resultIterator, Spliterator.ORDERED), false)
                .collect(Collectors.toList());
    } catch (Exception e) {
        throw new RuntimeException("Failed to materialize results", e);
    }
}
 

@Test
public void testBulkAddRestoredElements() throws Exception {
	final int testSize = 10;
	HashSet<TestElement> elementSet = new HashSet<>(testSize);
	for (int i = 0; i < testSize; ++i) {
		elementSet.add(new TestElement(i, i));
	}

	List<TestElement> twoTimesElementSet = new ArrayList<>(elementSet.size() * 2);

	for (TestElement testElement : elementSet) {
		twoTimesElementSet.add(testElement.deepCopy());
		twoTimesElementSet.add(testElement.deepCopy());
	}

	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	priorityQueue.addAll(twoTimesElementSet);
	priorityQueue.addAll(elementSet);

	final int expectedSize = testSetSemanticsAgainstDuplicateElements() ? elementSet.size() : 3 * elementSet.size();

	Assert.assertEquals(expectedSize, priorityQueue.size());
	try (final CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			if (testSetSemanticsAgainstDuplicateElements()) {
				Assert.assertTrue(elementSet.remove(iterator.next()));
			} else {
				Assert.assertTrue(elementSet.contains(iterator.next()));
			}
		}
	}
	if (testSetSemanticsAgainstDuplicateElements()) {
		Assert.assertTrue(elementSet.isEmpty());
	}
}
 

@Override
public void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, CloseableIterator<Buffer> iterator) {
	checkCheckpointId(checkpointId);
	try {
		iterator.close();
	} catch (Exception e) {
		rethrow(e);
	}
}
 

@Nonnull
@Override
public Set<T> getSubsetForKeyGroup(int keyGroupId) {
	HashSet<T> result = new HashSet<>();
	PQ partitionQueue = keyGroupedHeaps[globalKeyGroupToLocalIndex(keyGroupId)];
	try (CloseableIterator<T> iterator = partitionQueue.iterator()) {
		while (iterator.hasNext()) {
			result.add(iterator.next());
		}
	} catch (Exception e) {
		throw new FlinkRuntimeException("Exception while iterating key group.", e);
	}
	return result;
}
 

@Test
public void testBulkAddRestoredElements() throws Exception {
	final int testSize = 10;
	HashSet<TestElement> elementSet = new HashSet<>(testSize);
	for (int i = 0; i < testSize; ++i) {
		elementSet.add(new TestElement(i, i));
	}

	List<TestElement> twoTimesElementSet = new ArrayList<>(elementSet.size() * 2);

	for (TestElement testElement : elementSet) {
		twoTimesElementSet.add(testElement.deepCopy());
		twoTimesElementSet.add(testElement.deepCopy());
	}

	InternalPriorityQueue<TestElement> priorityQueue =
		newPriorityQueue(1);

	priorityQueue.addAll(twoTimesElementSet);
	priorityQueue.addAll(elementSet);

	final int expectedSize = testSetSemanticsAgainstDuplicateElements() ? elementSet.size() : 3 * elementSet.size();

	Assert.assertEquals(expectedSize, priorityQueue.size());
	try (final CloseableIterator<TestElement> iterator = priorityQueue.iterator()) {
		while (iterator.hasNext()) {
			if (testSetSemanticsAgainstDuplicateElements()) {
				Assert.assertTrue(elementSet.remove(iterator.next()));
			} else {
				Assert.assertTrue(elementSet.contains(iterator.next()));
			}
		}
	}
	if (testSetSemanticsAgainstDuplicateElements()) {
		Assert.assertTrue(elementSet.isEmpty());
	}
}
 

private static CloseableIterator<Buffer> toSingleBufferIterator(MemorySegment segment) {
	NetworkBuffer buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, Buffer.DataType.DATA_BUFFER, segment.size());
	return CloseableIterator.ofElement(buffer, Buffer::recycleBuffer);
}