Examples with BlockEncodingSerdeProvider io.prestosql.spi.snapshot.BlockEncodingSerdeProvider used on opensource projects

Example 1 with BlockEncodingSerdeProvider

use of io.prestosql.spi.snapshot.BlockEncodingSerdeProvider 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 BlockEncodingSerdeProvider

use of io.prestosql.spi.snapshot.BlockEncodingSerdeProvider in project hetu-core by openlookeng.

the class DummySpillerFactory method create.

@Override
public Spiller create(List<Type> types, SpillContext spillContext, AggregatedMemoryContext memoryContext) {
    return new Spiller() {

        @RestorableConfig(unsupported = true)
        private final RestorableConfig restorableConfig = null;

        private final List<Iterable<Page>> spills = new ArrayList<>();

        private final List<AtomicBoolean> spillCommitted = new ArrayList<>();

        @Override
        public ListenableFuture<?> spill(Iterator<Page> pageIterator) {
            spillsCount++;
            spills.add(ImmutableList.copyOf(pageIterator));
            spillCommitted.add(new AtomicBoolean(true));
            return immediateFuture(null);
        }

        @Override
        public Pair<ListenableFuture<?>, Runnable> spillUnCommit(Iterator<Page> pageIterator) {
            spillsCount++;
            spills.add(ImmutableList.copyOf(pageIterator));
            AtomicBoolean isCommitted = new AtomicBoolean(false);
            spillCommitted.add(isCommitted);
            return ImmutablePair.of(immediateFuture(null), () -> isCommitted.set(true));
        }

        @Override
        public List<Iterator<Page>> getSpills() {
            return spills.stream().map(Iterable::iterator).collect(toImmutableList());
        }

        @Override
        public void close() {
            spills.clear();
        }

        /**
         * Capture this object's internal state, so it can be used later to restore to the same state.
         *
         * @param serdeProvider
         * @return An object representing internal state of the current object
         */
        @Override
        public Object capture(BlockEncodingSerdeProvider serdeProvider) {
            DummySpillerState myState = new DummySpillerState();
            for (int i = 0; i < spills.size(); i++) {
                if (spillCommitted.get(i).get()) {
                    List<Page> pages = new ArrayList<>();
                    spills.get(i).forEach(pg -> pages.add(pg));
                    myState.spills.add(pages);
                }
            }
            return myState;
        }

        /**
         * Restore this object's internal state according to the snapshot
         *
         * @param state         an object that represents this object's snapshot state
         * @param serdeProvider
         */
        @Override
        public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
            DummySpillerState myState = (DummySpillerState) state;
            this.spills.clear();
            for (List<Page> s : myState.spills) {
                this.spills.add(s);
                this.spillCommitted.add(new AtomicBoolean(true));
            }
        }

        class DummySpillerState implements Serializable {

            List<List<Page>> spills = new ArrayList<>();
        }
    };
}

Example 3 with BlockEncodingSerdeProvider

use of io.prestosql.spi.snapshot.BlockEncodingSerdeProvider 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 4 with BlockEncodingSerdeProvider

use of io.prestosql.spi.snapshot.BlockEncodingSerdeProvider in project hetu-core by openlookeng.

the class DynamicFilterSourceOperator method restore.

@Override
public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) {
    DynamicFilterSourceOperatorState myState = (DynamicFilterSourceOperatorState) state;
    this.context.restore(myState.context, serdeProvider);
    if (myState.values == null) {
        this.values = null;
    } else {
        if (this.values == null) {
            this.values = new HashMap<>();
            for (Channel channel : channels) {
                this.values.put(channel, new HashSet<>());
            }
        }
        checkState(myState.values.length == channels.size());
        for (int i = 0; i < channels.size(); i++) {
            Type valueType = channels.get(i).type;
            Set<?> set = Arrays.stream(myState.values[i]).map(value -> SnapshotUtils.restoreHelper(value, valueType.getJavaType(), serdeProvider)).collect(Collectors.toSet());
            values.get(channels.get(i)).clear();
            values.get(channels.get(i)).addAll(set);
        }
    }
}

Example 5 with BlockEncodingSerdeProvider

use of io.prestosql.spi.snapshot.BlockEncodingSerdeProvider 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);
        }
    });
}