Examples with ColumnAdaptation io.trino.plugin.raptor.legacy.storage.RaptorPageSource.ColumnAdaptation used on opensource projects

Example 1 with ColumnAdaptation

use of io.trino.plugin.raptor.legacy.storage.RaptorPageSource.ColumnAdaptation in project trino by trinodb.

the class RaptorStorageManager method getPageSource.

@Override
public ConnectorPageSource getPageSource(UUID shardUuid, OptionalInt bucketNumber, List<Long> columnIds, List<Type> columnTypes, TupleDomain<RaptorColumnHandle> effectivePredicate, OrcReaderOptions orcReaderOptions, OptionalLong transactionId) {
    orcReaderOptions = orcReaderOptions.withMaxReadBlockSize(HUGE_MAX_READ_BLOCK_SIZE);
    OrcDataSource dataSource = openShard(shardUuid, orcReaderOptions);
    AggregatedMemoryContext memoryUsage = newSimpleAggregatedMemoryContext();
    try {
        OrcReader reader = OrcReader.createOrcReader(dataSource, orcReaderOptions).orElseThrow(() -> new TrinoException(RAPTOR_ERROR, "Data file is empty for shard " + shardUuid));
        Map<Long, OrcColumn> indexMap = columnIdIndex(reader.getRootColumn().getNestedColumns());
        List<OrcColumn> fileReadColumn = new ArrayList<>(columnIds.size());
        List<Type> fileReadTypes = new ArrayList<>(columnIds.size());
        List<ColumnAdaptation> columnAdaptations = new ArrayList<>(columnIds.size());
        for (int i = 0; i < columnIds.size(); i++) {
            long columnId = columnIds.get(i);
            if (isHiddenColumn(columnId)) {
                columnAdaptations.add(specialColumnAdaptation(columnId, shardUuid, bucketNumber));
                continue;
            }
            Type type = toOrcFileType(columnTypes.get(i), typeManager);
            OrcColumn fileColumn = indexMap.get(columnId);
            if (fileColumn == null) {
                columnAdaptations.add(ColumnAdaptation.nullColumn(type));
            } else {
                int sourceIndex = fileReadColumn.size();
                columnAdaptations.add(ColumnAdaptation.sourceColumn(sourceIndex));
                fileReadColumn.add(fileColumn);
                fileReadTypes.add(type);
            }
        }
        OrcPredicate predicate = getPredicate(effectivePredicate, indexMap);
        OrcRecordReader recordReader = reader.createRecordReader(fileReadColumn, fileReadTypes, predicate, UTC, memoryUsage, INITIAL_BATCH_SIZE, RaptorPageSource::handleException);
        Optional<ShardRewriter> shardRewriter = Optional.empty();
        if (transactionId.isPresent()) {
            shardRewriter = Optional.of(createShardRewriter(transactionId.getAsLong(), bucketNumber, shardUuid));
        }
        return new RaptorPageSource(shardRewriter, recordReader, columnAdaptations, dataSource, memoryUsage);
    } catch (IOException | RuntimeException e) {
        closeQuietly(dataSource);
        throw new TrinoException(RAPTOR_ERROR, "Failed to create page source for shard " + shardUuid, e);
    } catch (Throwable t) {
        closeQuietly(dataSource);
        throw t;
    }
}