Examples with TestingTempStorageManager com.facebook.presto.testing.TestingTempStorageManager used on opensource projects

Example 1 with TestingTempStorageManager

use of com.facebook.presto.testing.TestingTempStorageManager in project presto by prestodb.

the class TestTempStorageSingleStreamSpiller method assertSpill.

private void assertSpill(boolean compression, boolean encryption) throws Exception {
    File spillPath = new File(tempDirectory, UUID.randomUUID().toString());
    TempStorageSingleStreamSpillerFactory spillerFactory = new TempStorageSingleStreamSpillerFactory(new TestingTempStorageManager(spillPath.toString()), // executor won't be closed, because we don't call destroy() on the spiller factory
    executor, new BlockEncodingManager(), new SpillerStats(), compression, encryption, LocalTempStorage.NAME);
    LocalMemoryContext memoryContext = newSimpleAggregatedMemoryContext().newLocalMemoryContext("test");
    SingleStreamSpiller singleStreamSpiller = spillerFactory.create(TYPES, new TestingSpillContext(), memoryContext);
    assertTrue(singleStreamSpiller instanceof TempStorageSingleStreamSpiller);
    TempStorageSingleStreamSpiller spiller = (TempStorageSingleStreamSpiller) singleStreamSpiller;
    Page page = buildPage();
    // The spillers will reserve memory in their constructors
    int retainedSizeForEmptyDataSink = toIntExact(new OutputStreamDataSink(new DynamicSliceOutput(0)).getRetainedSizeInBytes());
    assertEquals(memoryContext.getBytes(), retainedSizeForEmptyDataSink);
    spiller.spill(page).get();
    spiller.spill(Iterators.forArray(page, page, page)).get();
    assertEquals(listFiles(spillPath.toPath()).size(), 1);
    // 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(), retainedSizeForEmptyDataSink);
    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));
    }
    // 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);
    }
    spiller.close();
    assertEquals(listFiles(spillPath.toPath()).size(), 0);
    assertEquals(memoryContext.getBytes(), 0);
}

Example 2 with TestingTempStorageManager

use of com.facebook.presto.testing.TestingTempStorageManager in project presto by prestodb.

the class TestAggregationOperator method testDistinctMaskWithNull.

@Test
public void testDistinctMaskWithNull() {
    AccumulatorFactory distinctFactory = COUNT.bind(ImmutableList.of(0), Optional.of(1), ImmutableList.of(BIGINT, BOOLEAN), ImmutableList.of(), ImmutableList.of(), null, // distinct
    true, new JoinCompiler(MetadataManager.createTestMetadataManager(), new FeaturesConfig()), ImmutableList.of(), false, TEST_SESSION, new TempStorageStandaloneSpillerFactory(new TestingTempStorageManager(), new BlockEncodingManager(), new NodeSpillConfig(), new FeaturesConfig(), new SpillerStats()));
    OperatorFactory operatorFactory = new AggregationOperatorFactory(0, new PlanNodeId("test"), Step.SINGLE, ImmutableList.of(distinctFactory), false);
    DriverContext driverContext = createTaskContext(executor, scheduledExecutor, TEST_SESSION).addPipelineContext(0, true, true, false).addDriverContext();
    ByteArrayBlock trueMaskAllNull = new ByteArrayBlock(4, Optional.of(new boolean[] { true, true, true, true }), /* all positions are null */
    new byte[] { 1, 1, 1, 1 });
    /* non-zero value is true, all masks are true */
    Block trueNullRleMask = new RunLengthEncodedBlock(trueMaskAllNull.getSingleValueBlock(0), 4);
    List<Page> input = ImmutableList.of(new Page(4, createLongsBlock(1, 2, 3, 4), trueMaskAllNull), new Page(4, createLongsBlock(5, 6, 7, 8), trueNullRleMask));
    MaterializedResult expected = resultBuilder(driverContext.getSession(), BIGINT).row(// all rows should be filtered by nulls
    0L).build();
    assertOperatorEquals(operatorFactory, driverContext, input, expected);
}