use of io.trino.orc.metadata.Stream in project trino by trinodb.
the class TestBooleanStream method testWriteMultiple.
@Test
public void testWriteMultiple() throws IOException {
BooleanOutputStream outputStream = createValueOutputStream();
for (int i = 0; i < 3; i++) {
outputStream.reset();
BooleanList expectedValues = new BooleanArrayList(1024);
outputStream.writeBooleans(32, true);
expectedValues.addAll(Collections.nCopies(32, true));
outputStream.writeBooleans(32, false);
expectedValues.addAll(Collections.nCopies(32, false));
outputStream.writeBooleans(1, true);
expectedValues.add(true);
outputStream.writeBooleans(1, false);
expectedValues.add(false);
outputStream.writeBooleans(34, true);
expectedValues.addAll(Collections.nCopies(34, true));
outputStream.writeBooleans(34, false);
expectedValues.addAll(Collections.nCopies(34, false));
outputStream.writeBoolean(true);
expectedValues.add(true);
outputStream.writeBoolean(false);
expectedValues.add(false);
outputStream.close();
DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1000);
StreamDataOutput streamDataOutput = outputStream.getStreamDataOutput(new OrcColumnId(33));
streamDataOutput.writeData(sliceOutput);
Stream stream = streamDataOutput.getStream();
assertEquals(stream.getStreamKind(), StreamKind.DATA);
assertEquals(stream.getColumnId(), new OrcColumnId(33));
assertEquals(stream.getLength(), sliceOutput.size());
BooleanInputStream valueStream = createValueStream(sliceOutput.slice());
for (int index = 0; index < expectedValues.size(); index++) {
boolean expectedValue = expectedValues.getBoolean(index);
boolean actualValue = readValue(valueStream);
assertEquals(actualValue, expectedValue);
}
}
}
use of io.trino.orc.metadata.Stream in project trino by trinodb.
the class AbstractTestValueStream method testWriteValue.
protected void testWriteValue(List<List<T>> groups) throws IOException {
W outputStream = createValueOutputStream();
for (int i = 0; i < 3; i++) {
outputStream.reset();
long retainedBytes = 0;
for (List<T> group : groups) {
outputStream.recordCheckpoint();
group.forEach(value -> writeValue(outputStream, value));
assertTrue(outputStream.getRetainedBytes() >= retainedBytes);
retainedBytes = outputStream.getRetainedBytes();
}
outputStream.close();
DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1000);
StreamDataOutput streamDataOutput = outputStream.getStreamDataOutput(new OrcColumnId(33));
streamDataOutput.writeData(sliceOutput);
Stream stream = streamDataOutput.getStream();
assertEquals(stream.getStreamKind(), StreamKind.DATA);
assertEquals(stream.getColumnId(), new OrcColumnId(33));
assertEquals(stream.getLength(), sliceOutput.size());
List<C> checkpoints = outputStream.getCheckpoints();
assertEquals(checkpoints.size(), groups.size());
R valueStream = createValueStream(sliceOutput.slice());
for (List<T> group : groups) {
int index = 0;
for (T expectedValue : group) {
index++;
T actualValue = readValue(valueStream);
if (!actualValue.equals(expectedValue)) {
assertEquals(actualValue, expectedValue, "index=" + index);
}
}
}
for (int groupIndex = groups.size() - 1; groupIndex >= 0; groupIndex--) {
valueStream.seekToCheckpoint(checkpoints.get(groupIndex));
for (T expectedValue : groups.get(groupIndex)) {
T actualValue = readValue(valueStream);
if (!actualValue.equals(expectedValue)) {
assertEquals(actualValue, expectedValue);
}
}
}
}
}
use of io.trino.orc.metadata.Stream in project trino by trinodb.
the class StripeReader method readStripe.
public Stripe readStripe(StripeInformation stripe, AggregatedMemoryContext memoryUsage) throws IOException {
// read the stripe footer
StripeFooter stripeFooter = readStripeFooter(stripe, memoryUsage);
ColumnMetadata<ColumnEncoding> columnEncodings = stripeFooter.getColumnEncodings();
if (writeValidation.isPresent()) {
writeValidation.get().validateTimeZone(orcDataSource.getId(), stripeFooter.getTimeZone());
}
ZoneId fileTimeZone = stripeFooter.getTimeZone();
// get streams for selected columns
Map<StreamId, Stream> streams = new HashMap<>();
for (Stream stream : stripeFooter.getStreams()) {
if (includedOrcColumnIds.contains(stream.getColumnId()) && isSupportedStreamType(stream, types.get(stream.getColumnId()).getOrcTypeKind())) {
streams.put(new StreamId(stream), stream);
}
}
// handle stripes with more than one row group
boolean invalidCheckPoint = false;
if (rowsInRowGroup.isPresent() && stripe.getNumberOfRows() > rowsInRowGroup.getAsInt()) {
// determine ranges of the stripe to read
Map<StreamId, DiskRange> diskRanges = getDiskRanges(stripeFooter.getStreams());
diskRanges = Maps.filterKeys(diskRanges, Predicates.in(streams.keySet()));
// read the file regions
Map<StreamId, OrcChunkLoader> streamsData = readDiskRanges(stripe.getOffset(), diskRanges, memoryUsage);
// read the bloom filter for each column
Map<OrcColumnId, List<BloomFilter>> bloomFilterIndexes = readBloomFilterIndexes(streams, streamsData);
// read the row index for each column
Map<StreamId, List<RowGroupIndex>> columnIndexes = readColumnIndexes(streams, streamsData, bloomFilterIndexes);
if (writeValidation.isPresent()) {
writeValidation.get().validateRowGroupStatistics(orcDataSource.getId(), stripe.getOffset(), columnIndexes);
}
// select the row groups matching the tuple domain
Set<Integer> selectedRowGroups = selectRowGroups(stripe, columnIndexes);
// if all row groups are skipped, return null
if (selectedRowGroups.isEmpty()) {
// set accounted memory usage to zero
memoryUsage.close();
return null;
}
// value streams
Map<StreamId, ValueInputStream<?>> valueStreams = createValueStreams(streams, streamsData, columnEncodings);
// build the dictionary streams
InputStreamSources dictionaryStreamSources = createDictionaryStreamSources(streams, valueStreams, columnEncodings);
// build the row groups
try {
List<RowGroup> rowGroups = createRowGroups(stripe.getNumberOfRows(), streams, valueStreams, columnIndexes, selectedRowGroups, columnEncodings);
return new Stripe(stripe.getNumberOfRows(), fileTimeZone, columnEncodings, rowGroups, dictionaryStreamSources);
} catch (InvalidCheckpointException e) {
// The ORC file contains a corrupt checkpoint stream treat the stripe as a single row group.
invalidCheckPoint = true;
}
}
// stripe only has one row group
ImmutableMap.Builder<StreamId, DiskRange> diskRangesBuilder = ImmutableMap.builder();
for (Entry<StreamId, DiskRange> entry : getDiskRanges(stripeFooter.getStreams()).entrySet()) {
StreamId streamId = entry.getKey();
if (streams.containsKey(streamId)) {
diskRangesBuilder.put(entry);
}
}
ImmutableMap<StreamId, DiskRange> diskRanges = diskRangesBuilder.buildOrThrow();
// read the file regions
Map<StreamId, OrcChunkLoader> streamsData = readDiskRanges(stripe.getOffset(), diskRanges, memoryUsage);
long minAverageRowBytes = 0;
for (Entry<StreamId, Stream> entry : streams.entrySet()) {
if (entry.getKey().getStreamKind() == ROW_INDEX) {
List<RowGroupIndex> rowGroupIndexes = metadataReader.readRowIndexes(hiveWriterVersion, new OrcInputStream(streamsData.get(entry.getKey())));
checkState(rowGroupIndexes.size() == 1 || invalidCheckPoint, "expect a single row group or an invalid check point");
long totalBytes = 0;
long totalRows = 0;
for (RowGroupIndex rowGroupIndex : rowGroupIndexes) {
ColumnStatistics columnStatistics = rowGroupIndex.getColumnStatistics();
if (columnStatistics.hasMinAverageValueSizeInBytes()) {
totalBytes += columnStatistics.getMinAverageValueSizeInBytes() * columnStatistics.getNumberOfValues();
totalRows += columnStatistics.getNumberOfValues();
}
}
if (totalRows > 0) {
minAverageRowBytes += totalBytes / totalRows;
}
}
}
// value streams
Map<StreamId, ValueInputStream<?>> valueStreams = createValueStreams(streams, streamsData, columnEncodings);
// build the dictionary streams
InputStreamSources dictionaryStreamSources = createDictionaryStreamSources(streams, valueStreams, columnEncodings);
// build the row group
ImmutableMap.Builder<StreamId, InputStreamSource<?>> builder = ImmutableMap.builder();
for (Entry<StreamId, ValueInputStream<?>> entry : valueStreams.entrySet()) {
builder.put(entry.getKey(), new ValueInputStreamSource<>(entry.getValue()));
}
RowGroup rowGroup = new RowGroup(0, 0, stripe.getNumberOfRows(), minAverageRowBytes, new InputStreamSources(builder.buildOrThrow()));
return new Stripe(stripe.getNumberOfRows(), fileTimeZone, columnEncodings, ImmutableList.of(rowGroup), dictionaryStreamSources);
}
use of io.trino.orc.metadata.Stream in project trino by trinodb.
the class StripeReader method getRowGroupStatistics.
private static ColumnMetadata<ColumnStatistics> getRowGroupStatistics(ColumnMetadata<OrcType> types, Map<StreamId, List<RowGroupIndex>> columnIndexes, int rowGroup) {
requireNonNull(columnIndexes, "columnIndexes is null");
checkArgument(rowGroup >= 0, "rowGroup is negative");
Map<Integer, List<RowGroupIndex>> rowGroupIndexesByColumn = columnIndexes.entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getColumnId().getId(), Entry::getValue));
List<ColumnStatistics> statistics = new ArrayList<>(types.size());
for (int columnIndex = 0; columnIndex < types.size(); columnIndex++) {
List<RowGroupIndex> rowGroupIndexes = rowGroupIndexesByColumn.get(columnIndex);
if (rowGroupIndexes != null) {
statistics.add(rowGroupIndexes.get(rowGroup).getColumnStatistics());
} else {
statistics.add(null);
}
}
return new ColumnMetadata<>(statistics);
}
use of io.trino.orc.metadata.Stream in project trino by trinodb.
the class StripeReader method createRowGroups.
private List<RowGroup> createRowGroups(int rowsInStripe, Map<StreamId, Stream> streams, Map<StreamId, ValueInputStream<?>> valueStreams, Map<StreamId, List<RowGroupIndex>> columnIndexes, Set<Integer> selectedRowGroups, ColumnMetadata<ColumnEncoding> encodings) throws InvalidCheckpointException {
int rowsInRowGroup = this.rowsInRowGroup.orElseThrow(() -> new IllegalStateException("Cannot create row groups if row group info is missing"));
ImmutableList.Builder<RowGroup> rowGroupBuilder = ImmutableList.builder();
for (int rowGroupId : selectedRowGroups) {
Map<StreamId, StreamCheckpoint> checkpoints = getStreamCheckpoints(includedOrcColumnIds, types, decompressor.isPresent(), rowGroupId, encodings, streams, columnIndexes);
int rowOffset = rowGroupId * rowsInRowGroup;
int rowsInGroup = Math.min(rowsInStripe - rowOffset, rowsInRowGroup);
long minAverageRowBytes = columnIndexes.entrySet().stream().mapToLong(e -> e.getValue().get(rowGroupId).getColumnStatistics().getMinAverageValueSizeInBytes()).sum();
rowGroupBuilder.add(createRowGroup(rowGroupId, rowOffset, rowsInGroup, minAverageRowBytes, valueStreams, checkpoints));
}
return rowGroupBuilder.build();
}
Aggregations