public static TypedKeyValueHeap deserialize(Block block, Type keyType, Type valueType, BlockComparator blockComparator) { int capacity = toIntExact(BIGINT.getLong(block, 0)); Block keysBlock = new ArrayType(keyType).getObject(block, 1); Block valuesBlock = new ArrayType(valueType).getObject(block, 2); TypedKeyValueHeap heap = new TypedKeyValueHeap(blockComparator, keyType, valueType, capacity); heap.addAll(keysBlock, valuesBlock); return heap; }
public static void output(ArrayType outputType, MinMaxByNState state, BlockBuilder out) { TypedKeyValueHeap heap = state.getTypedKeyValueHeap(); if (heap == null || heap.isEmpty()) { out.appendNull(); return; } Type elementType = outputType.getElementType(); BlockBuilder arrayBlockBuilder = out.beginBlockEntry(); BlockBuilder reversedBlockBuilder = elementType.createBlockBuilder(null, heap.getCapacity()); long startSize = heap.getEstimatedSize(); heap.popAll(reversedBlockBuilder); state.addMemoryUsage(heap.getEstimatedSize() - startSize); for (int i = reversedBlockBuilder.getPositionCount() - 1; i >= 0; i--) { elementType.appendTo(reversedBlockBuilder, i, arrayBlockBuilder); } out.closeEntry(); }
public static void input(BlockComparator comparator, Type valueType, Type keyType, MinMaxByNState state, Block value, Block key, int blockIndex, long n) { TypedKeyValueHeap heap = state.getTypedKeyValueHeap(); if (heap == null) { if (n <= 0) { throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "third argument of max_by/min_by must be a positive integer"); } checkCondition(n <= MAX_NUMBER_OF_VALUES, INVALID_FUNCTION_ARGUMENT, "third argument of max_by/min_by must be less than or equal to %s; found %s", MAX_NUMBER_OF_VALUES, n); heap = new TypedKeyValueHeap(comparator, keyType, valueType, toIntExact(n)); state.setTypedKeyValueHeap(heap); } long startSize = heap.getEstimatedSize(); if (!key.isNull(blockIndex)) { heap.add(key, value, blockIndex); } state.addMemoryUsage(heap.getEstimatedSize() - startSize); }
public static void combine(MinMaxByNState state, MinMaxByNState otherState) { TypedKeyValueHeap otherHeap = otherState.getTypedKeyValueHeap(); if (otherHeap == null) { return; } TypedKeyValueHeap heap = state.getTypedKeyValueHeap(); if (heap == null) { state.setTypedKeyValueHeap(otherHeap); return; } long startSize = heap.getEstimatedSize(); heap.addAll(otherHeap); state.addMemoryUsage(heap.getEstimatedSize() - startSize); }
private static void test(IntStream keyInputStream, Stream<String> valueInputStream, BlockComparator comparator, Iterator<String> outputIterator) { BlockBuilder keysBlockBuilder = BIGINT.createBlockBuilder(null, INPUT_SIZE); BlockBuilder valuesBlockBuilder = VARCHAR.createBlockBuilder(null, INPUT_SIZE); keyInputStream.forEach(x -> BIGINT.writeLong(keysBlockBuilder, x)); valueInputStream.forEach(x -> VARCHAR.writeString(valuesBlockBuilder, x)); TypedKeyValueHeap heap = new TypedKeyValueHeap(comparator, BIGINT, VARCHAR, OUTPUT_SIZE); heap.addAll(keysBlockBuilder, valuesBlockBuilder); BlockBuilder resultBlockBuilder = VARCHAR.createBlockBuilder(null, OUTPUT_SIZE); heap.popAll(resultBlockBuilder); Block resultBlock = resultBlockBuilder.build(); assertEquals(resultBlock.getPositionCount(), OUTPUT_SIZE); for (int i = 0; i < OUTPUT_SIZE; i++) { assertEquals(VARCHAR.getSlice(resultBlock, i).toStringUtf8(), outputIterator.next()); } } }
public void add(Block keyBlock, Block valueBlock, int position) { checkArgument(!keyBlock.isNull(position)); if (positionCount == capacity) { if (keyComparator.compareTo(keyBlockBuilder, heapIndex[0], keyBlock, position) >= 0) { return; // and new element is not larger than heap top: do not add } heapIndex[0] = keyBlockBuilder.getPositionCount(); keyType.appendTo(keyBlock, position, keyBlockBuilder); valueType.appendTo(valueBlock, position, valueBlockBuilder); siftDown(); } else { heapIndex[positionCount] = keyBlockBuilder.getPositionCount(); positionCount++; keyType.appendTo(keyBlock, position, keyBlockBuilder); valueType.appendTo(valueBlock, position, valueBlockBuilder); siftUp(); } compactIfNecessary(); }
public void addAll(TypedKeyValueHeap otherHeap) { addAll(otherHeap.keyBlockBuilder, otherHeap.valueBlockBuilder); }
@Override public long getEstimatedSize() { long estimatedSize = INSTANCE_SIZE; if (typedKeyValueHeap != null) { estimatedSize += typedKeyValueHeap.getEstimatedSize(); } return estimatedSize; }
public void addAll(Block keysBlock, Block valuesBlock) { for (int i = 0; i < keysBlock.getPositionCount(); i++) { add(keysBlock, valuesBlock, i); } }
public MinMaxByNStateSerializer(BlockComparator blockComparator, Type keyType, Type valueType) { this.blockComparator = blockComparator; this.keyType = keyType; this.valueType = valueType; this.serializedType = TypedKeyValueHeap.getSerializedType(keyType, valueType); }
@Override public void deserialize(Block block, int index, MinMaxByNState state) { Block currentBlock = (Block) serializedType.getObject(block, index); state.setTypedKeyValueHeap(TypedKeyValueHeap.deserialize(currentBlock, keyType, valueType, blockComparator)); } }
public void serialize(BlockBuilder out) { BlockBuilder blockBuilder = out.beginBlockEntry(); BIGINT.writeLong(blockBuilder, getCapacity()); BlockBuilder keyElements = blockBuilder.beginBlockEntry(); for (int i = 0; i < positionCount; i++) { keyType.appendTo(keyBlockBuilder, heapIndex[i], keyElements); } blockBuilder.closeEntry(); BlockBuilder valueElements = blockBuilder.beginBlockEntry(); for (int i = 0; i < positionCount; i++) { valueType.appendTo(valueBlockBuilder, heapIndex[i], valueElements); } blockBuilder.closeEntry(); out.closeEntry(); }
private void test(IntStream keyInputStream, Stream<String> valueInputStream, BlockComparator comparator, Iterator<String> outputIterator) { BlockBuilder keysBlockBuilder = BIGINT.createBlockBuilder(new BlockBuilderStatus(), INPUT_SIZE); BlockBuilder valuesBlockBuilder = VARCHAR.createBlockBuilder(new BlockBuilderStatus(), INPUT_SIZE); keyInputStream.forEach(x -> BIGINT.writeLong(keysBlockBuilder, x)); valueInputStream.forEach(x -> VARCHAR.writeString(valuesBlockBuilder, x)); TypedKeyValueHeap heap = new TypedKeyValueHeap(comparator, BIGINT, VARCHAR, OUTPUT_SIZE); heap.addAll(keysBlockBuilder, valuesBlockBuilder); BlockBuilder resultBlockBuilder = VARCHAR.createBlockBuilder(new BlockBuilderStatus(), OUTPUT_SIZE); heap.popAll(resultBlockBuilder); Block resultBlock = resultBlockBuilder.build(); assertEquals(resultBlock.getPositionCount(), OUTPUT_SIZE); for (int i = 0; i < OUTPUT_SIZE; i++) { assertEquals(VARCHAR.getSlice(resultBlock, i).toStringUtf8(), outputIterator.next()); } } }
public static void combine(MinMaxByNState state, MinMaxByNState otherState) { TypedKeyValueHeap otherHeap = otherState.getTypedKeyValueHeap(); if (otherHeap == null) { return; } TypedKeyValueHeap heap = state.getTypedKeyValueHeap(); if (heap == null) { state.setTypedKeyValueHeap(otherHeap); return; } long startSize = heap.getEstimatedSize(); heap.addAll(otherHeap); state.addMemoryUsage(heap.getEstimatedSize() - startSize); }
public void add(Block keyBlock, Block valueBlock, int position) { checkArgument(!keyBlock.isNull(position)); if (positionCount == capacity) { if (keyComparator.compareTo(keyBlockBuilder, heapIndex[0], keyBlock, position) >= 0) { return; // and new element is not larger than heap top: do not add } heapIndex[0] = keyBlockBuilder.getPositionCount(); keyType.appendTo(keyBlock, position, keyBlockBuilder); valueType.appendTo(valueBlock, position, valueBlockBuilder); siftDown(); } else { heapIndex[positionCount] = keyBlockBuilder.getPositionCount(); positionCount++; keyType.appendTo(keyBlock, position, keyBlockBuilder); valueType.appendTo(valueBlock, position, valueBlockBuilder); siftUp(); } compactIfNecessary(); }
public void addAll(TypedKeyValueHeap otherHeap) { addAll(otherHeap.keyBlockBuilder, otherHeap.valueBlockBuilder); }
@Override public void setTypedKeyValueHeap(TypedKeyValueHeap value) { TypedKeyValueHeap previous = getTypedKeyValueHeap(); if (previous != null) { size -= previous.getEstimatedSize(); } heaps.set(getGroupId(), value); size += value.getEstimatedSize(); }
public void addAll(Block keysBlock, Block valuesBlock) { for (int i = 0; i < keysBlock.getPositionCount(); i++) { add(keysBlock, valuesBlock, i); } }
public MinMaxByNStateSerializer(BlockComparator blockComparator, Type keyType, Type valueType) { this.blockComparator = blockComparator; this.keyType = keyType; this.valueType = valueType; this.serializedType = TypedKeyValueHeap.getSerializedType(keyType, valueType); }
@Override public void deserialize(Block block, int index, MinMaxByNState state) { if (block.isNull(index)) { return; } Block currentBlock = (Block) serializedType.getObject(block, index); state.setTypedKeyValueHeap(TypedKeyValueHeap.deserialize(currentBlock, keyType, valueType, blockComparator)); } }