Source:
<code>public void partitionPage(Page page) { int positionCount = page.getPositionCount(); int initialPositionCountForEachBuffer = min(positionCount, (positionCount / partitionFunction.getPartitionCount() + 1) * 2); for (int i = 0; i < partitionBuffers.length; i++) { partitionBuffers[i].resetPositions(initialPositionCountForEachBuffer); } Block nullBlock = nullChannel.isPresent() ? page.getBlock(nullChannel.getAsInt()) : null; Page partitionFunctionArgs = getPartitionFunctionArguments(page); for (int position = 0; position < positionCount; position++) { boolean shouldReplicate = (replicatesAnyRow && !hasAnyRowBeenReplicated) || nullBlock != null && nullBlock.isNull(position); if (shouldReplicate) { for (int i = 0; i < partitionBuffers.length; i++) { partitionBuffers[i].addPosition(position); } hasAnyRowBeenReplicated = true; } else { int partition = partitionFunction.getPartition(partitionFunctionArgs, position); partitionBuffers[partition].addPosition(position); } } decodedBlocksRetainedSize = 0; for (int i = 0; i < decodedBlocks.length; i++) { decodedBlocks[i] = decodeBlock(flattener, blockLeaseCloser, page.getBlock(i)); decodedBlocksRetainedSize += decodedBlocks[i].getRetainedSizeInBytes(); } for (int i = 0; i < partitionBuffers.length; i++) { <START> partitionBuffers[i].appendData(decodedBlocks, decodedBlocksRetainedSize, fixedWidthRowSize, variableWidthChannels, outputBuffer); <END> } try { blockLeaseCloser.close(); } catch (IOException e) { throw new UncheckedIOException(e); } }</code><technical_language>@yingsu00 I moving page.getRetainedSizeInBytes() estimated flattened size is right direction. allows for cleaner design. DecodedBlockNode needs original for retainedSizeInBytes is in memory tracking. Hence, I suggest introduce a new variable store estimate of flattened size of block</technical_language>

Target:
public void partitionPage(Page page) { int positionCount = page.getPositionCount(); int initialPositionCountForEachBuffer = min(positionCount, (positionCount / partitionFunction.getPartitionCount() + 1) * 2); for (int i = 0; i < partitionBuffers.length; i++) { partitionBuffers[i].resetPositions(initialPositionCountForEachBuffer); } Block nullBlock = nullChannel.isPresent() ? page.getBlock(nullChannel.getAsInt()) : null; Page partitionFunctionArgs = getPartitionFunctionArguments(page); for (int position = 0; position < positionCount; position++) { boolean shouldReplicate = (replicatesAnyRow && !hasAnyRowBeenReplicated) || nullBlock != null && nullBlock.isNull(position); if (shouldReplicate) { for (int i = 0; i < partitionBuffers.length; i++) { partitionBuffers[i].addPosition(position); } hasAnyRowBeenReplicated = true; } else { int partition = partitionFunction.getPartition(partitionFunctionArgs, position); partitionBuffers[partition].addPosition(position); } } long estimatedSerializedPageSize = 0; for (int i = 0; i < decodedBlocks.length; i++) { decodedBlocks[i] = decodeBlock(flattener, blockLeaseCloser, page.getBlock(i)); estimatedSerializedPageSize += decodedBlocks[i].getEstimatedSerializedSizeInBytes(); } for (int i = 0; i < partitionBuffers.length; i++) { partitionBuffers[i].appendData(decodedBlocks, estimatedSerializedPageSize, fixedWidthRowSize, variableWidthChannels, outputBuffer); } try { blockLeaseCloser.close(); } catch (IOException e) { throw new UncheckedIOException(e); } }