Examples with TransformationState io.prestosql.operator.WorkProcessor.TransformationState used on opensource projects

Example 1 with TransformationState

use of io.prestosql.operator.WorkProcessor.TransformationState in project hetu-core by openlookeng.

the class MergeSortedPages method buildPage.

private static WorkProcessor<Page> buildPage(WorkProcessor<PageWithPosition> pageWithPositions, List<Integer> outputChannels, List<Type> outputTypes, BiPredicate<PageBuilder, PageWithPosition> pageBreakPredicate, boolean updateMemoryAfterEveryPosition, AggregatedMemoryContext aggregatedMemoryContext, DriverYieldSignal yieldSignal) {
    LocalMemoryContext memoryContext = aggregatedMemoryContext.newLocalMemoryContext(MergeSortedPages.class.getSimpleName());
    PageBuilder pageBuilder = new PageBuilder(outputTypes);
    return pageWithPositions.yielding(yieldSignal::isSet).transform(new WorkProcessor.Transformation<PageWithPosition, Page>() {

        @RestorableConfig(stateClassName = "BuildPageState", uncapturedFields = { "val$outputTypes", "val$pageBreakPredicate", "val$outputChannels" })
        private final RestorableConfig restorableConfig = null;

        @Override
        public TransformationState<Page> process(@Nullable PageWithPosition pageWithPosition) {
            boolean finished = pageWithPosition == null;
            if (finished && pageBuilder.isEmpty()) {
                memoryContext.close();
                return TransformationState.finished();
            }
            if (finished || pageBreakPredicate.test(pageBuilder, pageWithPosition)) {
                if (!updateMemoryAfterEveryPosition) {
                    // update memory usage just before producing page to cap from top
                    memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
                }
                Page page = pageBuilder.build();
                pageBuilder.reset();
                if (!finished) {
                    pageWithPosition.appendTo(pageBuilder, outputChannels, outputTypes);
                }
                if (updateMemoryAfterEveryPosition) {
                    memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
                }
                return TransformationState.ofResult(page, !finished);
            }
            pageWithPosition.appendTo(pageBuilder, outputChannels, outputTypes);
            if (updateMemoryAfterEveryPosition) {
                memoryContext.setBytes(pageBuilder.getRetainedSizeInBytes());
            }
            return TransformationState.needsMoreData();
        }

        @Override
        public Object captureResult(Page result, BlockEncodingSerdeProvider serdeProvider) {
            if (result != null) {
                SerializedPage serializedPage = ((PagesSerde) serdeProvider).serialize(result);
                return serializedPage.capture(serdeProvider);
            }
            return null;
        }

        @Override
        public Page restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            if (resultState != null) {
                return ((PagesSerde) serdeProvider).deserialize(SerializedPage.restoreSerializedPage(resultState));
            }
            return null;
        }

        @Override
        public Object captureInput(PageWithPosition input, BlockEncodingSerdeProvider serdeProvider) {
            if (input != null) {
                return input.capture(serdeProvider);
            }
            return null;
        }

        @Override
        public PageWithPosition restoreInput(Object inputState, PageWithPosition input, BlockEncodingSerdeProvider serdeProvider) {
            if (inputState != null) {
                return PageWithPosition.restorePageWithPosition(inputState, serdeProvider);
            }
            return null;
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            BuildPageState myState = new BuildPageState();
            myState.pageBuilder = pageBuilder.capture(serdeProvider);
            myState.memoryContext = memoryContext.getBytes();
            return myState;
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            BuildPageState myState = (BuildPageState) state;
            pageBuilder.restore(myState.pageBuilder, serdeProvider);
            memoryContext.setBytes(myState.memoryContext);
        }
    });
}

Example 2 with TransformationState

use of io.prestosql.operator.WorkProcessor.TransformationState in project hetu-core by openlookeng.

the class WorkProcessorUtils method flatMap.

static <T, R> WorkProcessor<R> flatMap(WorkProcessor<T> processor, WorkProcessor.RestorableFunction<T, WorkProcessor<R>> mapper) {
    requireNonNull(processor, "processor is null");
    requireNonNull(mapper, "mapper is null");
    return processor.flatTransform(new Transformation<T, WorkProcessor<R>>() {

        @RestorableConfig(uncapturedFields = "mapper")
        private final RestorableConfig restorableConfig = null;

        T savedElement;

        @Override
        public TransformationState<WorkProcessor<R>> process(@org.jetbrains.annotations.Nullable T element) {
            if (element == null) {
                return TransformationState.finished();
            }
            savedElement = element;
            return TransformationState.ofResult(mapper.apply(element));
        }

        @Override
        public Object captureResult(WorkProcessor<R> result, BlockEncodingSerdeProvider serdeProvider) {
            return result.capture(serdeProvider);
        }

        @Override
        public WorkProcessor<R> restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            if (savedElement != null) {
                WorkProcessor<R> restoredProcessor = mapper.apply(savedElement);
                restoredProcessor.restore(resultState, serdeProvider);
                return restoredProcessor;
            }
            return null;
        }

        @Override
        public Object captureInput(T input, BlockEncodingSerdeProvider serdeProvider) {
            return mapper.captureInput(input, serdeProvider);
        }

        @Override
        public T restoreInput(Object inputState, T input, BlockEncodingSerdeProvider serdeProvider) {
            return mapper.restoreInput(inputState, serdeProvider);
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            return mapper.captureInput(savedElement, serdeProvider);
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            this.savedElement = mapper.restoreInput(state, serdeProvider);
        }
    });
}

Example 3 with TransformationState

use of io.prestosql.operator.WorkProcessor.TransformationState in project hetu-core by openlookeng.

the class WorkProcessorUtils method transform.

static <T, R> WorkProcessor<R> transform(WorkProcessor<T> processor, Transformation<T, R> transformation) {
    requireNonNull(processor, "processor is null");
    requireNonNull(transformation, "transformation is null");
    return create(new WorkProcessor.Process<R>() {

        @RestorableConfig(stateClassName = "TransformState", uncapturedFields = { "element" })
        private final RestorableConfig restorableConfig = null;

        T element;

        @Override
        public ProcessState<R> 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.yield();
                    }
                }
                TransformationState<R> 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.yield();
                    case RESULT:
                        return ProcessState.ofResult(state.getResult());
                    case FINISHED:
                        return ProcessState.finished();
                }
            }
        }

        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            TransformState myState = new TransformState();
            myState.transformation = transformation.capture(serdeProvider);
            myState.validElement = element != null;
            myState.processor = processor.capture(serdeProvider);
            return myState;
        }

        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            TransformState myState = (TransformState) state;
            transformation.restore(myState.transformation, serdeProvider);
            processor.restore(myState.processor, serdeProvider);
            if (myState.validElement) {
                checkArgument(processor.getResult() != null);
                element = processor.getResult();
            }
        }

        @Override
        public Object captureResult(R result, BlockEncodingSerdeProvider serdeProvider) {
            return transformation.captureResult(result, serdeProvider);
        }

        @Override
        public R restoreResult(Object resultState, BlockEncodingSerdeProvider serdeProvider) {
            return transformation.restoreResult(resultState, serdeProvider);
        }
    });
}