All Downloads are FREE. Search and download functionalities are using the official Maven repository.
Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
io.trino.operator.WorkProcessorUtils Maven / Gradle / Ivy
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.operator;
import com.google.common.collect.AbstractIterator;
import com.google.common.util.concurrent.ListenableFuture;
import io.trino.operator.WorkProcessor.ProcessState;
import io.trino.operator.WorkProcessor.Transformation;
import io.trino.operator.WorkProcessor.TransformationState;
import jakarta.annotation.Nullable;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Optional;
import java.util.PriorityQueue;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Comparator.comparing;
import static java.util.Objects.requireNonNull;
public final class WorkProcessorUtils
{
private WorkProcessorUtils() {}
static Iterator iteratorFrom(WorkProcessor processor)
{
requireNonNull(processor, "processor is null");
return new AbstractIterator<>()
{
final Iterator> yieldingIterator = yieldingIteratorFrom(processor);
@Override
protected T computeNext()
{
if (!yieldingIterator.hasNext()) {
return endOfData();
}
return yieldingIterator.next()
.orElseThrow(() -> new IllegalStateException("Cannot iterate over yielding WorkProcessor"));
}
};
}
static Iterator> yieldingIteratorFrom(WorkProcessor processor)
{
return new YieldingIterator<>(processor);
}
private static class YieldingIterator
extends AbstractIterator>
{
@Nullable
WorkProcessor processor;
YieldingIterator(WorkProcessor processor)
{
this.processor = requireNonNull(processor, "processorParameter is null");
}
@Override
protected Optional computeNext()
{
if (processor.process()) {
if (processor.isFinished()) {
processor = null;
return endOfData();
}
return Optional.of(processor.getResult());
}
if (processor.isBlocked()) {
throw new IllegalStateException("Cannot iterate over blocking WorkProcessor");
}
// yielded
return Optional.empty();
}
}
static WorkProcessor fromIterator(Iterator iterator)
{
requireNonNull(iterator, "iterator is null");
return create(() -> {
if (!iterator.hasNext()) {
return ProcessState.finished();
}
return ProcessState.ofResult(iterator.next());
});
}
static WorkProcessor mergeSorted(Iterable> processorIterable, Comparator comparator)
{
requireNonNull(comparator, "comparator is null");
Iterator> processorIterator = processorIterable.iterator();
checkArgument(processorIterator.hasNext(), "There must be at least one base processor");
PriorityQueue> queue = new PriorityQueue<>(2, comparing(ElementAndProcessor::getElement, comparator));
return create(new WorkProcessor.Process<>()
{
WorkProcessor processor = requireNonNull(processorIterator.next());
@Override
public ProcessState process()
{
while (true) {
if (processor.process()) {
if (!processor.isFinished()) {
queue.add(new ElementAndProcessor<>(processor.getResult(), processor));
}
}
else if (processor.isBlocked()) {
return ProcessState.blocked(processor.getBlockedFuture());
}
else {
return ProcessState.yielded();
}
if (processorIterator.hasNext()) {
processor = requireNonNull(processorIterator.next());
continue;
}
if (queue.isEmpty()) {
return ProcessState.finished();
}
ElementAndProcessor elementAndProcessor = queue.poll();
processor = elementAndProcessor.getProcessor();
return ProcessState.ofResult(elementAndProcessor.getElement());
}
}
});
}
static WorkProcessor yielding(WorkProcessor processor, BooleanSupplier yieldSignal)
{
return WorkProcessor.create(new YieldingProcess<>(processor, yieldSignal));
}
private static class YieldingProcess
implements WorkProcessor.Process
{
final WorkProcessor processor;
final BooleanSupplier yieldSignal;
boolean lastProcessYielded;
YieldingProcess(WorkProcessor processor, BooleanSupplier yieldSignal)
{
this.processor = requireNonNull(processor, "processor is null");
this.yieldSignal = requireNonNull(yieldSignal, "yieldSignal is null");
}
@Override
public ProcessState process()
{
if (!lastProcessYielded && yieldSignal.getAsBoolean()) {
lastProcessYielded = true;
return ProcessState.yielded();
}
lastProcessYielded = false;
return getNextState(processor);
}
}
static WorkProcessor blocking(WorkProcessor processor, Supplier> futureSupplier)
{
return WorkProcessor.create(new BlockingProcess<>(processor, futureSupplier));
}
private static class BlockingProcess
implements WorkProcessor.Process
{
final WorkProcessor processor;
final Supplier> futureSupplier;
ProcessState state;
BlockingProcess(WorkProcessor processor, Supplier> futureSupplier)
{
this.processor = requireNonNull(processor, "processor is null");
this.futureSupplier = requireNonNull(futureSupplier, "futureSupplier is null");
}
@Override
public ProcessState process()
{
if (state == null) {
state = getNextState(processor);
}
ListenableFuture future = futureSupplier.get();
if (!future.isDone()) {
if (state.getType() == ProcessState.Type.YIELD) {
// clear yielded state to continue computations in the next iteration
state = null;
}
return ProcessState.blocked(future);
}
ProcessState result = state;
state = null;
return result;
}
}
static WorkProcessor processEntryMonitor(WorkProcessor processor, Runnable monitor)
{
requireNonNull(processor, "processor is null");
requireNonNull(monitor, "monitor is null");
return WorkProcessor.create(() -> {
monitor.run();
return getNextState(processor);
});
}
static WorkProcessor processStateMonitor(WorkProcessor processor, Consumer> monitor)
{
requireNonNull(processor, "processor is null");
requireNonNull(monitor, "monitor is null");
return WorkProcessor.create(() -> {
ProcessState state = getNextState(processor);
monitor.accept(state);
return state;
});
}
static WorkProcessor finishWhen(WorkProcessor processor, BooleanSupplier finishSignal)
{
requireNonNull(processor, "processor is null");
requireNonNull(finishSignal, "finishSignal is null");
return WorkProcessor.create(() -> {
if (finishSignal.getAsBoolean()) {
return ProcessState.finished();
}
return getNextState(processor);
});
}
private static ProcessState getNextState(WorkProcessor processor)
{
if (processor.process()) {
if (processor.isFinished()) {
return ProcessState.finished();
}
return ProcessState.ofResult(processor.getResult());
}
if (processor.isBlocked()) {
return ProcessState.blocked(processor.getBlockedFuture());
}
return ProcessState.yielded();
}
static WorkProcessor flatMap(WorkProcessor processor, Function> mapper)
{
requireNonNull(processor, "processor is null");
requireNonNull(mapper, "mapper is null");
return processor.flatTransform(element -> {
if (element == null) {
return TransformationState.finished();
}
return TransformationState.ofResult(mapper.apply(element));
});
}
static WorkProcessor map(WorkProcessor processor, Function mapper)
{
requireNonNull(processor, "processor is null");
requireNonNull(mapper, "mapper is null");
return processor.transform(element -> {
if (element == null) {
return TransformationState.finished();
}
return TransformationState.ofResult(mapper.apply(element));
});
}
static WorkProcessor flatTransform(WorkProcessor processor, Transformation> transformation)
{
requireNonNull(processor, "processor is null");
requireNonNull(transformation, "transformation is null");
return processor.transform(transformation).transformProcessor(WorkProcessorUtils::flatten);
}
static WorkProcessor flatten(WorkProcessor> processor)
{
requireNonNull(processor, "processor is null");
return processor.transform(nestedProcessor -> {
if (nestedProcessor == null) {
return TransformationState.finished();
}
if (nestedProcessor.process()) {
if (nestedProcessor.isFinished()) {
return TransformationState.needsMoreData();
}
return TransformationState.ofResult(nestedProcessor.getResult(), false);
}
if (nestedProcessor.isBlocked()) {
return TransformationState.blocked(nestedProcessor.getBlockedFuture());
}
return TransformationState.yielded();
});
}
static WorkProcessor transform(WorkProcessor processor, Transformation transformation)
{
requireNonNull(processor, "processor is null");
requireNonNull(transformation, "transformation is null");
return create(new WorkProcessor.Process<>()
{
T element;
@Override
public ProcessState process()
{
while (true) {
if (element == null && !processor.isFinished()) {
if (processor.process()) {
if (!processor.isFinished()) {
element = requireNonNull(processor.getResult(), "result is null");
}
}
else if (processor.isBlocked()) {
return ProcessState.blocked(processor.getBlockedFuture());
}
else {
return ProcessState.yielded();
}
}
TransformationState state = requireNonNull(transformation.process(element), "state is null");
if (state.isNeedsMoreData()) {
checkState(!processor.isFinished(), "Cannot request more data when base processor is finished");
// set element to empty() in order to fetch a new one
element = null;
}
// pass-through transformation state if it doesn't require new data
switch (state.getType()) {
case NEEDS_MORE_DATA:
break;
case BLOCKED:
return ProcessState.blocked(state.getBlocked());
case YIELD:
return ProcessState.yielded();
case RESULT:
return ProcessState.ofResult(state.getResult());
case FINISHED:
return ProcessState.finished();
}
}
}
});
}
static WorkProcessor create(WorkProcessor.Process process)
{
return new ProcessWorkProcessor<>(process);
}
private static class ProcessWorkProcessor
implements WorkProcessor
{
@Nullable
WorkProcessor.Process process;
// set initial state to yield as it will cause processor computations to progress
ProcessState state = ProcessState.yielded();
ProcessWorkProcessor(WorkProcessor.Process process)
{
this.process = requireNonNull(process, "process is null");
}
@Override
public boolean process()
{
if (isBlocked()) {
return false;
}
if (isFinished()) {
return true;
}
state = requireNonNull(process.process());
if (state.getType() == ProcessState.Type.FINISHED) {
process = null;
return true;
}
return state.getType() == ProcessState.Type.RESULT;
}
@Override
public boolean isBlocked()
{
return state.getType() == ProcessState.Type.BLOCKED && !state.getBlocked().isDone();
}
@Override
public ListenableFuture getBlockedFuture()
{
checkState(state.getType() == ProcessState.Type.BLOCKED, "Must be blocked to get blocked future");
return state.getBlocked();
}
@Override
public boolean isFinished()
{
return state.getType() == ProcessState.Type.FINISHED;
}
@Override
public T getResult()
{
checkState(state.getType() == ProcessState.Type.RESULT, "process() must return true and must not be finished");
return state.getResult();
}
}
private static class ElementAndProcessor
{
@Nullable final T element;
final WorkProcessor processor;
ElementAndProcessor(T element, WorkProcessor processor)
{
this.element = element;
this.processor = requireNonNull(processor, "processor is null");
}
T getElement()
{
return element;
}
WorkProcessor getProcessor()
{
return processor;
}
}
}
