io.prestosql.operator.unnest.UnnestBlockBuilder Maven / Gradle / Ivy
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.prestosql.operator.unnest;
import com.google.common.annotations.VisibleForTesting;
import io.prestosql.spi.block.Block;
import io.prestosql.spi.block.BlockBuilder;
import io.prestosql.spi.block.DictionaryBlock;
import io.prestosql.spi.block.PageBuilderStatus;
import io.prestosql.spi.type.Type;
import java.util.Arrays;
import static com.google.common.base.Preconditions.checkElementIndex;
import static com.google.common.base.Preconditions.checkPositionIndexes;
import static com.google.common.base.Preconditions.checkState;
import static io.prestosql.operator.unnest.UnnestOperatorBlockUtil.calculateNewArraySize;
import static java.util.Objects.requireNonNull;
/**
* This class manages the low level details of building unnested blocks with a goal of minimizing data copying
*/
class UnnestBlockBuilder
{
private final Type type;
private Block source;
// checks for existence of null element in the source when required
private final NullElementFinder nullFinder = new NullElementFinder();
// flag indicating whether we are using copied block or a dictionary block
private boolean usingCopiedBlock;
// State for output dictionary block
private int[] ids;
private int positionCount;
// State for output copied block
private BlockBuilder outputBlockBuilder;
private PageBuilderStatus pageBuilderStatus;
public UnnestBlockBuilder(Type type)
{
this.type = requireNonNull(type, "type is null");
}
/**
* Replaces input source block with {@code block}. The old data structures for output have to be
* reset as well, because they are based on the source.
*/
public void resetInputBlock(Block block)
{
this.source = requireNonNull(block, "block is null");
// output and null-check have to be cleared because the source has changed
clearCurrentOutput();
nullFinder.resetCheck();
}
/**
* Resets all data structures used for producing current output block, except for the source block.
*/
@VisibleForTesting
void clearCurrentOutput()
{
// clear dictionary block state
ids = new int[0];
positionCount = 0;
usingCopiedBlock = false;
// clear copied block state
outputBlockBuilder = null;
pageBuilderStatus = null;
}
/**
* Prepares for a new output block.
* The caller has to ensure that the source is not null when this method is invoked.
*/
public void startNewOutput(PageBuilderStatus pageBuilderStatus, int expectedEntries)
{
checkState(source != null, "source is null");
this.pageBuilderStatus = pageBuilderStatus;
// Prepare for a new dictionary block
this.ids = new int[expectedEntries];
this.positionCount = 0;
this.usingCopiedBlock = false;
}
public Block buildOutputAndFlush()
{
Block outputBlock;
if (usingCopiedBlock) {
outputBlock = outputBlockBuilder.build();
}
else {
outputBlock = new DictionaryBlock(positionCount, source, ids);
}
// Flush stored state, so that ids cannot be modified after the dictionary has been constructed
clearCurrentOutput();
return outputBlock;
}
/**
* Append element at position {@code index} in the source block to output block
*/
public void appendElement(int index)
{
checkState(source != null, "source is null");
checkElementIndex(index, source.getPositionCount());
if (usingCopiedBlock) {
type.appendTo(source, index, outputBlockBuilder);
}
else {
appendId(index);
}
}
/**
* Append range of elements starting at {@code startPosition} and length {@code length}
* from the source block to output block.
*
* Purpose of this method is to avoid repeated range checks for every invocation of {@link #appendElement}.
*/
public void appendRange(int startPosition, int length)
{
// check range validity
checkState(source != null, "source is null");
checkPositionIndexes(startPosition, startPosition + length, source.getPositionCount());
if (usingCopiedBlock) {
for (int i = 0; i < length; i++) {
type.appendTo(source, startPosition + i, outputBlockBuilder);
}
}
else {
for (int i = 0; i < length; i++) {
appendId(startPosition + i);
}
}
}
public void appendNull()
{
if (usingCopiedBlock) {
outputBlockBuilder.appendNull();
}
else {
int nullIndex = nullFinder.getNullElementIndex();
if (nullIndex == -1) {
startUsingCopiedBlock();
appendNull();
}
else {
appendId(nullIndex);
}
}
}
/**
* This method is invoked when transition from dictionary block to copied block happens.
* - outputBlockBuilder is created and existing elements as indicated by ids[] are copied over.
* - ids[] and positionCount are reset, as they are not used anymore.
*/
private void startUsingCopiedBlock()
{
requireNonNull(pageBuilderStatus, "pageBuilderStatus is null");
// initialize state for copied block
outputBlockBuilder = type.createBlockBuilder(pageBuilderStatus.createBlockBuilderStatus(), ids.length);
// Transfer elements from ids[] to output page builder
for (int i = 0; i < positionCount; i++) {
type.appendTo(source, ids[i], outputBlockBuilder);
}
// reset state for dictionary block
ids = new int[0];
positionCount = 0;
usingCopiedBlock = true;
}
/**
* This method appends {@code sourcePosition} to ids array.
*
* The caller of this method has to ensure the following:
* - {@code sourcePosition} is a valid index in the source
* - {@code usingCopiedBlock} is false
*/
private void appendId(int sourcePosition)
{
if (positionCount == ids.length) {
// grow capacity
int newSize = calculateNewArraySize(ids.length);
ids = Arrays.copyOf(ids, newSize);
}
ids[positionCount] = sourcePosition;
positionCount++;
}
/**
* Get number of positions in the output block being prepared
*/
public int getPositionCount()
{
if (usingCopiedBlock) {
return outputBlockBuilder.getPositionCount();
}
return positionCount;
}
/**
* This class checks for the presence of a non-null element in {@code source}, and stores its position.
* The result is cached with the first invocation of {@link #getNullElementIndex} after the reset. The
* cache can be invalidated by invoking {@link #resetCheck}.
*/
private class NullElementFinder
{
private boolean checkedForNull;
private int nullElementPosition = -1;
void resetCheck()
{
this.checkedForNull = false;
this.nullElementPosition = -1;
}
private int getNullElementIndex()
{
if (checkedForNull) {
return nullElementPosition;
}
checkForNull();
return nullElementPosition;
}
private void checkForNull()
{
nullElementPosition = -1;
for (int i = 0; i < source.getPositionCount(); i++) {
if (source.isNull(i)) {
nullElementPosition = i;
break;
}
}
checkedForNull = true;
}
}
}