use of com.facebook.presto.common.block.RowBlock in project presto by prestodb.
the class StructBatchStreamReader method readBlock.
@Override
public Block readBlock() throws IOException {
if (!rowGroupOpen) {
openRowGroup();
}
if (readOffset > 0) {
if (presentStream != null) {
// skip ahead the present bit reader, but count the set bits
// and use this as the skip size for the field readers
readOffset = presentStream.countBitsSet(readOffset);
}
for (BatchStreamReader structField : structFields.values()) {
structField.prepareNextRead(readOffset);
}
}
boolean[] nullVector = null;
Block[] blocks;
if (presentStream == null) {
blocks = getBlocksForType(nextBatchSize);
} else {
nullVector = new boolean[nextBatchSize];
int nullValues = presentStream.getUnsetBits(nextBatchSize, nullVector);
if (nullValues != nextBatchSize) {
blocks = getBlocksForType(nextBatchSize - nullValues);
} else {
List<Type> typeParameters = type.getTypeParameters();
blocks = new Block[typeParameters.size()];
for (int i = 0; i < typeParameters.size(); i++) {
blocks[i] = typeParameters.get(i).createBlockBuilder(null, 0).build();
}
}
}
verify(Arrays.stream(blocks).mapToInt(Block::getPositionCount).distinct().count() == 1);
// Struct is represented as a row block
Block rowBlock = RowBlock.fromFieldBlocks(nextBatchSize, Optional.ofNullable(nullVector), blocks);
readOffset = 0;
nextBatchSize = 0;
return rowBlock;
}
use of com.facebook.presto.common.block.RowBlock in project presto by prestodb.
the class TestStructBatchStreamReader method testReaderLowerCasesFieldNamesFromStream.
/**
* The writer has fields with upper case characters, reader has same names downcased.
*/
@Test
public void testReaderLowerCasesFieldNamesFromStream() throws IOException {
List<String> readerFields = new ArrayList<>(Arrays.asList("field_a", "field_b", "field_c"));
List<String> writerFields = new ArrayList<>(Arrays.asList("field_A", "field_B", "field_C"));
List<String> writerData = new ArrayList<>(Arrays.asList("fieldAValue", "fieldBValue", "fieldCValue"));
Type readerType = getType(readerFields);
Type writerType = getType(writerFields);
write(tempFile, writerType, writerData);
RowBlock readBlock = read(tempFile, readerType);
List actual = (List) readerType.getObjectValue(SESSION.getSqlFunctionProperties(), readBlock, 0);
assertEquals(actual.size(), readerFields.size());
assertEquals(actual.get(0), "fieldAValue");
assertEquals(actual.get(1), "fieldBValue");
assertEquals(actual.get(2), "fieldCValue");
}
use of com.facebook.presto.common.block.RowBlock in project presto by prestodb.
the class TestStructBatchStreamReader method testExtraFieldsInWriter.
/**
* The ORC file has a field that is missing from the reader
*/
@Test
public void testExtraFieldsInWriter() throws IOException {
// field_b is missing
List<String> readerFields = new ArrayList<>(Arrays.asList("field_a", "field_c"));
List<String> writerFields = new ArrayList<>(Arrays.asList("field_a", "field_b", "field_c"));
List<String> writerData = new ArrayList<>(Arrays.asList("field_a_value", "field_b_value", "field_c_value"));
Type readerType = getType(readerFields);
Type writerType = getType(writerFields);
write(tempFile, writerType, writerData);
RowBlock readBlock = read(tempFile, readerType);
List actual = (List) readerType.getObjectValue(SESSION.getSqlFunctionProperties(), readBlock, 0);
assertEquals(actual.size(), readerFields.size());
assertEquals(actual.get(0), "field_a_value");
assertEquals(actual.get(1), "field_c_value");
}
use of com.facebook.presto.common.block.RowBlock in project presto by prestodb.
the class TestStructBatchStreamReader method read.
private RowBlock read(TempFile tempFile, Type readerType) throws IOException {
DataSize dataSize = new DataSize(1, MEGABYTE);
OrcDataSource orcDataSource = new FileOrcDataSource(tempFile.getFile(), dataSize, dataSize, dataSize, true);
OrcReader orcReader = new OrcReader(orcDataSource, ORC, new StorageOrcFileTailSource(), new StorageStripeMetadataSource(), NOOP_ORC_AGGREGATED_MEMORY_CONTEXT, new OrcReaderOptions(dataSize, dataSize, dataSize, false), false, NO_ENCRYPTION, DwrfKeyProvider.EMPTY, new RuntimeStats());
Map<Integer, Type> includedColumns = new HashMap<>();
includedColumns.put(0, readerType);
OrcBatchRecordReader recordReader = orcReader.createBatchRecordReader(includedColumns, OrcPredicate.TRUE, UTC, new TestingHiveOrcAggregatedMemoryContext(), OrcReader.INITIAL_BATCH_SIZE);
recordReader.nextBatch();
RowBlock block = (RowBlock) recordReader.readBlock(0);
recordReader.close();
return block;
}
use of com.facebook.presto.common.block.RowBlock in project presto by prestodb.
the class TestStructBatchStreamReader method testReaderLowerCasesFieldNamesFromType.
/**
* Reader has fields with upper case characters, writer has same names downcased.
*/
@Test
public void testReaderLowerCasesFieldNamesFromType() throws IOException {
List<String> readerFields = new ArrayList<>(Arrays.asList("field_A", "field_B", "field_C"));
List<String> writerFields = new ArrayList<>(Arrays.asList("field_a", "field_b", "field_c"));
List<String> writerData = new ArrayList<>(Arrays.asList("fieldAValue", "fieldBValue", "fieldCValue"));
Type readerType = getType(readerFields);
Type writerType = getType(writerFields);
write(tempFile, writerType, writerData);
RowBlock readBlock = read(tempFile, readerType);
List actual = (List) readerType.getObjectValue(SESSION.getSqlFunctionProperties(), readBlock, 0);
assertEquals(actual.size(), readerFields.size());
assertEquals(actual.get(0), "fieldAValue");
assertEquals(actual.get(1), "fieldBValue");
assertEquals(actual.get(2), "fieldCValue");
}
Aggregations