From 0b4eec3afa648861f4370dfa87a9ec1a8dc5a80d Mon Sep 17 00:00:00 2001 From: Jack Guo Date: Wed, 19 May 2021 23:53:55 -0400 Subject: [PATCH] Support spill-to-disk feature for snapshot --- .../execution/SqlQueryExecution.java | 5 +- .../operator/HashBuilderOperator.java | 49 +++- .../operator/LookupJoinOperator.java | 90 +++++- .../prestosql/operator/OrderByOperator.java | 43 ++- .../io/prestosql/operator/WindowOperator.java | 100 ++++++- .../SpillableHashAggregationBuilder.java | 72 ++++- .../spiller/FileSingleStreamSpiller.java | 1 + .../operator/TestHashJoinOperator.java | 2 + .../operator/TestOrderByOperator.java | 152 +++++++++- .../operator/TestWindowOperator.java | 267 +++++++++++++++++- 10 files changed, 744 insertions(+), 37 deletions(-) diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java index d5fbddfee..6bde77c78 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java @@ -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 diff --git a/presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java b/presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java index 885303fb7..fac616f64 100644 --- a/presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java +++ b/presto-main/src/main/java/io/prestosql/operator/HashBuilderOperator.java @@ -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 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; } } diff --git a/presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java b/presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java index c8cc6e0f0..012ca9924 100644 --- a/presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java +++ b/presto-main/src/main/java/io/prestosql/operator/LookupJoinOperator.java @@ -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 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 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 savedRows; } } diff --git a/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java b/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java index 7b609cfaa..caae803a7 100644 --- a/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java +++ b/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java @@ -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 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; } } diff --git a/presto-main/src/main/java/io/prestosql/operator/WindowOperator.java b/presto-main/src/main/java/io/prestosql/operator/WindowOperator.java index c2d18d2c9..cc32466ba 100644 --- a/presto-main/src/main/java/io/prestosql/operator/WindowOperator.java +++ b/presto-main/src/main/java/io/prestosql/operator/WindowOperator.java @@ -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> pagesIndexToWindowPartitions = new WorkProcessor.RestorableFunction>() + @RestorableConfig(uncapturedFields = {"this$0"}) + private class PagesIndexToWindowPartitions + implements WorkProcessor.RestorableFunction> { - @RestorableConfig(uncapturedFields = {"this$0"}) - private final RestorableConfig restorableConfig = null; - @Override public WorkProcessor apply(PagesIndexWithHashStrategies pagesIndexWithHashStrategies) { @@ -595,7 +598,7 @@ public class WindowOperator windowInfo.addIndex(pagesIndex); - return WorkProcessor.create(new WorkProcessor.Process() + return WorkProcessor.create(new Process() { @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> { @@ -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 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 getSpilledFilePaths() + { + if (isSpilled()) { + return spillablePagesToPagesIndexes.get().spiller.get().getSpilledFilePaths(); + } + return ImmutableList.of(); + } + @Override public void close() { diff --git a/presto-main/src/main/java/io/prestosql/operator/aggregation/builder/SpillableHashAggregationBuilder.java b/presto-main/src/main/java/io/prestosql/operator/aggregation/builder/SpillableHashAggregationBuilder.java index b97af27f3..9cdb31c23 100644 --- a/presto-main/src/main/java/io/prestosql/operator/aggregation/builder/SpillableHashAggregationBuilder.java +++ b/presto-main/src/main/java/io/prestosql/operator/aggregation/builder/SpillableHashAggregationBuilder.java @@ -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 mergedSpilledPages = mergeHashSort.get().merge( groupByTypes, hashAggregationBuilder.buildIntermediateTypes(), - //TODO-cp-I39B76 need snapshot support for the iterator? ImmutableList.>builder() .addAll(spiller.get().getSpills().stream() .map(WorkProcessor::fromIterator) @@ -295,7 +300,6 @@ public class SpillableHashAggregationBuilder WorkProcessor 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; + } } diff --git a/presto-main/src/main/java/io/prestosql/spiller/FileSingleStreamSpiller.java b/presto-main/src/main/java/io/prestosql/spiller/FileSingleStreamSpiller.java index da544a337..9e7be0acb 100644 --- a/presto-main/src/main/java/io/prestosql/spiller/FileSingleStreamSpiller.java +++ b/presto-main/src/main/java/io/prestosql/spiller/FileSingleStreamSpiller.java @@ -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) { diff --git a/presto-main/src/test/java/io/prestosql/operator/TestHashJoinOperator.java b/presto-main/src/test/java/io/prestosql/operator/TestHashJoinOperator.java index 6deafaeb3..a71e1d695 100644 --- a/presto-main/src/test/java/io/prestosql/operator/TestHashJoinOperator.java +++ b/presto-main/src/test/java/io/prestosql/operator/TestHashJoinOperator.java @@ -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; } diff --git a/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java b/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java index 44600d757..0575597b6 100644 --- a/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java +++ b/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java @@ -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 input1 = rowPagesBuilder(VARCHAR, BIGINT) + .row("a", 1L) + .row("b", 2L) + .pageBreak() + .row("b", 3L) + .row("a", 4L) + .build(); + + List 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 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); + } } diff --git a/presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java b/presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java index eacc0a2f7..bb78250d7 100644 --- a/presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java +++ b/presto-main/src/test/java/io/prestosql/operator/TestWindowOperator.java @@ -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 outputPages = ImmutableList.builder(); + + List 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 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 outputPages = ImmutableList.builder(); + + List 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 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); + } }