Examples with LocalMemoryContext com.facebook.presto.memory.context.LocalMemoryContext used on opensource projects

Example 31 with LocalMemoryContext

use of com.facebook.presto.memory.context.LocalMemoryContext in project presto by prestodb.

the class TestFileSingleStreamSpiller method assertSpill.

private void assertSpill(boolean compression, boolean encryption) throws Exception {
    File spillPath = new File(tempDirectory, UUID.randomUUID().toString());
    FileSingleStreamSpillerFactory spillerFactory = new FileSingleStreamSpillerFactory(// executor won't be closed, because we don't call destroy() on the spiller factory
    executor, new BlockEncodingManager(), new SpillerStats(), ImmutableList.of(spillPath.toPath()), 1.0, compression, encryption);
    LocalMemoryContext memoryContext = newSimpleAggregatedMemoryContext().newLocalMemoryContext("test");
    SingleStreamSpiller singleStreamSpiller = spillerFactory.create(TYPES, new TestingSpillContext(), memoryContext);
    assertTrue(singleStreamSpiller instanceof FileSingleStreamSpiller);
    FileSingleStreamSpiller spiller = (FileSingleStreamSpiller) singleStreamSpiller;
    Page page = buildPage();
    // The spillers will reserve memory in their constructors
    assertEquals(memoryContext.getBytes(), 4096);
    spiller.spill(page).get();
    spiller.spill(Iterators.forArray(page, page, page)).get();
    assertEquals(listFiles(spillPath.toPath()).size(), 1);
    // Assert the spill codec flags match the expected configuration
    try (InputStream is = newInputStream(listFiles(spillPath.toPath()).get(0))) {
        Iterator<SerializedPage> serializedPages = PagesSerdeUtil.readSerializedPages(new InputStreamSliceInput(is));
        assertTrue(serializedPages.hasNext(), "at least one page should be successfully read back");
        byte markers = serializedPages.next().getPageCodecMarkers();
        assertEquals(PageCodecMarker.COMPRESSED.isSet(markers), compression);
        assertEquals(PageCodecMarker.ENCRYPTED.isSet(markers), encryption);
    }
    // The spillers release their memory reservations when they are closed, therefore at this point
    // they will have non-zero memory reservation.
    // assertEquals(memoryContext.getBytes(), 0);
    Iterator<Page> spilledPagesIterator = spiller.getSpilledPages();
    assertEquals(memoryContext.getBytes(), FileSingleStreamSpiller.BUFFER_SIZE);
    ImmutableList<Page> spilledPages = ImmutableList.copyOf(spilledPagesIterator);
    // The spillers release their memory reservations when they are closed, therefore at this point
    // they will have non-zero memory reservation.
    // assertEquals(memoryContext.getBytes(), 0);
    assertEquals(4, spilledPages.size());
    for (int i = 0; i < 4; ++i) {
        PageAssertions.assertPageEquals(TYPES, page, spilledPages.get(i));
    }
    spiller.close();
    assertEquals(listFiles(spillPath.toPath()).size(), 0);
    assertEquals(memoryContext.getBytes(), 0);
}