jffwang
/
hetu-core
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support spill-to-disk feature for snapshot

This commit is contained in:
Jack Guo 2021-05-19 23:53:55 -04:00 committed by Sandy Gao
parent 116f8e84e1
commit 0b4eec3afa
10 changed files with 744 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java
View File

@ -743,9 +743,8 @@ public class SqlQueryExecution
// Doesn't work with the following features
if (SystemSessionProperties.isReuseTableScanEnabled(session)
|| SystemSessionProperties.isCTEReuseEnabled(session)
|| SystemSessionProperties.isSpillEnabled(session)) {
reasons.add("No support along with reuse_table_scan or cte_reuse_enabled or spill_enabled features");
|| SystemSessionProperties.isCTEReuseEnabled(session)) {
reasons.add("No support along with reuse_table_scan or cte_reuse_enabled features");
}
// All input tables must support snapshotting

49
presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java
View File

@ -21,6 +21,7 @@ import io.prestosql.execution.Lifespan;
import io.prestosql.memory.context.LocalMemoryContext;
import io.prestosql.snapshot.SingleInputSnapshotState;
import io.prestosql.snapshot.SnapshotStateId;
import io.prestosql.snapshot.Spillable;
import io.prestosql.snapshot.TaskSnapshotManager;
import io.prestosql.spi.Page;
import io.prestosql.spi.plan.PlanNodeId;
@ -36,6 +37,7 @@ import javax.annotation.concurrent.ThreadSafe;
import java.io.IOException;
import java.io.Serializable;
import java.nio.file.Path;
import java.util.ArrayDeque;
import java.util.HashMap;
import java.util.List;
@ -54,13 +56,20 @@ import static io.airlift.concurrent.MoreFutures.getDone;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
// Snapshot: Most of these fields are immutable objects, excpet for:
// - spilledLookupSourceHandle: content is only changed after index is fully built
// - spillInProgress: must be "done" when markers are received (see needsInput)
// - unspillInProgress: unspill can only happen after "finish", so no marker can be received after that
// - lookupSourceSupplier: only becomes non-null after "finish"
// - lookupSourceChecksum: only used when unspilling lookupSourceSupplier
// - finishMemoryRevoke: must be empty, because new input (including markers) can't be added until finishMemoryRevoke is called
@ThreadSafe
@RestorableConfig(uncapturedFields = {"lookupSourceFactory", "lookupSourceFactoryDestroyed", "outputChannels",
"hashChannels", "filterFunctionFactory", "sortChannel", "searchFunctionFactories", "singleStreamSpillerFactory",
"lookupSourceNotNeeded", "spilledLookupSourceHandle", "spiller", "spillInProgress", "unspillInProgress", "lookupSourceSupplier", "lookupSourceChecksum",
"lookupSourceNotNeeded", "spilledLookupSourceHandle", "spillInProgress", "unspillInProgress", "lookupSourceSupplier", "lookupSourceChecksum",
"finishMemoryRevoke", "snapshotState", "lastMarker"})
public class HashBuilderOperator
implements SinkOperator
implements SinkOperator, Spillable
{
public static class HashBuilderOperatorFactory
implements OperatorFactory
@ -698,6 +707,21 @@ public class HashBuilderOperator
}
}
@Override
public boolean isSpilled()
{
return state == State.SPILLING_INPUT || state == State.INPUT_SPILLED || state == State.INPUT_UNSPILLING;
}
@Override
public List<Path> getSpilledFilePaths()
{
if (isSpilled()) {
return ImmutableList.of(spiller.get().getFile());
}
return ImmutableList.of();
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
@ -709,6 +733,11 @@ public class HashBuilderOperator
myState.hashCollisionsCounter = hashCollisionsCounter.capture(serdeProvider);
myState.state = state.toString();
myState.alreadyFinished = alreadyFinished;
// Capture spill related to spill
if (spiller.isPresent()) {
myState.spiller = spiller.get().capture(serdeProvider);
}
return myState;
}
@ -723,8 +752,23 @@ public class HashBuilderOperator
this.index.restore(myState.index, serdeProvider);
this.hashCollisionsCounter.restore(myState.hashCollisionsCounter, serdeProvider);
State oldState = this.state;
this.state = State.valueOf(myState.state);
this.alreadyFinished = myState.alreadyFinished;
// Restore spill related fields
if (myState.spiller != null) {
if (!spiller.isPresent()) {
spiller = Optional.of(singleStreamSpillerFactory.create(
index.getTypes(),
operatorContext.getSpillContext().newLocalSpillContext(),
operatorContext.newLocalSystemMemoryContext(HashBuilderOperator.class.getSimpleName())));
}
this.spiller.get().restore(myState.spiller, serdeProvider);
}
if (oldState == State.CONSUMING_INPUT && this.state == State.SPILLING_INPUT) {
lookupSourceFactory.setPartitionSpilledLookupSourceHandle(partitionIndex, spilledLookupSourceHandle);
}
}
private static class HashBuilderOperatorState
@ -739,5 +783,6 @@ public class HashBuilderOperator
private Object hashCollisionsCounter;
private String state;
private boolean alreadyFinished;
private Object spiller;
}
}

90
presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java
View File

@ -16,6 +16,8 @@ package io.prestosql.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.io.Closer;
import com.google.common.util.concurrent.ListenableFuture;
import io.hetu.core.transport.execution.buffer.PagesSerde;
import io.hetu.core.transport.execution.buffer.SerializedPage;
import io.prestosql.operator.JoinProbe.JoinProbeFactory;
import io.prestosql.operator.LookupJoinOperators.JoinType;
import io.prestosql.operator.LookupSourceProvider.LookupSourceLease;
@ -25,6 +27,7 @@ import io.prestosql.snapshot.SingleInputSnapshotState;
import io.prestosql.spi.Page;
import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider;
import io.prestosql.spi.snapshot.MarkerPage;
import io.prestosql.spi.snapshot.Restorable;
import io.prestosql.spi.snapshot.RestorableConfig;
import io.prestosql.spi.type.Type;
import io.prestosql.spiller.PartitioningSpiller;
@ -57,9 +60,19 @@ import static java.util.Collections.emptyIterator;
import static java.util.Collections.emptyList;
import static java.util.Objects.requireNonNull;
// Snapshot: Most of these fields are immutable objects, excpet for:
// - lookupSourceProvider: will be established after restore when build side finishes
// - probe: must be null when addInput is called
// - outputPage: must be null when addInput is called
// - partitionGenerator: stateless
// - spillInProgress: must be "done" when markers are received (see needsInput)
// - unspilling: can only be true after finishing becomes true
// - currentPartition: only set after finishing becomes true (when unspilling is true)
// - unspilledLookupSource: only set after finishing becomes true (when unspilling is true)
// - unspilledInputPages: only set after finishing becomes true (when unspilling is true)
@RestorableConfig(uncapturedFields = {"probeTypes", "joinProbeFactory", "afterClose", "hashGenerator", "lookupSourceFactory",
"partitioningSpillerFactory", "lookupSourceProviderFuture", "lookupSourceProvider", "probe", "outputPage", "spiller",
"partitionGenerator", "spillInProgress", "savedRows", "unspilling", "currentPartition",
"partitioningSpillerFactory", "lookupSourceProviderFuture", "lookupSourceProvider", "probe", "outputPage",
"partitionGenerator", "spillInProgress", "unspilling", "currentPartition",
"unspilledLookupSource", "unspilledInputPages", "snapshotState"})
public class LookupJoinOperator
implements Operator
@ -254,12 +267,6 @@ public class LookupJoinOperator
{
requireNonNull(spillInfoSnapshot, "spillInfoSnapshot is null");
if (snapshotState != null) {
if (snapshotState.processPage(page)) {
return;
}
}
if (spillInfoSnapshot.hasSpilled()) {
page = spillAndMaskSpilledPositions(page, spillInfoSnapshot.getSpillMask());
if (page.getPositionCount() == 0) {
@ -693,6 +700,7 @@ public class LookupJoinOperator
// This class must be public because LookupJoinOperator is isolated.
public static class SavedRow
implements Restorable
{
/**
* A page with exactly one {@link Page#getPositionCount}, representing saved row.
@ -717,12 +725,48 @@ public class LookupJoinOperator
public SavedRow(Page page, int position, long joinPositionWithinPartition, boolean currentProbePositionProducedRow, int joinSourcePositions)
{
this.row = page.getSingleValuePage(position);
this(page.getSingleValuePage(position), joinPositionWithinPartition, currentProbePositionProducedRow, joinSourcePositions);
}
public SavedRow(Page row, long joinPositionWithinPartition, boolean currentProbePositionProducedRow, int joinSourcePositions)
{
this.row = row;
this.joinPositionWithinPartition = joinPositionWithinPartition;
this.currentProbePositionProducedRow = currentProbePositionProducedRow;
this.joinSourcePositions = joinSourcePositions;
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
SavedRowState myState = new SavedRowState();
PagesSerde serde = (PagesSerde) serdeProvider;
myState.row = serde.serialize(row).capture(serdeProvider);
myState.currentProbePositionProducedRow = currentProbePositionProducedRow;
myState.joinPositionWithinPartition = joinPositionWithinPartition;
myState.joinSourcePositions = joinSourcePositions;
return myState;
}
private static SavedRow restoreSavedRow(Object state, BlockEncodingSerdeProvider serdeProvider)
{
SavedRowState savedRowsState = (SavedRowState) state;
PagesSerde serde = (PagesSerde) serdeProvider;
SerializedPage sp = SerializedPage.restoreSerializedPage(savedRowsState.row);
return new SavedRow(serde.deserialize(sp),
savedRowsState.joinPositionWithinPartition,
savedRowsState.currentProbePositionProducedRow,
savedRowsState.joinSourcePositions);
}
private static class SavedRowState
implements Serializable
{
private Object row;
private long joinPositionWithinPartition;
private boolean currentProbePositionProducedRow;
private int joinSourcePositions;
}
}
private boolean tryBuildPage()
@ -770,6 +814,13 @@ public class LookupJoinOperator
myState.currentProbePositionProducedRow = currentProbePositionProducedRow;
myState.partitionedConsumption = partitionedConsumption != null ? true : false;
myState.lookupPartitions = lookupPartitions != null ? true : false;
if (spiller.isPresent()) {
myState.spiller = spiller.get().capture(serdeProvider);
}
myState.savedRows = new HashMap<>();
for (Map.Entry<Integer, SavedRow> entry : savedRows.entrySet()) {
myState.savedRows.put(entry.getKey(), entry.getValue().capture(serdeProvider));
}
return myState;
}
@ -807,12 +858,28 @@ public class LookupJoinOperator
else {
this.lookupPartitions = null;
}
if (myState.spiller != null) {
if (!spiller.isPresent()) {
spiller = Optional.of(partitioningSpillerFactory.create(
probeTypes,
getPartitionGenerator(),
operatorContext.getSpillContext().newLocalSpillContext(),
operatorContext.newAggregateSystemMemoryContext()));
}
this.spiller.get().restore(myState.spiller, serdeProvider);
}
this.savedRows.clear();
for (Map.Entry<Integer, Object> entry : myState.savedRows.entrySet()) {
SavedRow savedRow = SavedRow.restoreSavedRow(entry.getValue(), serdeProvider);
this.savedRows.put(entry.getKey(), savedRow);
}
}
private static class LookupJoinOperatorState
implements Serializable
{
//TODO-cp-I2EATR: Work to be done for spilled situations
private Object operatorContext;
private Object statisticsCounter;
private Object pageBuilder;
@ -828,5 +895,8 @@ public class LookupJoinOperator
private boolean partitionedConsumption;
private boolean lookupPartitions;
private Object spiller;
private Map<Integer, Object> savedRows;
}
}

43
presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java
View File

@ -18,6 +18,7 @@ import com.google.common.primitives.Ints;
import com.google.common.util.concurrent.ListenableFuture;
import io.prestosql.memory.context.LocalMemoryContext;
import io.prestosql.snapshot.SingleInputSnapshotState;
import io.prestosql.snapshot.Spillable;
import io.prestosql.spi.Page;
import io.prestosql.spi.block.Block;
import io.prestosql.spi.block.SortOrder;
@ -30,6 +31,7 @@ import io.prestosql.spiller.SpillerFactory;
import io.prestosql.sql.gen.OrderingCompiler;
import java.io.Serializable;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
@ -45,10 +47,14 @@ import static io.airlift.concurrent.MoreFutures.getFutureValue;
import static io.prestosql.util.MergeSortedPages.mergeSortedPages;
import static java.util.Objects.requireNonNull;
// Snapshot: Most of these fields are immutable objects, excpet for:
// - spillInProgress: must be "done" when markers are received (see needsInput)
// - finishMemoryRevoke: must be empty, because new input (including markers) can't be added until finishMemoryRevoke is called
// - sortedPages: only set after "finish" is called
@RestorableConfig(uncapturedFields = {"sortChannels", "sortOrder", "outputChannels", "sourceTypes", "spillerFactory",
"orderingCompiler", "spiller", "spillInProgress", "finishMemoryRevoke", "sortedPages", "state", "snapshotState"})
"orderingCompiler", "spillInProgress", "finishMemoryRevoke", "sortedPages", "state", "snapshotState"})
public class OrderByOperator
implements Operator
implements Operator, Spillable
{
public static class OrderByOperatorFactory
implements OperatorFactory
@ -417,6 +423,21 @@ public class OrderByOperator
spiller.ifPresent(Spiller::close);
}
@Override
public boolean isSpilled()
{
return spiller.isPresent();
}
@Override
public List<Path> getSpilledFilePaths()
{
if (isSpilled()) {
return spiller.get().getSpilledFilePaths();
}
return ImmutableList.of();
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
@ -425,6 +446,11 @@ public class OrderByOperator
myState.revocableMemoryContext = revocableMemoryContext.getBytes();
myState.localUserMemoryContext = localUserMemoryContext.getBytes();
myState.pageIndex = pageIndex.capture(serdeProvider);
// Capture spill related fields
if (spiller.isPresent()) {
myState.spiller = spiller.get().capture(serdeProvider);
}
return myState;
}
@ -436,6 +462,17 @@ public class OrderByOperator
this.revocableMemoryContext.setBytes(myState.revocableMemoryContext);
this.localUserMemoryContext.setBytes(myState.localUserMemoryContext);
this.pageIndex.restore(myState.pageIndex, serdeProvider);
// Restore spill related fields
if (myState.spiller != null) {
if (!spiller.isPresent()) {
spiller = Optional.of(spillerFactory.get().create(
sourceTypes,
operatorContext.getSpillContext(),
operatorContext.newAggregateSystemMemoryContext()));
}
this.spiller.get().restore(myState.spiller, serdeProvider);
}
}
private static class OrderByOperatorState
@ -446,5 +483,7 @@ public class OrderByOperator
private long localUserMemoryContext;
private Object pageIndex;
private Object spiller;
}
}

100
presto-main/src/main/java/io/prestosql/operator/WindowOperator.java
View File

@ -30,6 +30,7 @@ import io.prestosql.operator.WorkProcessor.TransformationState;
import io.prestosql.operator.window.FramedWindowFunction;
import io.prestosql.operator.window.WindowPartition;
import io.prestosql.snapshot.SingleInputSnapshotState;
import io.prestosql.snapshot.Spillable;
import io.prestosql.spi.Page;
import io.prestosql.spi.PageBuilder;
import io.prestosql.spi.block.Block;
@ -47,6 +48,7 @@ import javax.annotation.Nullable;
import javax.validation.constraints.NotNull;
import java.io.Serializable;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
@ -63,16 +65,18 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.Iterables.concat;
import static com.google.common.collect.Iterators.peekingIterator;
import static io.airlift.concurrent.MoreFutures.checkSuccess;
import static io.prestosql.operator.WorkProcessor.Process;
import static io.prestosql.operator.WorkProcessor.TransformationState.needsMoreData;
import static io.prestosql.spi.block.SortOrder.ASC_NULLS_LAST;
import static io.prestosql.util.MergeSortedPages.mergeSortedPages;
import static java.util.Collections.nCopies;
import static java.util.Objects.requireNonNull;
@RestorableConfig(uncapturedFields = {"outputTypes", "outputChannels",
"driverWindowInfo", "pageBuffer", "snapshotState", "pagesIndexToWindowPartitions"})
// - driverWindowInfo: only set after "close" is called
// - pageBuffer: needsInput requires pageBuffer to have null page and not finished
@RestorableConfig(uncapturedFields = {"outputTypes", "outputChannels", "driverWindowInfo", "pageBuffer", "snapshotState"})
public class WindowOperator
implements Operator
implements Operator, Spillable
{
public static class WindowOperatorFactory
implements OperatorFactory
@ -311,14 +315,14 @@ public class WindowOperator
this.outputPages = pageBuffer.pages()
.flatTransform(spillablePagesToPagesIndexes.get())
.flatMap(pagesIndexToWindowPartitions)
.flatMap(new PagesIndexToWindowPartitions())
.transform(new WindowPartitionsToOutputPages());
}
else {
this.spillablePagesToPagesIndexes = Optional.empty();
this.outputPages = pageBuffer.pages()
.transform(new PagesToPagesIndexes(inMemoryPagesIndexWithHashStrategies, orderChannels, ordering))
.flatMap(pagesIndexToWindowPartitions)
.flatMap(new PagesIndexToWindowPartitions())
.transform(new WindowPartitionsToOutputPages());
}
@ -581,11 +585,10 @@ public class WindowOperator
int pendingInputPosition;
}
private WorkProcessor.RestorableFunction<PagesIndexWithHashStrategies, WorkProcessor<WindowPartition>> pagesIndexToWindowPartitions = new WorkProcessor.RestorableFunction<PagesIndexWithHashStrategies, WorkProcessor<WindowPartition>>()
@RestorableConfig(uncapturedFields = {"this$0"})
private class PagesIndexToWindowPartitions
implements WorkProcessor.RestorableFunction<PagesIndexWithHashStrategies, WorkProcessor<WindowPartition>>
{
@RestorableConfig(uncapturedFields = {"this$0"})
private final RestorableConfig restorableConfig = null;
@Override
public WorkProcessor<WindowPartition> apply(PagesIndexWithHashStrategies pagesIndexWithHashStrategies)
{
@ -595,7 +598,7 @@ public class WindowOperator
windowInfo.addIndex(pagesIndex);
return WorkProcessor.create(new WorkProcessor.Process<WindowPartition>()
return WorkProcessor.create(new Process<WindowPartition>()
{
@RestorableConfig(uncapturedFields = {"val$pagesIndex", "val$pagesIndexWithHashStrategies", "this$1"})
private final RestorableConfig restorableConfig = null;
@ -681,7 +684,7 @@ public class WindowOperator
}
return inMemoryPagesIndexWithHashStrategies;
}
};
}
@RestorableConfig(uncapturedFields = {"this$0"})
private class WindowPartitionsToOutputPages
@ -760,7 +763,10 @@ public class WindowOperator
}
}
@RestorableConfig(unsupported = true)
// - spillInProgress: transformation is blocked until spillInProgress is done, so can't receive markers
@RestorableConfig(uncapturedFields = {
"this$0", "sourceTypes", "orderChannels", "ordering",
"spillerFactory", "pageWithPositionComparator", "spillInProgress"})
private class SpillablePagesToPagesIndexes
implements Transformation<Page, WorkProcessor<PagesIndexWithHashStrategies>>
{
@ -1047,12 +1053,51 @@ public class WindowOperator
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
return 0;
SpillablePagesToPagesIndexesState myState = new SpillablePagesToPagesIndexesState();
myState.localUserMemoryContext = localUserMemoryContext.getBytes();
myState.localRevocableMemoryContext = localRevocableMemoryContext.getBytes();
if (spiller.isPresent()) {
myState.spiller = spiller.get().capture(serdeProvider);
}
myState.inMemoryPagesIndexWithHashStrategies = inMemoryPagesIndexWithHashStrategies.capture(serdeProvider);
myState.mergedPagesIndexWithHashStrategies = mergedPagesIndexWithHashStrategies.capture(serdeProvider);
myState.spillingWhenConvertingRevocableMemory = spillingWhenConvertingRevocableMemory;
myState.resetPagesIndex = resetPagesIndex;
myState.pendingInputPosition = pendingInputPosition;
if (currentSpillGroupRowPage.isPresent()) {
PagesSerde serde = (PagesSerde) serdeProvider;
myState.currentSpillGroupRowPage = serde.serialize(currentSpillGroupRowPage.get()).capture(serdeProvider);
}
return myState;
}
@Override
public void restore(Object state, BlockEncodingSerdeProvider serdeProvider)
{
SpillablePagesToPagesIndexesState myState =
(SpillablePagesToPagesIndexesState) state;
this.localUserMemoryContext.setBytes(myState.localUserMemoryContext);
this.localRevocableMemoryContext.setBytes(myState.localRevocableMemoryContext);
if (myState.spiller != null) {
if (!spiller.isPresent()) {
this.spiller = Optional.of(spillerFactory.create(
sourceTypes,
operatorContext.getSpillContext(),
operatorContext.newAggregateSystemMemoryContext()));
}
this.spiller.get().restore(myState.spiller, serdeProvider);
}
this.inMemoryPagesIndexWithHashStrategies.restore(myState.inMemoryPagesIndexWithHashStrategies, serdeProvider);
this.mergedPagesIndexWithHashStrategies.restore(myState.mergedPagesIndexWithHashStrategies, serdeProvider);
this.spillingWhenConvertingRevocableMemory = myState.spillingWhenConvertingRevocableMemory;
this.resetPagesIndex = myState.resetPagesIndex;
this.pendingInputPosition = myState.pendingInputPosition;
if (myState.currentSpillGroupRowPage != null) {
PagesSerde serde = (PagesSerde) serdeProvider;
SerializedPage sp = SerializedPage.restoreSerializedPage(myState.currentSpillGroupRowPage);
currentSpillGroupRowPage = Optional.of(serde.deserialize(sp));
}
}
@Override
@ -1090,6 +1135,20 @@ public class WindowOperator
}
}
private static class SpillablePagesToPagesIndexesState
implements Serializable
{
private long localUserMemoryContext;
private long localRevocableMemoryContext;
private Object spiller;
private Object inMemoryPagesIndexWithHashStrategies;
private Object mergedPagesIndexWithHashStrategies;
private boolean spillingWhenConvertingRevocableMemory;
private boolean resetPagesIndex;
private int pendingInputPosition;
private Object currentSpillGroupRowPage;
}
private int updatePagesIndex(PagesIndexWithHashStrategies pagesIndexWithHashStrategies, Page page, int startPosition, Optional<Page> currentSpillGroupRowPage)
{
checkArgument(page.getPositionCount() > startPosition);
@ -1195,6 +1254,21 @@ public class WindowOperator
return right;
}
@Override
public boolean isSpilled()
{
return spillEnabled && spillablePagesToPagesIndexes.get().spiller.isPresent();
}
@Override
public List<Path> getSpilledFilePaths()
{
if (isSpilled()) {
return spillablePagesToPagesIndexes.get().spiller.get().getSpilledFilePaths();
}
return ImmutableList.of();
}
@Override
public void close()
{

72
presto-main/src/main/java/io/prestosql/operator/aggregation/builder/SpillableHashAggregationBuilder.java
View File

@ -26,6 +26,8 @@ import io.prestosql.operator.WorkProcessor;
import io.prestosql.operator.aggregation.AccumulatorFactory;
import io.prestosql.spi.Page;
import io.prestosql.spi.plan.AggregationNode;
import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider;
import io.prestosql.spi.snapshot.Restorable;
import io.prestosql.spi.snapshot.RestorableConfig;
import io.prestosql.spi.type.Type;
import io.prestosql.spiller.Spiller;
@ -33,6 +35,7 @@ import io.prestosql.spiller.SpillerFactory;
import io.prestosql.sql.gen.JoinCompiler;
import java.io.IOException;
import java.io.Serializable;
import java.util.List;
import java.util.Optional;
@ -44,10 +47,13 @@ import static io.airlift.concurrent.MoreFutures.getFutureValue;
import static io.prestosql.operator.Operator.NOT_BLOCKED;
import static java.lang.Math.max;
//TODO-cp-I39B76 should be covered in supporting spill, unsupported for now
@RestorableConfig(unsupported = true)
// - merger: this variable is used after created from either captured fields or final fields, so no need to capture.
// - mergeHashSort: this variable is used after created from operatorContext. amd operatorContext is captured, so no need to capture.
// - spillInProgress: must be "done" when markers are received.
@RestorableConfig(uncapturedFields = {"spillerFactory", "accumulatorFactories", "groupByTypes", "groupByChannels",
"hashChannel", "merger", "mergeHashSort", "spillInProgress", "joinCompiler"})
public class SpillableHashAggregationBuilder
implements AggregationBuilder
implements AggregationBuilder, Restorable
{
private InMemoryHashAggregationBuilder hashAggregationBuilder;
private final SpillerFactory spillerFactory;
@ -274,7 +280,6 @@ public class SpillableHashAggregationBuilder
WorkProcessor<Page> mergedSpilledPages = mergeHashSort.get().merge(
groupByTypes,
hashAggregationBuilder.buildIntermediateTypes(),
//TODO-cp-I39B76 need snapshot support for the iterator?
ImmutableList.<WorkProcessor<Page>>builder()
.addAll(spiller.get().getSpills().stream()
.map(WorkProcessor::fromIterator)
@ -295,7 +300,6 @@ public class SpillableHashAggregationBuilder
WorkProcessor<Page> mergedSpilledPages = mergeHashSort.get().merge(
groupByTypes,
hashAggregationBuilder.buildIntermediateTypes(),
//TODO-cp-I39B76 need snapshot support for the iterator?
spiller.get().getSpills().stream()
.map(WorkProcessor::fromIterator)
.collect(toImmutableList()),
@ -347,4 +351,62 @@ public class SpillableHashAggregationBuilder
});
emptyHashAggregationBuilderSize = hashAggregationBuilder.getSizeInMemory();
}
@Override
public Object capture(BlockEncodingSerdeProvider serdeProvider)
{
SpillableHashAggregationBuilderState myState = new SpillableHashAggregationBuilderState();
myState.operatorContext = operatorContext.capture(serdeProvider);
myState.hashAggregationBuilder = hashAggregationBuilder.capture(serdeProvider);
myState.localUserMemoryContext = localUserMemoryContext.getBytes();
myState.localRevocableMemoryContext = localRevocableMemoryContext.getBytes();
myState.emptyHashAggregationBuilderSize = emptyHashAggregationBuilderSize;
myState.hashCollisions = hashCollisions;
myState.expectedHashCollisions = expectedHashCollisions;
myState.producingOutput = producingOutput;
if (spiller.isPresent()) {
myState.spiller = spiller.get().capture(serdeProvider);
}
return myState;
}
@Override
public void restore(Object state, BlockEncodingSerdeProvider serdeProvider)
{
SpillableHashAggregationBuilderState myState = (SpillableHashAggregationBuilderState) state;
operatorContext.restore(myState.operatorContext, serdeProvider);
hashAggregationBuilder.restore(myState.hashAggregationBuilder, serdeProvider);
emptyHashAggregationBuilderSize = myState.emptyHashAggregationBuilderSize;
hashCollisions = myState.hashCollisions;
expectedHashCollisions = myState.expectedHashCollisions;
producingOutput = myState.producingOutput;
localRevocableMemoryContext.setBytes(myState.localRevocableMemoryContext);
localUserMemoryContext.setBytes(myState.localUserMemoryContext);
if (myState.spiller != null) {
if (!spiller.isPresent()) {
spiller = Optional.of(spillerFactory.create(
hashAggregationBuilder.buildTypes(),
operatorContext.getSpillContext(),
operatorContext.newAggregateSystemMemoryContext()));
}
this.spiller.get().restore(myState.spiller, serdeProvider);
}
}
private static class SpillableHashAggregationBuilderState
implements Serializable
{
private Object operatorContext;
private Object hashAggregationBuilder;
private long localRevocableMemoryContext;
private long localUserMemoryContext;
private long emptyHashAggregationBuilderSize;
private long hashCollisions;
private double expectedHashCollisions;
private boolean producingOutput;
private Object spiller;
}
}

1
presto-main/src/main/java/io/prestosql/spiller/FileSingleStreamSpiller.java
View File

@ -257,6 +257,7 @@ public class FileSingleStreamSpiller
this.pageSizeList = myState.pageSizeList;
this.targetFile.close();
Path path = Paths.get(myState.targetFile);
// Actual file content is restored after this returns, in SingleInputSnapshotState.loadSpilledFiles
Files.deleteIfExists(path);
this.targetFile = closer.register(new FileHolder(Files.createFile(path)));
for (Long pageSize : pageSizeList) {

2
presto-main/src/test/java/io/prestosql/operator/TestHashJoinOperator.java
View File

@ -349,6 +349,8 @@ public class TestHashJoinOperator
lookupJoinOperatorMapping.put("currentProbePositionProducedRow", false);
lookupJoinOperatorMapping.put("partitionedConsumption", false);
lookupJoinOperatorMapping.put("lookupPartitions", false);
lookupJoinOperatorMapping.put("spiller", null);
lookupJoinOperatorMapping.put("savedRows", 0); // Only compare map size
return lookupJoinOperatorMapping;
}

152
presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java
View File

@ -14,34 +14,52 @@
package io.prestosql.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.io.Files;
import io.airlift.units.DataSize;
import io.airlift.units.DataSize.Unit;
import io.hetu.core.filesystem.HetuLocalFileSystemClient;
import io.hetu.core.filesystem.LocalConfig;
import io.prestosql.ExceededMemoryLimitException;
import io.prestosql.Session;
import io.prestosql.filesystem.FileSystemClientManager;
import io.prestosql.metadata.InMemoryNodeManager;
import io.prestosql.operator.OrderByOperator.OrderByOperatorFactory;
import io.prestosql.snapshot.SnapshotConfig;
import io.prestosql.snapshot.SnapshotUtils;
import io.prestosql.spi.Page;
import io.prestosql.spi.plan.PlanNodeId;
import io.prestosql.spi.snapshot.MarkerPage;
import io.prestosql.spiller.FileSingleStreamSpillerFactory;
import io.prestosql.spiller.GenericSpillerFactory;
import io.prestosql.spiller.SpillerStats;
import io.prestosql.sql.gen.OrderingCompiler;
import io.prestosql.testing.MaterializedResult;
import io.prestosql.testing.TestingTaskContext;
import io.prestosql.testing.assertions.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
import static io.airlift.concurrent.MoreFutures.getFutureValue;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.airlift.testing.Assertions.assertGreaterThan;
import static io.airlift.units.DataSize.succinctBytes;
import static io.prestosql.RowPagesBuilder.rowPagesBuilder;
import static io.prestosql.SessionTestUtils.TEST_SESSION;
import static io.prestosql.SessionTestUtils.TEST_SNAPSHOT_SESSION;
import static io.prestosql.metadata.MetadataManager.createTestMetadataManager;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEquals;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEqualsWithSimpleSelfStateComparison;
import static io.prestosql.operator.OperatorAssertion.toMaterializedResult;
@ -57,6 +75,9 @@ import static io.prestosql.testing.TestingTaskContext.createTaskContext;
import static java.lang.String.format;
import static java.util.concurrent.Executors.newCachedThreadPool;
import static java.util.concurrent.Executors.newScheduledThreadPool;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
@ -66,7 +87,7 @@ public class TestOrderByOperator
private ExecutorService executor;
private ScheduledExecutorService scheduledExecutor;
private DummySpillerFactory spillerFactory;
private final SnapshotUtils snapshotUtils = NOOP_SNAPSHOT_UTILS;
private SnapshotUtils snapshotUtils = NOOP_SNAPSHOT_UTILS;
@DataProvider
public static Object[][] spillEnabled()
@ -322,4 +343,133 @@ public class TestOrderByOperator
.addPipelineContext(0, true, true, false)
.addDriverContext();
}
private static GenericSpillerFactory createGenericSpillerFactory(Path spillPath)
{
FileSingleStreamSpillerFactory streamSpillerFactory = new FileSingleStreamSpillerFactory(
listeningDecorator(newCachedThreadPool()),
createTestMetadataManager().getFunctionAndTypeManager().getBlockEncodingSerde(),
new SpillerStats(),
ImmutableList.of(spillPath),
1.0, false, false);
return new GenericSpillerFactory(streamSpillerFactory);
}
/**
* This test is supposed to consume 4 pages and produce the output page with sorted ordering.
* The spilling and capturing('capture1') happened after the first 2 pages added into the operator.
* The operator is rescheduled after 4 pages added (but before finish() is called).
*
* @throws Exception
*/
@Test
public void testCaptureRestoreWithSpill()
throws Exception
{
// Initialization
Path spillPath = Files.createTempDir().toPath();
GenericSpillerFactory spillerFactory = createGenericSpillerFactory(spillPath);
FileSystemClientManager fileSystemClientManager = mock(FileSystemClientManager.class);
when(fileSystemClientManager.getFileSystemClient(any(Path.class))).thenReturn(new HetuLocalFileSystemClient(new LocalConfig(new Properties()), Paths.get("/tmp/hetu/snapshot/")));
SnapshotConfig snapshotConfig = new SnapshotConfig();
snapshotUtils = new SnapshotUtils(fileSystemClientManager, snapshotConfig, new InMemoryNodeManager());
snapshotUtils.initialize();
List<Page> input1 = rowPagesBuilder(VARCHAR, BIGINT)
.row("a", 1L)
.row("b", 2L)
.pageBreak()
.row("b", 3L)
.row("a", 4L)
.build();
List<Page> input2 = rowPagesBuilder(VARCHAR, BIGINT)
.row("c", 4L)
.row("d", 6L)
.pageBreak()
.row("c", 2L)
.row("d", 3L)
.build();
OrderByOperatorFactory operatorFactory = new OrderByOperatorFactory(
0,
new PlanNodeId("test"),
ImmutableList.of(VARCHAR, BIGINT),
ImmutableList.of(0, 1),
10,
ImmutableList.of(0, 1),
ImmutableList.of(ASC_NULLS_LAST, DESC_NULLS_LAST),
new PagesIndex.TestingFactory(false),
true,
Optional.of(spillerFactory),
new OrderingCompiler());
DriverContext driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
OrderByOperator orderByOperator = (OrderByOperator) operatorFactory.createOperator(driverContext);
// Step1: add the first 2 pages
for (Page page : input1) {
orderByOperator.addInput(page);
}
// Step2: spilling happened here
getFutureValue(orderByOperator.startMemoryRevoke());
orderByOperator.finishMemoryRevoke();
// Step3: add a marker page to make 'capture1' happened
MarkerPage marker = MarkerPage.snapshotPage(1);
orderByOperator.addInput(marker);
// Step4: add another 2 pages
for (Page page : input2) {
orderByOperator.addInput(page);
}
// Step5: assume the task is rescheduled due to failure and everything is re-constructed
driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
operatorFactory = new OrderByOperatorFactory(
0,
new PlanNodeId("test"),
ImmutableList.of(VARCHAR, BIGINT),
ImmutableList.of(0, 1),
10,
ImmutableList.of(0, 1),
ImmutableList.of(ASC_NULLS_LAST, DESC_NULLS_LAST),
new PagesIndex.TestingFactory(false),
true,
Optional.of(spillerFactory),
new OrderingCompiler());
orderByOperator = (OrderByOperator) operatorFactory.createOperator(driverContext);
// Step6: restore to 'capture1', the spiller should contains the reference of the first 2 pages for now.
MarkerPage resumeMarker = MarkerPage.resumePage(1);
orderByOperator.addInput(resumeMarker);
// Step7: continue to add another 2 pages
for (Page page : input2) {
orderByOperator.addInput(page);
}
orderByOperator.finish();
// Compare the results
MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BIGINT)
.row("a", 4L)
.row("a", 1L)
.row("b", 3L)
.row("b", 2L)
.row("c", 4L)
.row("c", 2L)
.row("d", 6L)
.row("d", 3L)
.build();
ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
Page p = orderByOperator.getOutput();
while (p instanceof MarkerPage) {
p = orderByOperator.getOutput();
}
outputPages.add(p);
MaterializedResult actual = toMaterializedResult(driverContext.getSession(), expected.getTypes(), outputPages.build());
Assert.assertEquals(actual, expected);
}
}

267
presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java
View File

@ -14,11 +14,16 @@
package io.prestosql.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.io.Files;
import com.google.common.primitives.Ints;
import io.airlift.units.DataSize;
import io.airlift.units.DataSize.Unit;
import io.hetu.core.filesystem.HetuLocalFileSystemClient;
import io.hetu.core.filesystem.LocalConfig;
import io.prestosql.ExceededMemoryLimitException;
import io.prestosql.Session;
import io.prestosql.filesystem.FileSystemClientManager;
import io.prestosql.metadata.InMemoryNodeManager;
import io.prestosql.operator.WindowOperator.WindowOperatorFactory;
import io.prestosql.operator.window.FirstValueFunction;
import io.prestosql.operator.window.FrameInfo;
@ -28,30 +33,43 @@ import io.prestosql.operator.window.LeadFunction;
import io.prestosql.operator.window.NthValueFunction;
import io.prestosql.operator.window.ReflectionWindowFunctionSupplier;
import io.prestosql.operator.window.RowNumberFunction;
import io.prestosql.snapshot.SnapshotConfig;
import io.prestosql.snapshot.SnapshotUtils;
import io.prestosql.spi.Page;
import io.prestosql.spi.block.SortOrder;
import io.prestosql.spi.plan.PlanNodeId;
import io.prestosql.spi.snapshot.MarkerPage;
import io.prestosql.spi.type.Type;
import io.prestosql.spiller.FileSingleStreamSpillerFactory;
import io.prestosql.spiller.GenericSpillerFactory;
import io.prestosql.spiller.SpillerFactory;
import io.prestosql.spiller.SpillerStats;
import io.prestosql.sql.gen.OrderingCompiler;
import io.prestosql.testing.MaterializedResult;
import io.prestosql.testing.TestingTaskContext;
import io.prestosql.testing.assertions.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.Optional;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator;
import static io.airlift.concurrent.MoreFutures.getFutureValue;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.airlift.testing.Assertions.assertGreaterThan;
import static io.airlift.units.DataSize.succinctBytes;
import static io.prestosql.RowPagesBuilder.rowPagesBuilder;
import static io.prestosql.SessionTestUtils.TEST_SESSION;
import static io.prestosql.SessionTestUtils.TEST_SNAPSHOT_SESSION;
import static io.prestosql.metadata.MetadataManager.createTestMetadataManager;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEquals;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEqualsIgnoreOrder;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEqualsIgnoreOrderWithRestoreToNewOperator;
@ -73,6 +91,9 @@ import static io.prestosql.testing.TestingTaskContext.createTaskContext;
import static java.lang.String.format;
import static java.util.concurrent.Executors.newCachedThreadPool;
import static java.util.concurrent.Executors.newScheduledThreadPool;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
@ -102,7 +123,7 @@ public class TestWindowOperator
private ExecutorService executor;
private ScheduledExecutorService scheduledExecutor;
private DummySpillerFactory spillerFactory;
private final SnapshotUtils snapshotUtils = NOOP_SNAPSHOT_UTILS;
private SnapshotUtils snapshotUtils = NOOP_SNAPSHOT_UTILS;
@BeforeMethod
public void setUp()
@ -1142,4 +1163,248 @@ public class TestWindowOperator
.addPipelineContext(0, true, true, false)
.addDriverContext();
}
private static GenericSpillerFactory createGenericSpillerFactory(Path spillPath)
{
FileSingleStreamSpillerFactory streamSpillerFactory = new FileSingleStreamSpillerFactory(
listeningDecorator(newCachedThreadPool()),
createTestMetadataManager().getFunctionAndTypeManager().getBlockEncodingSerde(),
new SpillerStats(),
ImmutableList.of(spillPath),
1.0, false, false);
return new GenericSpillerFactory(streamSpillerFactory);
}
@Test
public void testCaptureRestoreWithSpill()
throws Exception
{
// Initialization
Path spillPath = Files.createTempDir().toPath();
GenericSpillerFactory spillerFactory = createGenericSpillerFactory(spillPath);
FileSystemClientManager fileSystemClientManager = mock(FileSystemClientManager.class);
when(fileSystemClientManager.getFileSystemClient(any(Path.class))).thenReturn(new HetuLocalFileSystemClient(new LocalConfig(new Properties()), Paths.get("/tmp/hetu/snapshot/")));
SnapshotConfig snapshotConfig = new SnapshotConfig();
snapshotUtils = new SnapshotUtils(fileSystemClientManager, snapshotConfig, new InMemoryNodeManager());
snapshotUtils.initialize();
ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
List<Page> input1 = rowPagesBuilder(VARCHAR, BIGINT, DOUBLE, BOOLEAN)
.row("b", -1L, -0.1, true)
.row("a", 2L, 0.3, false)
.row("a", 4L, 0.2, true)
.pageBreak()
.row("b", 5L, 0.4, false)
.row("a", 6L, 0.1, true)
.build();
List<Page> input2 = rowPagesBuilder(VARCHAR, BIGINT, DOUBLE, BOOLEAN)
.row("c", -1L, -0.1, true)
.row("d", 2L, 0.3, false)
.row("c", 4L, 0.2, true)
.pageBreak()
.row("d", 5L, 0.4, false)
.build();
WindowOperatorFactory operatorFactory = new WindowOperatorFactory(
0,
new PlanNodeId("test"),
ImmutableList.of(VARCHAR, BIGINT, DOUBLE, BOOLEAN),
Ints.asList(0, 1, 2, 3),
ROW_NUMBER,
Ints.asList(0),
ImmutableList.of(),
Ints.asList(1),
ImmutableList.copyOf(new SortOrder[] {SortOrder.ASC_NULLS_LAST}),
0,
10,
new PagesIndex.TestingFactory(false),
true,
spillerFactory,
new OrderingCompiler());
DriverContext driverContext = createDriverContext(0, TEST_SNAPSHOT_SESSION);
WindowOperator windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step1: add the first 2 pages
for (Page page : input1) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
// Step2: spilling happened here
getFutureValue(windowOperator.startMemoryRevoke());
windowOperator.finishMemoryRevoke();
// Step3: add a marker page to make 'capture1' happened
MarkerPage marker = MarkerPage.snapshotPage(1);
windowOperator.addInput(marker);
windowOperator.getOutput();
// Step4: add another 2 pages
for (Page page : input2) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
// Step5: assume the task is rescheduled due to failure and everything is re-constructed
driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
operatorFactory = new WindowOperatorFactory(
0,
new PlanNodeId("test"),
ImmutableList.of(VARCHAR, BIGINT, DOUBLE, BOOLEAN),
Ints.asList(0, 1, 2, 3),
ROW_NUMBER,
Ints.asList(0),
ImmutableList.of(),
Ints.asList(1),
ImmutableList.copyOf(new SortOrder[] {SortOrder.ASC_NULLS_LAST}),
0,
10,
new PagesIndex.TestingFactory(false),
true,
spillerFactory,
new OrderingCompiler());
windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step6: restore to 'capture1', the spiller should contains the reference of the first 2 pages for now.
MarkerPage resumeMarker = MarkerPage.resumePage(1);
windowOperator.addInput(resumeMarker);
windowOperator.getOutput();
// Step7: continue to add another 2 pages
for (Page page : input2) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
windowOperator.finish();
// Compare the results
MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BIGINT, DOUBLE, BOOLEAN, BIGINT)
.row("a", 2L, 0.3, false, 1L)
.row("a", 4L, 0.2, true, 2L)
.row("a", 6L, 0.1, true, 3L)
.row("b", -1L, -0.1, true, 1L)
.row("b", 5L, 0.4, false, 2L)
.row("c", -1L, -0.1, true, 1L)
.row("c", 4L, 0.2, true, 2L)
.row("d", 2L, 0.3, false, 1L)
.row("d", 5L, 0.4, false, 2L)
.build();
Page p = windowOperator.getOutput();
while (p == null) {
p = windowOperator.getOutput();
}
outputPages.add(p);
MaterializedResult actual = toMaterializedResult(driverContext.getSession(), expected.getTypes(), outputPages.build());
Assert.assertEquals(actual, expected);
}
@Test
public void testCaptureRestoreWithoutSpill()
throws Exception
{
FileSystemClientManager fileSystemClientManager = mock(FileSystemClientManager.class);
when(fileSystemClientManager.getFileSystemClient(any(Path.class))).thenReturn(new HetuLocalFileSystemClient(new LocalConfig(new Properties()), Paths.get("/tmp/hetu/snapshot/")));
SnapshotConfig snapshotConfig = new SnapshotConfig();
snapshotUtils = new SnapshotUtils(fileSystemClientManager, snapshotConfig, new InMemoryNodeManager());
snapshotUtils.initialize();
ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
List<Page> input1 = rowPagesBuilder(VARCHAR, BIGINT, DOUBLE, BOOLEAN)
.row("b", -1L, -0.1, true)
.row("a", 2L, 0.3, false)
.row("a", 4L, 0.2, true)
.pageBreak()
.row("b", 5L, 0.4, false)
.row("a", 6L, 0.1, true)
.build();
List<Page> input2 = rowPagesBuilder(VARCHAR, BIGINT, DOUBLE, BOOLEAN)
.row("c", -1L, -0.1, true)
.row("d", 2L, 0.3, false)
.row("c", 4L, 0.2, true)
.pageBreak()
.row("d", 5L, 0.4, false)
.build();
WindowOperatorFactory operatorFactory = createFactoryUnbounded(
ImmutableList.of(VARCHAR, BIGINT, DOUBLE, BOOLEAN),
Ints.asList(0, 1, 2, 3),
ROW_NUMBER,
Ints.asList(0),
Ints.asList(1),
ImmutableList.copyOf(new SortOrder[] {SortOrder.ASC_NULLS_LAST}),
false);
DriverContext driverContext = createDriverContext(0, TEST_SNAPSHOT_SESSION);
WindowOperator windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step1: add the first 2 pages
for (Page page : input1) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
// Step2: add a marker page to make 'capture1' happened
MarkerPage marker = MarkerPage.snapshotPage(1);
windowOperator.addInput(marker);
windowOperator.getOutput();
// Step3: add another 2 pages
for (Page page : input2) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
// Step4: assume the task is rescheduled due to failure and everything is re-constructed
driverContext = createDriverContext(8, TEST_SNAPSHOT_SESSION);
operatorFactory = createFactoryUnbounded(
ImmutableList.of(VARCHAR, BIGINT, DOUBLE, BOOLEAN),
Ints.asList(0, 1, 2, 3),
ROW_NUMBER,
Ints.asList(0),
Ints.asList(1),
ImmutableList.copyOf(new SortOrder[] {SortOrder.ASC_NULLS_LAST}),
false);
windowOperator = (WindowOperator) operatorFactory.createOperator(driverContext);
// Step5: restore to 'capture1'
MarkerPage resumeMarker = MarkerPage.resumePage(1);
windowOperator.addInput(resumeMarker);
windowOperator.getOutput();
// Step6: continue to add another 2 pages
for (Page page : input2) {
windowOperator.addInput(page);
windowOperator.getOutput();
}
windowOperator.finish();
// Compare the results
MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BIGINT, DOUBLE, BOOLEAN, BIGINT)
.row("a", 2L, 0.3, false, 1L)
.row("a", 4L, 0.2, true, 2L)
.row("a", 6L, 0.1, true, 3L)
.row("b", -1L, -0.1, true, 1L)
.row("b", 5L, 0.4, false, 2L)
.row("c", -1L, -0.1, true, 1L)
.row("c", 4L, 0.2, true, 2L)
.row("d", 2L, 0.3, false, 1L)
.row("d", 5L, 0.4, false, 2L)
.build();
Page p = windowOperator.getOutput();
while (p == null) {
p = windowOperator.getOutput();
}
outputPages.add(p);
MaterializedResult actual = toMaterializedResult(driverContext.getSession(), expected.getTypes(), outputPages.build());
Assert.assertEquals(actual, expected);
}
}