private static PrestoThriftPageResult toThriftPage(Page page, List<Type> columnTypes, @Nullable PrestoThriftId nextToken) { if (page == null) { checkState(nextToken == null, "there must be no more data when page is null"); return new PrestoThriftPageResult(ImmutableList.of(), 0, null); } checkState(page.getChannelCount() == columnTypes.size(), "number of columns in a page doesn't match the one in requested types"); int numberOfColumns = columnTypes.size(); List<PrestoThriftBlock> columnBlocks = new ArrayList<>(numberOfColumns); for (int i = 0; i < numberOfColumns; i++) { columnBlocks.add(fromBlock(page.getBlock(i), columnTypes.get(i))); } return new PrestoThriftPageResult(columnBlocks, page.getPositionCount(), nextToken); }
this.outputColumnTypes = outputColumnTypes.build(); this.keys = fromRecordSet(requireNonNull(keys, "keys is null"));
private static List<List<String>> thriftPageToList(PrestoThriftPageResult page, int begin, int end) { checkArgument(begin <= end, "invalid interval"); if (begin == end) { // empty interval return ImmutableList.of(); } List<PrestoThriftBlock> blocks = page.getColumnBlocks(); List<List<String>> result = new ArrayList<>(blocks.size()); for (PrestoThriftBlock block : blocks) { result.add(blockAsList(block, begin, end)); } return result; }
public static PrestoThriftPageResult fromRecordSet(RecordSet recordSet) { List<Type> types = recordSet.getColumnTypes(); int numberOfColumns = types.size(); int positions = totalRecords(recordSet); if (numberOfColumns == 0) { return new PrestoThriftPageResult(ImmutableList.of(), positions, null); } List<PrestoThriftBlock> thriftBlocks = new ArrayList<>(numberOfColumns); for (int columnIndex = 0; columnIndex < numberOfColumns; columnIndex++) { thriftBlocks.add(fromRecordSetColumn(recordSet, columnIndex, positions)); } return new PrestoThriftPageResult(thriftBlocks, positions, null); }
private Page processBatch(PrestoThriftPageResult rowsBatch) { firstCall = false; nextToken = rowsBatch.getNextToken(); Page page = rowsBatch.toPage(columnTypes); if (page != null) { long pageSize = page.getSizeInBytes(); completedBytes += pageSize; stats.addScanPageSize(pageSize); } else { stats.addScanPageSize(0); } return page; }
@Override public ListenableFuture<PrestoThriftSplitBatch> getIndexSplits(PrestoThriftSchemaTableName schemaTableName, List<String> indexColumnNames, List<String> outputColumnNames, PrestoThriftPageResult keys, PrestoThriftTupleDomain outputConstraint, int maxSplitCount, PrestoThriftNullableToken nextToken) if (keys.getRowCount() == 0) { return immediateFuture(new PrestoThriftSplitBatch(ImmutableList.of(), null)); int[] values = keys.getColumnBlocks().get(0).getIntegerData().getInts(); int begin; int end;
private void testReadWrite(Random random, int records, Function<List<Block>, PrestoThriftPageResult> convert) { // generate columns data List<Block> inputBlocks = new ArrayList<>(columns.size()); for (ColumnDefinition column : columns) { inputBlocks.add(generateColumn(column, random, records)); } // convert column data to thrift ("write step") PrestoThriftPageResult batch = convert.apply(inputBlocks); // convert thrift data to page/blocks ("read step") Page page = batch.toPage(columns.stream().map(ColumnDefinition::getType).collect(toImmutableList())); // compare the result with original input assertNotNull(page); assertEquals(page.getChannelCount(), columns.size()); for (int i = 0; i < columns.size(); i++) { Block actual = page.getBlock(i); Block expected = inputBlocks.get(i); assertBlock(actual, expected, columns.get(i)); } }
int totalKeys = keys.getRowCount(); int partSize = totalKeys / NUMBER_OF_INDEX_SPLITS; List<PrestoThriftSplit> splits = new ArrayList<>(NUMBER_OF_INDEX_SPLITS);
@ThriftConstructor public PrestoThriftPageResult(List<PrestoThriftBlock> columnBlocks, int rowCount, @Nullable PrestoThriftId nextToken) { this.columnBlocks = requireNonNull(columnBlocks, "columnBlocks is null"); checkArgument(rowCount >= 0, "rowCount is negative"); checkAllColumnsAreOfExpectedSize(columnBlocks, rowCount); this.rowCount = rowCount; this.nextToken = nextToken; }
checkState(resultContext != null, "no associated context for the request"); PrestoThriftPageResult pageResult = getFutureValue(resultFuture); Page page = pageResult.toPage(outputColumnTypes); if (page != null) { long pageSize = page.getSizeInBytes(); stats.addIndexPageSize(0); if (pageResult.getNextToken() != null) { sendDataRequest(resultContext, pageResult.getNextToken()); updateSignalAndStatusFutures(); return page;
private void testReadWrite(Random random, int records, Function<List<Block>, PrestoThriftPageResult> convert) { // generate columns data List<Block> inputBlocks = new ArrayList<>(columns.size()); for (ColumnDefinition column : columns) { inputBlocks.add(generateColumn(column, random, records)); } // convert column data to thrift ("write step") PrestoThriftPageResult batch = convert.apply(inputBlocks); // convert thrift data to page/blocks ("read step") Page page = batch.toPage(columns.stream().map(ColumnDefinition::getType).collect(toImmutableList())); // compare the result with original input assertNotNull(page); assertEquals(page.getChannelCount(), columns.size()); for (int i = 0; i < columns.size(); i++) { Block actual = page.getBlock(i); Block expected = inputBlocks.get(i); assertBlock(actual, expected, columns.get(i)); } }
int totalKeys = keys.getRowCount(); int partSize = totalKeys / NUMBER_OF_INDEX_SPLITS; List<PrestoThriftSplit> splits = new ArrayList<>(NUMBER_OF_INDEX_SPLITS);
private void testPageReadWrite(Random random, int records) { testReadWrite(random, records, blocks -> { List<PrestoThriftBlock> columnBlocks = new ArrayList<>(columns.size()); for (int i = 0; i < columns.size(); i++) { columnBlocks.add(fromBlock(blocks.get(i), columns.get(i).getType())); } return new PrestoThriftPageResult(columnBlocks, records, null); }); }
private void testRecordSetReadWrite(Random random, int records) { testReadWrite(random, records, blocks -> { List<Type> types = columns.stream().map(ColumnDefinition::getType).collect(toImmutableList()); PageRecordSet inputRecordSet = new PageRecordSet(types, new Page(blocks.toArray(new Block[blocks.size()]))); return fromRecordSet(inputRecordSet); }); }
private static List<List<String>> thriftPageToList(PrestoThriftPageResult page, int begin, int end) { checkArgument(begin <= end, "invalid interval"); if (begin == end) { // empty interval return ImmutableList.of(); } List<PrestoThriftBlock> blocks = page.getColumnBlocks(); List<List<String>> result = new ArrayList<>(blocks.size()); for (PrestoThriftBlock block : blocks) { result.add(blockAsList(block, begin, end)); } return result; }
private static PrestoThriftPageResult pageResult(int value, PrestoThriftId nextToken) { return new PrestoThriftPageResult(ImmutableList.of(integerData(new PrestoThriftInteger(null, new int[] {value}))), 1, nextToken); } }
private void testRecordSetReadWrite(Random random, int records) { testReadWrite(random, records, blocks -> { List<Type> types = columns.stream().map(ColumnDefinition::getType).collect(toImmutableList()); PageRecordSet inputRecordSet = new PageRecordSet(types, new Page(blocks.toArray(new Block[blocks.size()]))); return fromRecordSet(inputRecordSet); }); }
@Override public ListenableFuture<PrestoThriftPageResult> getRows(PrestoThriftId splitId, List<String> columns, long maxBytes, PrestoThriftNullableToken nextToken) { if (rowsPerSplit == 0) { return immediateFuture(new PrestoThriftPageResult(ImmutableList.of(), 0, null)); } int key = Ints.fromByteArray(splitId.getId()); int offset = nextToken.getToken() != null ? Ints.fromByteArray(nextToken.getToken().getId()) : 0; PrestoThriftId newNextToken = offset + 1 < rowsPerSplit ? new PrestoThriftId(Ints.toByteArray(offset + 1)) : null; return immediateFuture(pageResult(key * 10 + offset, newNextToken)); }
private static PrestoThriftPageResult toThriftPage(Page page, List<Type> columnTypes, @Nullable PrestoThriftId nextToken) { if (page == null) { checkState(nextToken == null, "there must be no more data when page is null"); return new PrestoThriftPageResult(ImmutableList.of(), 0, null); } checkState(page.getChannelCount() == columnTypes.size(), "number of columns in a page doesn't match the one in requested types"); int numberOfColumns = columnTypes.size(); List<PrestoThriftBlock> columnBlocks = new ArrayList<>(numberOfColumns); for (int i = 0; i < numberOfColumns; i++) { columnBlocks.add(fromBlock(page.getBlock(i), columnTypes.get(i))); } return new PrestoThriftPageResult(columnBlocks, page.getPositionCount(), nextToken); }
private void testPageReadWrite(Random random, int records) { testReadWrite(random, records, blocks -> { List<PrestoThriftBlock> columnBlocks = new ArrayList<>(columns.size()); for (int i = 0; i < columns.size(); i++) { columnBlocks.add(fromBlock(blocks.get(i), columns.get(i).getType())); } return new PrestoThriftPageResult(columnBlocks, records, null); }); }