org.apache.flink.util.CloseableIterator Java Examples

The following examples show how to use org.apache.flink.util.CloseableIterator. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: RemoteInputChannel.java    From flink with Apache License 2.0 6 votes vote down vote up
@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));
	}
}
 
Example #2
Source File: ChannelStateWriteRequest.java    From flink with Apache License 2.0 6 votes vote down vote up
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);
}
 
Example #3
Source File: InternalPriorityQueueTestBase.java    From flink with Apache License 2.0 6 votes vote down vote up
@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());
	}
}
 
Example #4
Source File: SpanningWrapperTest.java    From flink with Apache License 2.0 6 votes vote down vote up
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);
	}
}
 
Example #5
Source File: SpanningWrapperTest.java    From flink with Apache License 2.0 6 votes vote down vote up
@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));
}
 
Example #6
Source File: InternalPriorityQueueTestBase.java    From Flink-CEPplus with Apache License 2.0 6 votes vote down vote up
@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());
	}
}
 
Example #7
Source File: InternalPriorityQueueTestBase.java    From flink with Apache License 2.0 6 votes vote down vote up
@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());
	}
}
 
Example #8
Source File: SpanningWrapper.java    From flink with Apache License 2.0 5 votes vote down vote up
@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))
	);
}
 
Example #9
Source File: InternalTimerServiceImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
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());
		}
	}
}
 
Example #10
Source File: InternalTimerServiceImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
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;
}
 
Example #11
Source File: StreamSelectTableSink.java    From flink with Apache License 2.0 5 votes vote down vote up
public SelectResultProvider getSelectResultProvider() {
	return new SelectResultProvider() {

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

		@Override
		public CloseableIterator<Row> getResultIterator() {
			return new RowIteratorWrapper(iterator);
		}
	};
}
 
Example #12
Source File: TableResultImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
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");
}
 
Example #13
Source File: NonSpanningWrapper.java    From flink with Apache License 2.0 5 votes vote down vote up
CloseableIterator<Buffer> getUnconsumedSegment() {
	if (!hasRemaining()) {
		return CloseableIterator.empty();
	}
	MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(remaining());
	this.segment.copyTo(position, segment, 0, remaining());
	return singleBufferIterator(segment);
}
 
Example #14
Source File: InternalTimerServiceImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
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;
}
 
Example #15
Source File: SpanningWrapper.java    From flink with Apache License 2.0 5 votes vote down vote up
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());
	}
}
 
Example #16
Source File: SelectTableSinkBase.java    From flink with Apache License 2.0 5 votes vote down vote up
public SelectResultProvider getSelectResultProvider() {
	return new SelectResultProvider() {
		@Override
		public void setJobClient(JobClient jobClient) {
			iterator.setJobClient(jobClient);
		}

		@Override
		public CloseableIterator<Row> getResultIterator() {
			return new RowIteratorWrapper(iterator);
		}
	};
}
 
Example #17
Source File: KeyGroupPartitionedPriorityQueue.java    From flink with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #18
Source File: ChannelStateWriterImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example #19
Source File: RecordingChannelStateWriter.java    From flink with Apache License 2.0 5 votes vote down vote up
@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);
	}
}
 
Example #20
Source File: InternalTimerServiceImpl.java    From flink with Apache License 2.0 5 votes vote down vote up
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());
		}
	}
}
 
Example #21
Source File: KeyGroupPartitionedPriorityQueue.java    From Flink-CEPplus with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #22
Source File: SpanningRecordSerializationTest.java    From flink with Apache License 2.0 5 votes vote down vote up
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();
}
 
Example #23
Source File: InternalPriorityQueueTestBase.java    From Flink-CEPplus with Apache License 2.0 5 votes vote down vote up
@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());
	}
}
 
Example #24
Source File: InternalTimerServiceImpl.java    From Flink-CEPplus with Apache License 2.0 5 votes vote down vote up
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;
}
 
Example #25
Source File: SpendReportTest.java    From flink-playgrounds with Apache License 2.0 5 votes vote down vote up
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);
    }
}
 
Example #26
Source File: InternalPriorityQueueTestBase.java    From flink with Apache License 2.0 5 votes vote down vote up
@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());
	}
}
 
Example #27
Source File: MockChannelStateWriter.java    From flink with Apache License 2.0 5 votes vote down vote up
@Override
public void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, CloseableIterator<Buffer> iterator) {
	checkCheckpointId(checkpointId);
	try {
		iterator.close();
	} catch (Exception e) {
		rethrow(e);
	}
}
 
Example #28
Source File: KeyGroupPartitionedPriorityQueue.java    From flink with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #29
Source File: InternalPriorityQueueTestBase.java    From flink with Apache License 2.0 5 votes vote down vote up
@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());
	}
}
 
Example #30
Source File: SpanningWrapper.java    From flink with Apache License 2.0 4 votes vote down vote up
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);
}