Examples with GroupedAccumulatorState io.trino.spi.function.GroupedAccumulatorState used on opensource projects

Example 1 with GroupedAccumulatorState

use of io.trino.spi.function.GroupedAccumulatorState in project trino by trinodb.

the class TestStateCompiler method testComplexStateEstimatedSize.

@Test(invocationCount = 100, successPercentage = 90)
public void testComplexStateEstimatedSize() {
    Map<String, Type> fieldMap = ImmutableMap.of("Block", new ArrayType(BIGINT), "AnotherBlock", mapType(BIGINT, VARCHAR));
    AccumulatorStateFactory<TestComplexState> factory = StateCompiler.generateStateFactory(TestComplexState.class, fieldMap);
    TestComplexState groupedState = factory.createGroupedState();
    long initialRetainedSize = getComplexStateRetainedSize(groupedState);
    assertEquals(groupedState.getEstimatedSize(), initialRetainedSize);
    // BlockBigArray or SliceBigArray has an internal map that can grow in size when getting more blocks
    // need to handle the map overhead separately
    initialRetainedSize -= getReferenceCountMapOverhead(groupedState);
    for (int i = 0; i < 1000; i++) {
        long retainedSize = 0;
        ((GroupedAccumulatorState) groupedState).setGroupId(i);
        groupedState.setBoolean(true);
        groupedState.setLong(1);
        groupedState.setDouble(2.0);
        groupedState.setByte((byte) 3);
        groupedState.setInt(4);
        Slice slice = utf8Slice("test");
        retainedSize += slice.getRetainedSize();
        groupedState.setSlice(slice);
        slice = wrappedDoubleArray(1.0, 2.0, 3.0);
        retainedSize += slice.getRetainedSize();
        groupedState.setAnotherSlice(slice);
        groupedState.setYetAnotherSlice(null);
        Block array = createLongsBlock(45);
        retainedSize += array.getRetainedSizeInBytes();
        groupedState.setBlock(array);
        BlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1);
        BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry();
        BIGINT.writeLong(singleMapBlockWriter, 123L);
        VARCHAR.writeSlice(singleMapBlockWriter, utf8Slice("testBlock"));
        mapBlockBuilder.closeEntry();
        Block map = mapBlockBuilder.build();
        retainedSize += map.getRetainedSizeInBytes();
        groupedState.setAnotherBlock(map);
        assertEquals(groupedState.getEstimatedSize(), initialRetainedSize + retainedSize * (i + 1) + getReferenceCountMapOverhead(groupedState));
    }
    for (int i = 0; i < 1000; i++) {
        long retainedSize = 0;
        ((GroupedAccumulatorState) groupedState).setGroupId(i);
        groupedState.setBoolean(true);
        groupedState.setLong(1);
        groupedState.setDouble(2.0);
        groupedState.setByte((byte) 3);
        groupedState.setInt(4);
        Slice slice = utf8Slice("test");
        retainedSize += slice.getRetainedSize();
        groupedState.setSlice(slice);
        slice = wrappedDoubleArray(1.0, 2.0, 3.0);
        retainedSize += slice.getRetainedSize();
        groupedState.setAnotherSlice(slice);
        groupedState.setYetAnotherSlice(null);
        Block array = createLongsBlock(45);
        retainedSize += array.getRetainedSizeInBytes();
        groupedState.setBlock(array);
        BlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1);
        BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry();
        BIGINT.writeLong(singleMapBlockWriter, 123L);
        VARCHAR.writeSlice(singleMapBlockWriter, utf8Slice("testBlock"));
        mapBlockBuilder.closeEntry();
        Block map = mapBlockBuilder.build();
        retainedSize += map.getRetainedSizeInBytes();
        groupedState.setAnotherBlock(map);
        assertEquals(groupedState.getEstimatedSize(), initialRetainedSize + retainedSize * 1000 + getReferenceCountMapOverhead(groupedState));
    }
}