@Benchmark
public Page benchmark() {
PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(prestoType));
int count = processor.process(null, inputPage, 0, inputPage.getPositionCount(), pageBuilder);
checkState(count == inputPage.getPositionCount());
return pageBuilder.build();
}
@Override
public ConnectorPageSource createPageSource(
ConnectorTransactionHandle transactionHandle,
ConnectorSession session,
ConnectorSplit split,
List<ColumnHandle> columns) {
InternalTable table = getInternalTable(transactionHandle, session, split, columns);
List<Integer> channels = new ArrayList<>();
for (ColumnHandle column : columns) {
String columnName =
checkType(column, InformationSchemaColumnHandle.class, "column").getColumnName();
int columnIndex = table.getColumnIndex(columnName);
channels.add(columnIndex);
}
ImmutableList.Builder<Page> pages = ImmutableList.builder();
for (Page page : table.getPages()) {
Block[] blocks = new Block[channels.size()];
for (int index = 0; index < blocks.length; index++) {
blocks[index] = page.getBlock(channels.get(index));
}
pages.add(new Page(page.getPositionCount(), blocks));
}
return new FixedPageSource(pages.build());
}
private static Page rearrangePage(Page page, int[] channels) {
Block[] newBlocks = new Block[channels.length];
for (int i = 0; i < channels.length; i++) {
newBlocks[i] = page.getBlock(channels[i]);
}
return new Page(page.getPositionCount(), newBlocks);
}
public ListenableFuture<?> partitionPage(Page page) {
requireNonNull(page, "page is null");
Page partitionFunctionArgs = getPartitionFunctionArguments(page);
for (int position = 0; position < page.getPositionCount(); position++) {
if (nullChannel.isPresent() && page.getBlock(nullChannel.getAsInt()).isNull(position)) {
for (PageBuilder pageBuilder : pageBuilders) {
pageBuilder.declarePosition();
for (int channel = 0; channel < sourceTypes.size(); channel++) {
Type type = sourceTypes.get(channel);
type.appendTo(page.getBlock(channel), position, pageBuilder.getBlockBuilder(channel));
}
}
} else {
int partition = partitionFunction.getPartition(partitionFunctionArgs, position);
PageBuilder pageBuilder = pageBuilders.get(partition);
pageBuilder.declarePosition();
for (int channel = 0; channel < sourceTypes.size(); channel++) {
Type type = sourceTypes.get(channel);
type.appendTo(page.getBlock(channel), position, pageBuilder.getBlockBuilder(channel));
}
}
}
return flush(false);
}
@Override
public Page getOutput() {
Page page = exchangeClient.pollPage();
if (page != null) {
operatorContext.recordGeneratedInput(page.getSizeInBytes(), page.getPositionCount());
}
return page;
}
@Override
public Page getOutput() {
if (!pages.hasNext()) {
return null;
}
Page page = pages.next();
if (page != null) {
operatorContext.recordGeneratedInput(page.getSizeInBytes(), page.getPositionCount());
}
return page;
}
private Page getPartitionFunctionArguments(Page page) {
Block[] blocks = new Block[partitionChannels.size()];
for (int i = 0; i < blocks.length; i++) {
Optional<Block> partitionConstant = partitionConstants.get(i);
if (partitionConstant.isPresent()) {
blocks[i] = new RunLengthEncodedBlock(partitionConstant.get(), page.getPositionCount());
} else {
blocks[i] = page.getBlock(partitionChannels.get(i));
}
}
return new Page(page.getPositionCount(), blocks);
}
private static boolean executeFilterWithNoInputColumns(Operator operator) {
Page page = getAtMostOnePage(operator, ZERO_CHANNEL_PAGE);
boolean value;
if (page != null) {
assertEquals(page.getPositionCount(), 1);
assertEquals(page.getChannelCount(), 0);
value = true;
} else {
value = false;
}
return value;
}
private Object selectSingleValue(Operator operator) {
Page output = getAtMostOnePage(operator, SOURCE_PAGE);
assertNotNull(output);
assertEquals(output.getPositionCount(), 1);
assertEquals(output.getChannelCount(), 1);
Type type = operator.getTypes().get(0);
Block block = output.getBlock(0);
assertEquals(block.getPositionCount(), 1);
return type.getObjectValue(session.toConnectorSession(), block, 0);
}
@Override
public void addInput(Page page) {
requireNonNull(page, "page is null");
checkState(isBlocked().isDone(), "output is already blocked");
if (page.getPositionCount() == 0) {
return;
}
page = pagePreprocessor.apply(page);
blocked = partitionFunction.partitionPage(page);
operatorContext.recordGeneratedOutput(page.getSizeInBytes(), page.getPositionCount());
}
private static boolean executeFilter(Operator operator) {
Page page = getAtMostOnePage(operator, SOURCE_PAGE);
boolean value;
if (page != null) {
assertEquals(page.getPositionCount(), 1);
assertEquals(page.getChannelCount(), 1);
assertTrue(operator.getTypes().get(0).getBoolean(page.getBlock(0), 0));
value = true;
} else {
value = false;
}
return value;
}
/**
* @return the unused section of the page, or null if fully applied. pagesIndex guaranteed to have
* at least one row after this method returns
*/
private Page updatePagesIndex(Page page) {
checkArgument(page.getPositionCount() > 0);
// TODO: Fix pagesHashStrategy to allow specifying channels for comparison, it currently
// requires us to rearrange the right side blocks in consecutive channel order
Page preGroupedPage = rearrangePage(page, preGroupedChannels);
if (pagesIndex.getPositionCount() == 0
|| pagesIndex.positionEqualsRow(
preGroupedPartitionHashStrategy, 0, 0, preGroupedPage.getBlocks())) {
// Find the position where the pre-grouped columns change
int groupEnd = findGroupEnd(preGroupedPage, preGroupedPartitionHashStrategy, 0);
// Add the section of the page that contains values for the current group
pagesIndex.addPage(page.getRegion(0, groupEnd));
if (page.getPositionCount() - groupEnd > 0) {
// Save the remaining page, which may contain multiple partitions
return page.getRegion(groupEnd, page.getPositionCount() - groupEnd);
} else {
// Page fully consumed
return null;
}
} else {
// We had previous results buffered, but the new page starts with new group values
return page;
}
}
@SuppressWarnings("NumericCastThatLosesPrecision")
@Override
public int getBucket(Page page, int position) {
long hash = 0;
for (Block block : page.getBlocks()) {
long value = BIGINT.getLong(block, position);
hash = (hash * 31) + XxHash64.hash(value);
}
int value = (int) (hash & Integer.MAX_VALUE);
return value % bucketCount;
}
public ListenableFuture<?> flush(boolean force) {
// add all full pages to output buffer
List<ListenableFuture<?>> blockedFutures = new ArrayList<>();
for (int partition = 0; partition < pageBuilders.size(); partition++) {
PageBuilder partitionPageBuilder = pageBuilders.get(partition);
if (!partitionPageBuilder.isEmpty() && (force || partitionPageBuilder.isFull())) {
Page pagePartition = partitionPageBuilder.build();
partitionPageBuilder.reset();
blockedFutures.add(outputBuffer.enqueue(partition, pagePartition));
pagesAdded.incrementAndGet();
rowsAdded.addAndGet(pagePartition.getPositionCount());
}
}
ListenableFuture<?> future = Futures.allAsList(blockedFutures);
if (future.isDone()) {
return NOT_BLOCKED;
}
return future;
}
public static Row extractRow(Page page, int position, List<Type> types) {
checkArgument(page.getChannelCount() == types.size(), "channelCount does not match");
checkArgument(
position < page.getPositionCount(),
"Requested position %s from a page with positionCount %s ",
position,
page.getPositionCount());
RowBuilder rowBuilder = new RowBuilder();
for (int channel = 0; channel < page.getChannelCount(); channel++) {
Block block = page.getBlock(channel);
Type type = types.get(channel);
int size;
Object value = getNativeContainerValue(type, block, position);
if (value == null) {
size = SIZE_OF_BYTE;
} else if (type.getJavaType() == boolean.class) {
size = SIZE_OF_BYTE;
} else if (type.getJavaType() == long.class) {
size = SIZE_OF_LONG;
} else if (type.getJavaType() == double.class) {
size = SIZE_OF_DOUBLE;
} else if (type.getJavaType() == Slice.class) {
size = ((Slice) value).length();
} else if (type.getJavaType() == Block.class) {
size = ((Block) value).getSizeInBytes();
} else {
throw new AssertionError("Unimplemented type: " + type);
}
rowBuilder.add(nativeContainerToOrcValue(type, value), size);
}
Row row = rowBuilder.build();
verify(
row.getColumns().size() == types.size(),
"Column count in row: %s Expected column count: %s",
row.getColumns().size(),
types.size());
return row;
}
// Assumes input grouped on relevant pagesHashStrategy columns
private static int findGroupEnd(
Page page, PagesHashStrategy pagesHashStrategy, int startPosition) {
checkArgument(page.getPositionCount() > 0, "Must have at least one position");
checkPositionIndex(startPosition, page.getPositionCount(), "startPosition out of bounds");
// Short circuit if the whole page has the same value
if (pagesHashStrategy.rowEqualsRow(
startPosition, page.getBlocks(), page.getPositionCount() - 1, page.getBlocks())) {
return page.getPositionCount();
}
// TODO: do position binary search
int endPosition = startPosition + 1;
while (endPosition < page.getPositionCount()
&& pagesHashStrategy.rowEqualsRow(
endPosition - 1, page.getBlocks(), endPosition, page.getBlocks())) {
endPosition++;
}
return endPosition;
}
@Override
public void addInput(Page page) {
checkState(state == State.NEEDS_INPUT, "Operator can not take input at this time");
requireNonNull(page, "page is null");
checkState(pendingInput == null, "Operator already has pending input");
if (page.getPositionCount() == 0) {
return;
}
pendingInput = page;
if (processPendingInput()) {
state = State.HAS_OUTPUT;
}
operatorContext.setMemoryReservation(pagesIndex.getEstimatedSize().toBytes());
}
@Override
public void addInput(Page page) {
requireNonNull(page, "page is null");
checkState(!finishing, "Operator is finishing");
checkState(channelSet != null, "Set has not been built yet");
checkState(outputPage == null, "Operator still has pending output");
// create the block builder for the new boolean column
// we know the exact size required for the block
BlockBuilder blockBuilder = BOOLEAN.createFixedSizeBlockBuilder(page.getPositionCount());
Page probeJoinPage = new Page(page.getBlock(probeJoinChannel));
// update hashing strategy to use probe cursor
for (int position = 0; position < page.getPositionCount(); position++) {
if (probeJoinPage.getBlock(0).isNull(position)) {
throw new PrestoException(
NOT_SUPPORTED,
"NULL values are not allowed on the probe side of SemiJoin operator. See the query plan for details.");
} else {
boolean contains = channelSet.contains(position, probeJoinPage);
if (!contains && channelSet.containsNull()) {
blockBuilder.appendNull();
} else {
BOOLEAN.writeBoolean(blockBuilder, contains);
}
}
}
// add the new boolean column to the page
Block[] sourceBlocks = page.getBlocks();
Block[] outputBlocks =
new Block[sourceBlocks.length + 1]; // +1 for the single boolean output channel
System.arraycopy(sourceBlocks, 0, outputBlocks, 0, sourceBlocks.length);
outputBlocks[sourceBlocks.length] = blockBuilder.build();
outputPage = new Page(outputBlocks);
}
private static void assertPageEquals(Page expectedPage, Page actualPage) {
assertEquals(actualPage.getPositionCount(), expectedPage.getPositionCount());
assertEquals(actualPage.getChannelCount(), expectedPage.getChannelCount());
}