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 cefe7fcc3..79e21ac18 100644 --- a/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java +++ b/presto-main/src/main/java/io/prestosql/operator/OrderByOperator.java @@ -31,6 +31,7 @@ import io.prestosql.spi.type.Type; import io.prestosql.spiller.Spiller; import io.prestosql.spiller.SpillerFactory; import io.prestosql.sql.gen.OrderingCompiler; +import org.apache.commons.lang3.tuple.Pair; import java.io.Serializable; import java.nio.file.Path; @@ -413,12 +414,14 @@ public class OrderByOperator } primarySpillRunning = true; pageIndex.sort(sortChannels, sortOrder); - spillInProgress = spiller.get().spill(pageIndex.getSortedPages()); + Pair, Runnable> spillState = spiller.get().spillUnCommit(pageIndex.getSortedPages()); + spillInProgress = spillState.getLeft(); LOG.debug("spilling to disk initiated by Order by operator using primary spiller"); finishMemoryRevoke = () -> { pageIndex.clear(); updateMemoryUsage(true); primarySpillRunning = false; + spillState.getRight().run(); }; return spillInProgress; } @@ -523,12 +526,14 @@ public class OrderByOperator verify(spiller.isPresent(), "spiller not present"); secondarySpillRunning = true; secondaryPageIndex.sort(sortChannels, sortOrder); - spill2InProgress = spiller.get().spill(secondaryPageIndex.getSortedPages()); + Pair, Runnable> spillState = spiller.get().spillUnCommit(secondaryPageIndex.getSortedPages()); + spill2InProgress = spillState.getLeft(); LOG.debug("spilling to disk initiated by Order by operator using secondary spiller"); finishMemoryRevoke2 = () -> { secondaryPageIndex.clear(); updateMemoryUsage(false); secondarySpillRunning = false; + spillState.getRight().run(); }; } diff --git a/presto-main/src/main/java/io/prestosql/spiller/GenericSpiller.java b/presto-main/src/main/java/io/prestosql/spiller/GenericSpiller.java index fd4c23bf3..982766848 100644 --- a/presto-main/src/main/java/io/prestosql/spiller/GenericSpiller.java +++ b/presto-main/src/main/java/io/prestosql/spiller/GenericSpiller.java @@ -22,6 +22,8 @@ import io.prestosql.spi.Page; import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider; import io.prestosql.spi.snapshot.RestorableConfig; import io.prestosql.spi.type.Type; +import org.apache.commons.lang3.tuple.ImmutablePair; +import org.apache.commons.lang3.tuple.Pair; import javax.annotation.concurrent.NotThreadSafe; @@ -31,14 +33,15 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Iterator; import java.util.List; -import java.util.stream.Collectors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.IntStream; import static com.google.common.base.Preconditions.checkState; import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toList; @NotThreadSafe -@RestorableConfig(uncapturedFields = {"types", "spillContext", "aggregatedMemoryContext", "singleStreamSpillerFactory", "closer", "previousSpill"}) +@RestorableConfig(uncapturedFields = {"types", "spillContext", "aggregatedMemoryContext", "singleStreamSpillerFactory", "closer", "previousSpill", "spillCommitted"}) public class GenericSpiller implements Spiller { @@ -49,6 +52,7 @@ public class GenericSpiller private final Closer closer = Closer.create(); private ListenableFuture previousSpill = Futures.immediateFuture(null); private final List singleStreamSpillers = new ArrayList<>(); + private final List spillCommitted = new ArrayList<>(); public GenericSpiller( List types, @@ -68,10 +72,28 @@ public class GenericSpiller SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName())); closer.register(singleStreamSpiller); singleStreamSpillers.add(singleStreamSpiller); + spillCommitted.add(new AtomicBoolean(true)); previousSpill = singleStreamSpiller.spill(pageIterator); return previousSpill; } + /** + * Initiate spilling of pages stream. Returns completed future once spilling has finished with commit function. + * + * @param pageIterator + */ + @Override + public Pair, Runnable> spillUnCommit(Iterator pageIterator) + { + SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName())); + closer.register(singleStreamSpiller); + singleStreamSpillers.add(singleStreamSpiller); + AtomicBoolean isCommitted = new AtomicBoolean(false); + spillCommitted.add(isCommitted); + previousSpill = singleStreamSpiller.spill(pageIterator); + return ImmutablePair.of(previousSpill, () -> isCommitted.set(true)); + } + @Override public List> getSpills() { @@ -110,15 +132,18 @@ public class GenericSpiller @Override public List getSpilledFilePaths() { - return singleStreamSpillers.stream().map(s -> s.getFile()).collect(Collectors.toList()); + return IntStream.range(0, spillCommitted.size()).filter(i -> spillCommitted.get(i).get()).mapToObj(o -> singleStreamSpillers.get(o).getFile()).collect(toList()); } @Override public Object capture(BlockEncodingSerdeProvider serdeProvider) { GenericSpillerState myState = new GenericSpillerState(); - for (SingleStreamSpiller s : singleStreamSpillers) { - myState.singleStreamSpillers.add(s.capture(serdeProvider)); + for (int i = 0; i < singleStreamSpillers.size(); i++) { + if (spillCommitted.get(i).get()) { + SingleStreamSpiller s = singleStreamSpillers.get(i); + myState.singleStreamSpillers.add(s.capture(serdeProvider)); + } } return myState; } @@ -131,6 +156,7 @@ public class GenericSpiller SingleStreamSpiller singleStreamSpiller = singleStreamSpillerFactory.create(types, spillContext, aggregatedMemoryContext.newLocalMemoryContext(GenericSpiller.class.getSimpleName())); singleStreamSpiller.restore(s, serdeProvider); this.singleStreamSpillers.add(singleStreamSpiller); + this.spillCommitted.add(new AtomicBoolean(true)); this.closer.register(singleStreamSpiller); } } diff --git a/presto-main/src/main/java/io/prestosql/spiller/SingleStreamSpiller.java b/presto-main/src/main/java/io/prestosql/spiller/SingleStreamSpiller.java index 99f201f70..e4fec61f5 100644 --- a/presto-main/src/main/java/io/prestosql/spiller/SingleStreamSpiller.java +++ b/presto-main/src/main/java/io/prestosql/spiller/SingleStreamSpiller.java @@ -69,4 +69,14 @@ public interface SingleStreamSpiller void deleteFile(); Path getFile(); + + default void commit() + { + /* do nothing */ + } + + default boolean isCommitted() + { + return true; + } } diff --git a/presto-main/src/main/java/io/prestosql/spiller/Spiller.java b/presto-main/src/main/java/io/prestosql/spiller/Spiller.java index 240eac168..450a8eda6 100644 --- a/presto-main/src/main/java/io/prestosql/spiller/Spiller.java +++ b/presto-main/src/main/java/io/prestosql/spiller/Spiller.java @@ -17,6 +17,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.ListenableFuture; import io.prestosql.spi.Page; import io.prestosql.spi.snapshot.Restorable; +import org.apache.commons.lang3.tuple.ImmutablePair; +import org.apache.commons.lang3.tuple.Pair; import java.io.Closeable; import java.nio.file.Path; @@ -31,6 +33,14 @@ public interface Spiller */ ListenableFuture spill(Iterator pageIterator); + /** + * Initiate spilling of pages stream. Returns completed future once spilling has finished with commit function. + */ + default Pair, Runnable> spillUnCommit(Iterator pageIterator) + { + return ImmutablePair.of(spill(pageIterator), () -> {}); + } + /** * Returns list of previously spilled Pages streams. */ diff --git a/presto-main/src/test/java/io/prestosql/operator/DummySpillerFactory.java b/presto-main/src/test/java/io/prestosql/operator/DummySpillerFactory.java index 77033b685..4c0ee855e 100644 --- a/presto-main/src/test/java/io/prestosql/operator/DummySpillerFactory.java +++ b/presto-main/src/test/java/io/prestosql/operator/DummySpillerFactory.java @@ -17,14 +17,19 @@ import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.ListenableFuture; import io.prestosql.memory.context.AggregatedMemoryContext; import io.prestosql.spi.Page; +import io.prestosql.spi.snapshot.BlockEncodingSerdeProvider; import io.prestosql.spi.snapshot.RestorableConfig; import io.prestosql.spi.type.Type; import io.prestosql.spiller.Spiller; import io.prestosql.spiller.SpillerFactory; +import org.apache.commons.lang3.tuple.ImmutablePair; +import org.apache.commons.lang3.tuple.Pair; +import java.io.Serializable; import java.util.ArrayList; import java.util.Iterator; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.util.concurrent.Futures.immediateFuture; @@ -43,15 +48,27 @@ public class DummySpillerFactory private final RestorableConfig restorableConfig = null; private final List> spills = new ArrayList<>(); + private final List spillCommitted = new ArrayList<>(); @Override public ListenableFuture spill(Iterator pageIterator) { spillsCount++; spills.add(ImmutableList.copyOf(pageIterator)); + spillCommitted.add(new AtomicBoolean(true)); return immediateFuture(null); } + @Override + public Pair, Runnable> spillUnCommit(Iterator pageIterator) + { + spillsCount++; + spills.add(ImmutableList.copyOf(pageIterator)); + AtomicBoolean isCommitted = new AtomicBoolean(false); + spillCommitted.add(isCommitted); + return ImmutablePair.of(immediateFuture(null), () -> isCommitted.set(true)); + } + @Override public List> getSpills() { @@ -65,6 +82,49 @@ public class DummySpillerFactory { spills.clear(); } + + /** + * Capture this object's internal state, so it can be used later to restore to the same state. + * + * @param serdeProvider + * @return An object representing internal state of the current object + */ + @Override + public Object capture(BlockEncodingSerdeProvider serdeProvider) + { + DummySpillerState myState = new DummySpillerState(); + for (int i = 0; i < spills.size(); i++) { + if (spillCommitted.get(i).get()) { + List pages = new ArrayList<>(); + spills.get(i).forEach(pg -> pages.add(pg)); + myState.spills.add(pages); + } + } + return myState; + } + + /** + * Restore this object's internal state according to the snapshot + * + * @param state an object that represents this object's snapshot state + * @param serdeProvider + */ + @Override + public void restore(Object state, BlockEncodingSerdeProvider serdeProvider) + { + DummySpillerState myState = (DummySpillerState) state; + this.spills.clear(); + for (List s : myState.spills) { + this.spills.add(s); + this.spillCommitted.add(new AtomicBoolean(true)); + } + } + + class DummySpillerState + implements Serializable + { + List> spills = new ArrayList<>(); + } }; } 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 37a8b544f..8021f4e92 100644 --- a/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java +++ b/presto-main/src/test/java/io/prestosql/operator/TestOrderByOperator.java @@ -309,6 +309,54 @@ public class TestOrderByOperator assertOperatorEquals(operatorFactory, driverContext, input, expected, revokeMemoryWhenAddingPages); } + @Test + public void testReverseOrderWithSnapshot() + { + List input = rowPagesBuilder(BIGINT, DOUBLE) + .row(1L, 0.1) + .row(2L, 0.2) + .pageBreak() + .row(-1L, -0.1) + .row(4L, 0.4) + .build(); + + OrderByOperatorFactory operatorFactory = new OrderByOperatorFactory( + 0, + new PlanNodeId("test"), + ImmutableList.of(BIGINT, DOUBLE), + ImmutableList.of(0), + 10, + ImmutableList.of(0), + ImmutableList.of(DESC_NULLS_LAST), + new PagesIndex.TestingFactory(false), + true, + Optional.of(spillerFactory), + new OrderingCompiler(), + true); + + DriverContext driverContext = createDriverContext(8, TEST_SESSION); + MaterializedResult expected = resultBuilder(driverContext.getSession(), BIGINT) + .row(4L) + .row(2L) + .row(1L) + .row(-1L) + .build(); + + assertOperatorEqualsWithSimpleSelfStateComparison(operatorFactory, driverContext, input, expected, true, createExpectedMappingRevoke()); + } + + private Map createExpectedMappingRevoke() + { + Map expectedMapping = new HashMap<>(); + expectedMapping.put("operatorContext", 0); + expectedMapping.put("revocableMemoryContext", 1288L); + expectedMapping.put("localUserMemoryContext", 0L); + expectedMapping.put("secondaryMemoryContext", 0L); + expectedMapping.put("secondarySpillRunning", false); + expectedMapping.put("primarySpillRunning", false); + return expectedMapping; + } + @Test(expectedExceptions = ExceededMemoryLimitException.class, expectedExceptionsMessageRegExp = "Query exceeded per-node user memory limit of 10B.*") public void testMemoryLimit() { diff --git a/presto-main/src/test/java/io/prestosql/spiller/TestBinaryFileSpiller.java b/presto-main/src/test/java/io/prestosql/spiller/TestBinaryFileSpiller.java index a5212b929..fa4756abb 100644 --- a/presto-main/src/test/java/io/prestosql/spiller/TestBinaryFileSpiller.java +++ b/presto-main/src/test/java/io/prestosql/spiller/TestBinaryFileSpiller.java @@ -14,6 +14,7 @@ package io.prestosql.spiller; import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.ListenableFuture; import io.hetu.core.transport.execution.buffer.PagesSerde; import io.hetu.core.transport.execution.buffer.PagesSerdeFactory; import io.prestosql.RowPagesBuilder; @@ -23,12 +24,18 @@ import io.prestosql.spi.Page; import io.prestosql.spi.block.BlockBuilder; import io.prestosql.spi.type.Type; import io.prestosql.sql.analyzer.FeaturesConfig; +import io.prestosql.testing.TestingPagesSerdeFactory; +import org.apache.commons.lang3.tuple.Pair; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import java.io.File; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.concurrent.ExecutionException; @@ -44,6 +51,7 @@ import static io.prestosql.spi.type.VarbinaryType.VARBINARY; import static io.prestosql.spi.type.VarcharType.VARCHAR; import static java.lang.Double.doubleToLongBits; import static java.nio.file.Files.createTempDirectory; +import static java.nio.file.StandardCopyOption.REPLACE_EXISTING; import static org.testng.Assert.assertEquals; @Test(singleThreaded = true) @@ -52,6 +60,7 @@ public class TestBinaryFileSpiller private static final List TYPES = ImmutableList.of(BIGINT, VARCHAR, DOUBLE, BIGINT); private final File spillPath = createTempDirectory(getClass().getName()).toFile(); + private final File spillUploadPath = createTempDirectory(getClass().getName()).toFile(); private SpillerStats spillerStats; private FileSingleStreamSpillerFactory singleStreamSpillerFactory; private SpillerFactory factory; @@ -89,6 +98,9 @@ public class TestBinaryFileSpiller { singleStreamSpillerFactory.destroy(); deleteRecursively(spillPath.toPath(), ALLOW_INSECURE); + if (spillUploadPath.exists()) { + deleteRecursively(spillUploadPath.toPath(), ALLOW_INSECURE); + } } @Test @@ -172,4 +184,114 @@ public class TestBinaryFileSpiller spiller.close(); assertEquals(memoryContext.getBytes(), 0); } + + @Test + public void testFileVarbinaryComittableSpiller() + throws Exception + { + List types = ImmutableList.of(BIGINT, DOUBLE, VARBINARY); + + BlockBuilder col1 = BIGINT.createBlockBuilder(null, 1); + BlockBuilder col2 = DOUBLE.createBlockBuilder(null, 1); + BlockBuilder col3 = VARBINARY.createBlockBuilder(null, 1); + + col1.writeLong(42).closeEntry(); + col2.writeLong(doubleToLongBits(43.0)).closeEntry(); + col3.writeLong(doubleToLongBits(43.0)).writeLong(1).closeEntry(); + + Page page = new Page(col1.build(), col2.build(), col3.build()); + + testSpillerUnCommit(types, + ImmutableList.of(page), + ImmutableList.of(page, page), + ImmutableList.of(page, page, page)); + } + + private void testSpillerUnCommit(List types, List... spills) + throws ExecutionException, InterruptedException, IOException + { + long spilledBytesBefore = spillerStats.getTotalSpilledBytes(); + long spilledBytes = 0; + + assertEquals(memoryContext.getBytes(), 0); + List runners = new ArrayList<>(); + PagesSerde serde = TestingPagesSerdeFactory.testingPagesSerde(); + + Spiller spiller = factory.create(TYPES, bytes -> {}, memoryContext); + spilledBytes = doSpill(spiller, spilledBytes, runners, spills, 0); + spillUploadPath.mkdirs(); + + int counter = 1; + int runCount = runners.size(); + List uploadedFile = new ArrayList<>(); + for (counter = 1; counter <= runCount; counter++) { + runners.remove(0).run(); + + assertEquals(spiller.getSpilledFilePaths().size(), counter); + assertEquals(runners.size(), runCount - counter); + + Object snapshot = spiller.capture(serde); + spiller.getSpilledFilePaths().stream().forEach(path -> { + try { + Files.copy(path, Paths.get(spillUploadPath.getPath(), path.getFileName().toString()), REPLACE_EXISTING); + uploadedFile.add(Paths.get(spillUploadPath.getPath(), path.getFileName().toString())); + } + catch (IOException e) { + e.printStackTrace(); + } + }); + + spiller.close(); + spiller = factory.create(TYPES, bytes -> {}, memoryContext); + spiller.restore(snapshot, serde); + uploadedFile.stream().forEach(path -> { + try { + Files.move(path, Paths.get(spillPath.getPath(), path.getFileName().toString()), REPLACE_EXISTING); + } + catch (IOException e) { + e.printStackTrace(); + } + }); + + uploadedFile.clear(); + doSpill(spiller, spilledBytes, runners, spills, counter); + } + //assertEquals(spillerStats.getTotalSpilledBytes() - spilledBytesBefore, spilledBytes); + // At this point, the buffers should still be accounted for in the memory context, because + // the spiller (FileSingleStreamSpiller) doesn't release its memory reservation until it's closed. + //assertEquals(memoryContext.getBytes(), spills.length * FileSingleStreamSpiller.BUFFER_SIZE); + + List> actualSpills = spiller.getSpills(); + assertEquals(actualSpills.size(), spills.length); + + for (int i = 0; i < actualSpills.size(); i++) { + List actualSpill = ImmutableList.copyOf(actualSpills.get(i)); + List expectedSpill = spills[i]; + + assertEquals(actualSpill.size(), expectedSpill.size()); + for (int j = 0; j < actualSpill.size(); j++) { + assertPageEquals(types, actualSpill.get(j), expectedSpill.get(j)); + } + } + spiller.close(); + assertEquals(memoryContext.getBytes(), 0); + } + + private long doSpill(Spiller spiller, long spilledBytes, List runners, List[] spills, int count) throws InterruptedException, ExecutionException + { + runners.clear(); + for (List spill : spills) { + if (count > 0) { + count--; + continue; + } + spilledBytes += spill.stream() + .mapToLong(page -> pagesSerde.serialize(page).getSizeInBytes()) + .sum(); + Pair, Runnable> spillState = spiller.spillUnCommit(spill.iterator()); + runners.add(spillState.getRight()); + spillState.getLeft().get(); + } + return spilledBytes; + } }