Examples with PartitioningSpillResult io.prestosql.spiller.PartitioningSpiller.PartitioningSpillResult used on opensource projects

Example 1 with PartitioningSpillResult

use of io.prestosql.spiller.PartitioningSpiller.PartitioningSpillResult in project hetu-core by openlookeng.

the class TestGenericPartitioningSpiller method testCloseDuringReading.

@Test
public void testCloseDuringReading() throws Exception {
    Iterator<Page> readingInProgress;
    try (PartitioningSpiller spiller = factory.create(TYPES, new ModuloPartitionFunction(0, 4), mockSpillContext(), mockMemoryContext(scheduledExecutor))) {
        Page page = SequencePageBuilder.createSequencePage(TYPES, 10, FIRST_PARTITION_START, 5, 10, 15);
        PartitioningSpillResult spillResult = spiller.partitionAndSpill(page, partition -> true);
        assertEquals(spillResult.getRetained().getPositionCount(), 0);
        getFutureValue(spillResult.getSpillingFuture());
        // We get the iterator but we do not exhaust it, so that close happens during reading
        readingInProgress = spiller.getSpilledPages(0);
    }
    try {
        readingInProgress.hasNext();
        fail("Iterator.hasNext() should fail since underlying resources are closed");
    } catch (UncheckedIOException ignored) {
    // expected
    // could be ignored
    }
}

Example 2 with PartitioningSpillResult

use of io.prestosql.spiller.PartitioningSpiller.PartitioningSpillResult in project hetu-core by openlookeng.

the class TestGenericPartitioningSpiller method testWriteManyPartitions.

@Test
public void testWriteManyPartitions() throws Exception {
    List<Type> types = ImmutableList.of(BIGINT);
    int partitionCount = 4;
    AggregatedMemoryContext memoryContext = mockMemoryContext(scheduledExecutor);
    try (GenericPartitioningSpiller spiller = (GenericPartitioningSpiller) factory.create(types, new ModuloPartitionFunction(0, partitionCount), mockSpillContext(), memoryContext)) {
        for (int i = 0; i < 50_000; i++) {
            Page page = SequencePageBuilder.createSequencePage(types, partitionCount, 0);
            PartitioningSpillResult spillResult = spiller.partitionAndSpill(page, partition -> true);
            assertEquals(spillResult.getRetained().getPositionCount(), 0);
            getFutureValue(spillResult.getSpillingFuture());
            getFutureValue(spiller.flush());
        }
    }
    assertEquals(memoryContext.getBytes(), 0, "Reserved bytes should be zeroed after spiller is closed");
}

Example 3 with PartitioningSpillResult

use of io.prestosql.spiller.PartitioningSpiller.PartitioningSpillResult in project hetu-core by openlookeng.

the class TestGenericPartitioningSpiller method testFileSpiller.

@Test
public void testFileSpiller() throws Exception {
    try (PartitioningSpiller spiller = factory.create(TYPES, new FourFixedPartitionsPartitionFunction(0), mockSpillContext(), mockMemoryContext(scheduledExecutor))) {
        RowPagesBuilder builder = RowPagesBuilder.rowPagesBuilder(TYPES);
        builder.addSequencePage(10, SECOND_PARTITION_START, 5, 10, 15);
        builder.addSequencePage(10, FIRST_PARTITION_START, -5, 0, 5);
        List<Page> firstSpill = builder.build();
        builder = RowPagesBuilder.rowPagesBuilder(TYPES);
        builder.addSequencePage(10, THIRD_PARTITION_START, 15, 20, 25);
        builder.addSequencePage(10, FOURTH_PARTITION_START, 25, 30, 35);
        List<Page> secondSpill = builder.build();
        IntPredicate spillPartitionMask = ImmutableSet.of(1, 2)::contains;
        PartitioningSpillResult result = spiller.partitionAndSpill(firstSpill.get(0), spillPartitionMask);
        result.getSpillingFuture().get();
        assertEquals(result.getRetained().getPositionCount(), 0);
        result = spiller.partitionAndSpill(firstSpill.get(1), spillPartitionMask);
        result.getSpillingFuture().get();
        assertEquals(result.getRetained().getPositionCount(), 10);
        result = spiller.partitionAndSpill(secondSpill.get(0), spillPartitionMask);
        result.getSpillingFuture().get();
        assertEquals(result.getRetained().getPositionCount(), 0);
        result = spiller.partitionAndSpill(secondSpill.get(1), spillPartitionMask);
        result.getSpillingFuture().get();
        assertEquals(result.getRetained().getPositionCount(), 10);
        builder = RowPagesBuilder.rowPagesBuilder(TYPES);
        builder.addSequencePage(10, SECOND_PARTITION_START, 5, 10, 15);
        List<Page> secondPartition = builder.build();
        builder = RowPagesBuilder.rowPagesBuilder(TYPES);
        builder.addSequencePage(10, THIRD_PARTITION_START, 15, 20, 25);
        List<Page> thirdPartition = builder.build();
        assertSpilledPages(TYPES, spiller, ImmutableList.of(ImmutableList.of(), secondPartition, thirdPartition, ImmutableList.of()));
    }
}

Example 4 with PartitioningSpillResult

use of io.prestosql.spiller.PartitioningSpiller.PartitioningSpillResult in project hetu-core by openlookeng.

the class LookupJoinOperator method spillAndMaskSpilledPositions.

private Page spillAndMaskSpilledPositions(Page page, IntPredicate spillMask, BiPredicate<Integer, Long> spillMatcher) {
    checkState(spillInProgress.isDone(), "Previous spill still in progress");
    checkSuccess(spillInProgress, "spilling failed");
    if (!spiller.isPresent()) {
        spiller = Optional.of(partitioningSpillerFactory.create(probeTypes, getPartitionGenerator(), operatorContext.getSpillContext().newLocalSpillContext(), operatorContext.newAggregateSystemMemoryContext(), hashGenerator::hashPosition));
    }
    PartitioningSpillResult result = spiller.get().partitionAndSpill(page, spillMask, spillMatcher);
    spillInProgress = result.getSpillingFuture();
    return result.getRetained();
}