Examples with AggregatedMemoryContext io.trino.memory.context.AggregatedMemoryContext used on opensource projects

Example 1 with AggregatedMemoryContext

use of io.trino.memory.context.AggregatedMemoryContext in project trino by trinodb.

the class OrcPageSourceFactory method createOrcPageSource.

private ConnectorPageSource createOrcPageSource(HdfsEnvironment hdfsEnvironment, ConnectorIdentity identity, Configuration configuration, Path path, long start, long length, long estimatedFileSize, List<HiveColumnHandle> columns, List<HiveColumnHandle> projections, boolean useOrcColumnNames, boolean isFullAcid, TupleDomain<HiveColumnHandle> effectivePredicate, DateTimeZone legacyFileTimeZone, OrcReaderOptions options, Optional<AcidInfo> acidInfo, OptionalInt bucketNumber, boolean originalFile, AcidTransaction transaction, FileFormatDataSourceStats stats) {
    for (HiveColumnHandle column : columns) {
        checkArgument(column.getColumnType() == REGULAR, "column type must be regular: %s", column);
    }
    checkArgument(!effectivePredicate.isNone());
    OrcDataSource orcDataSource;
    boolean originalFilesPresent = acidInfo.isPresent() && !acidInfo.get().getOriginalFiles().isEmpty();
    try {
        FileSystem fileSystem = hdfsEnvironment.getFileSystem(identity, path, configuration);
        FSDataInputStream inputStream = hdfsEnvironment.doAs(identity, () -> fileSystem.open(path));
        orcDataSource = new HdfsOrcDataSource(new OrcDataSourceId(path.toString()), estimatedFileSize, options, inputStream, stats);
    } catch (Exception e) {
        if (nullToEmpty(e.getMessage()).trim().equals("Filesystem closed") || e instanceof FileNotFoundException) {
            throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, e);
        }
        throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, splitError(e, path, start, length), e);
    }
    AggregatedMemoryContext memoryUsage = newSimpleAggregatedMemoryContext();
    try {
        Optional<OrcReader> optionalOrcReader = OrcReader.createOrcReader(orcDataSource, options);
        if (optionalOrcReader.isEmpty()) {
            return new EmptyPageSource();
        }
        OrcReader reader = optionalOrcReader.get();
        if (!originalFile && acidInfo.isPresent() && !acidInfo.get().isOrcAcidVersionValidated()) {
            validateOrcAcidVersion(path, reader);
        }
        List<OrcColumn> fileColumns = reader.getRootColumn().getNestedColumns();
        int actualColumnCount = columns.size() + (isFullAcid ? 3 : 0);
        List<OrcColumn> fileReadColumns = new ArrayList<>(actualColumnCount);
        List<Type> fileReadTypes = new ArrayList<>(actualColumnCount);
        List<OrcReader.ProjectedLayout> fileReadLayouts = new ArrayList<>(actualColumnCount);
        if (isFullAcid && !originalFilesPresent) {
            verifyAcidSchema(reader, path);
            Map<String, OrcColumn> acidColumnsByName = uniqueIndex(fileColumns, orcColumn -> orcColumn.getColumnName().toLowerCase(ENGLISH));
            fileColumns = ensureColumnNameConsistency(acidColumnsByName.get(AcidSchema.ACID_COLUMN_ROW_STRUCT.toLowerCase(ENGLISH)).getNestedColumns(), columns);
            fileReadColumns.add(acidColumnsByName.get(AcidSchema.ACID_COLUMN_ORIGINAL_TRANSACTION.toLowerCase(ENGLISH)));
            fileReadTypes.add(BIGINT);
            fileReadLayouts.add(fullyProjectedLayout());
            fileReadColumns.add(acidColumnsByName.get(AcidSchema.ACID_COLUMN_BUCKET.toLowerCase(ENGLISH)));
            fileReadTypes.add(INTEGER);
            fileReadLayouts.add(fullyProjectedLayout());
            fileReadColumns.add(acidColumnsByName.get(AcidSchema.ACID_COLUMN_ROW_ID.toLowerCase(ENGLISH)));
            fileReadTypes.add(BIGINT);
            fileReadLayouts.add(fullyProjectedLayout());
        }
        Map<String, OrcColumn> fileColumnsByName = ImmutableMap.of();
        if (useOrcColumnNames || isFullAcid) {
            verifyFileHasColumnNames(fileColumns, path);
            // Convert column names read from ORC files to lower case to be consistent with those stored in Hive Metastore
            fileColumnsByName = uniqueIndex(fileColumns, orcColumn -> orcColumn.getColumnName().toLowerCase(ENGLISH));
        }
        Map<String, List<List<String>>> projectionsByColumnName = ImmutableMap.of();
        Map<Integer, List<List<String>>> projectionsByColumnIndex = ImmutableMap.of();
        if (useOrcColumnNames || isFullAcid) {
            projectionsByColumnName = projections.stream().collect(Collectors.groupingBy(HiveColumnHandle::getBaseColumnName, mapping(OrcPageSourceFactory::getDereferencesAsList, toList())));
        } else {
            projectionsByColumnIndex = projections.stream().collect(Collectors.groupingBy(HiveColumnHandle::getBaseHiveColumnIndex, mapping(OrcPageSourceFactory::getDereferencesAsList, toList())));
        }
        TupleDomainOrcPredicateBuilder predicateBuilder = TupleDomainOrcPredicate.builder().setBloomFiltersEnabled(options.isBloomFiltersEnabled()).setDomainCompactionThreshold(domainCompactionThreshold);
        Map<HiveColumnHandle, Domain> effectivePredicateDomains = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Effective predicate is none"));
        List<ColumnAdaptation> columnAdaptations = new ArrayList<>(columns.size());
        for (HiveColumnHandle column : columns) {
            OrcColumn orcColumn = null;
            OrcReader.ProjectedLayout projectedLayout = null;
            Map<Optional<HiveColumnProjectionInfo>, Domain> columnDomains = null;
            if (useOrcColumnNames || isFullAcid) {
                String columnName = column.getName().toLowerCase(ENGLISH);
                orcColumn = fileColumnsByName.get(columnName);
                if (orcColumn != null) {
                    projectedLayout = createProjectedLayout(orcColumn, projectionsByColumnName.get(columnName));
                    columnDomains = effectivePredicateDomains.entrySet().stream().filter(columnDomain -> columnDomain.getKey().getBaseColumnName().toLowerCase(ENGLISH).equals(columnName)).collect(toImmutableMap(columnDomain -> columnDomain.getKey().getHiveColumnProjectionInfo(), Map.Entry::getValue));
                }
            } else if (column.getBaseHiveColumnIndex() < fileColumns.size()) {
                orcColumn = fileColumns.get(column.getBaseHiveColumnIndex());
                if (orcColumn != null) {
                    projectedLayout = createProjectedLayout(orcColumn, projectionsByColumnIndex.get(column.getBaseHiveColumnIndex()));
                    columnDomains = effectivePredicateDomains.entrySet().stream().filter(columnDomain -> columnDomain.getKey().getBaseHiveColumnIndex() == column.getBaseHiveColumnIndex()).collect(toImmutableMap(columnDomain -> columnDomain.getKey().getHiveColumnProjectionInfo(), Map.Entry::getValue));
                }
            }
            Type readType = column.getType();
            if (orcColumn != null) {
                int sourceIndex = fileReadColumns.size();
                columnAdaptations.add(ColumnAdaptation.sourceColumn(sourceIndex));
                fileReadColumns.add(orcColumn);
                fileReadTypes.add(readType);
                fileReadLayouts.add(projectedLayout);
                // Add predicates on top-level and nested columns
                for (Map.Entry<Optional<HiveColumnProjectionInfo>, Domain> columnDomain : columnDomains.entrySet()) {
                    OrcColumn nestedColumn = getNestedColumn(orcColumn, columnDomain.getKey());
                    if (nestedColumn != null) {
                        predicateBuilder.addColumn(nestedColumn.getColumnId(), columnDomain.getValue());
                    }
                }
            } else {
                columnAdaptations.add(ColumnAdaptation.nullColumn(readType));
            }
        }
        OrcRecordReader recordReader = reader.createRecordReader(fileReadColumns, fileReadTypes, fileReadLayouts, predicateBuilder.build(), start, length, legacyFileTimeZone, memoryUsage, INITIAL_BATCH_SIZE, exception -> handleException(orcDataSource.getId(), exception), NameBasedFieldMapper::create);
        Optional<OrcDeletedRows> deletedRows = acidInfo.map(info -> new OrcDeletedRows(path.getName(), new OrcDeleteDeltaPageSourceFactory(options, identity, configuration, hdfsEnvironment, stats), identity, configuration, hdfsEnvironment, info, bucketNumber, memoryUsage));
        Optional<Long> originalFileRowId = acidInfo.filter(OrcPageSourceFactory::hasOriginalFiles).map(info -> OriginalFilesUtils.getPrecedingRowCount(acidInfo.get().getOriginalFiles(), path, hdfsEnvironment, identity, options, configuration, stats));
        if (transaction.isDelete()) {
            if (originalFile) {
                int bucket = bucketNumber.orElse(0);
                long startingRowId = originalFileRowId.orElse(0L);
                columnAdaptations.add(ColumnAdaptation.originalFileRowIdColumn(startingRowId, bucket));
            } else {
                columnAdaptations.add(ColumnAdaptation.rowIdColumn());
            }
        } else if (transaction.isUpdate()) {
            HiveUpdateProcessor updateProcessor = transaction.getUpdateProcessor().orElseThrow(() -> new IllegalArgumentException("updateProcessor not present"));
            List<HiveColumnHandle> dependencyColumns = projections.stream().filter(HiveColumnHandle::isBaseColumn).collect(toImmutableList());
            if (originalFile) {
                int bucket = bucketNumber.orElse(0);
                long startingRowId = originalFileRowId.orElse(0L);
                columnAdaptations.add(updatedRowColumnsWithOriginalFiles(startingRowId, bucket, updateProcessor, dependencyColumns));
            } else {
                columnAdaptations.add(updatedRowColumns(updateProcessor, dependencyColumns));
            }
        }
        return new OrcPageSource(recordReader, columnAdaptations, orcDataSource, deletedRows, originalFileRowId, memoryUsage, stats);
    } catch (Exception e) {
        try {
            orcDataSource.close();
        } catch (IOException ignored) {
        }
        if (e instanceof TrinoException) {
            throw (TrinoException) e;
        }
        String message = splitError(e, path, start, length);
        if (e instanceof BlockMissingException) {
            throw new TrinoException(HIVE_MISSING_DATA, message, e);
        }
        throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, message, e);
    }
}

Example 2 with AggregatedMemoryContext

use of io.trino.memory.context.AggregatedMemoryContext in project trino by trinodb.

the class TestBroadcastOutputBuffer method testSharedBufferBlocking.

@Test
public void testSharedBufferBlocking() {
    SettableFuture<Void> blockedFuture = SettableFuture.create();
    MockMemoryReservationHandler reservationHandler = new MockMemoryReservationHandler(blockedFuture);
    AggregatedMemoryContext memoryContext = newRootAggregatedMemoryContext(reservationHandler, 0L);
    Page page = createPage(1);
    long pageSize = serializePage(page).getRetainedSize();
    // create a buffer that can only hold two pages
    BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), DataSize.ofBytes(pageSize * 2), memoryContext, directExecutor());
    OutputBufferMemoryManager memoryManager = buffer.getMemoryManager();
    // adding the first page will block as no memory is available (MockMemoryReservationHandler will return a future that is not done)
    enqueuePage(buffer, page);
    // more memory is available
    blockedFuture.set(null);
    memoryManager.onMemoryAvailable();
    assertTrue(memoryManager.getBufferBlockedFuture().isDone(), "buffer shouldn't be blocked");
    // we should be able to add one more page after more memory is available
    addPage(buffer, page);
    // the buffer is full now
    enqueuePage(buffer, page);
}

Example 3 with AggregatedMemoryContext

use of io.trino.memory.context.AggregatedMemoryContext in project trino by trinodb.

the class TestBroadcastOutputBuffer method testSharedBufferBlockingNoBlockOnFull.

@Test
public void testSharedBufferBlockingNoBlockOnFull() {
    SettableFuture<Void> blockedFuture = SettableFuture.create();
    MockMemoryReservationHandler reservationHandler = new MockMemoryReservationHandler(blockedFuture);
    AggregatedMemoryContext memoryContext = newRootAggregatedMemoryContext(reservationHandler, 0L);
    Page page = createPage(1);
    long pageSize = serializePage(page).getRetainedSize();
    // create a buffer that can only hold two pages
    BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), DataSize.ofBytes(pageSize * 2), memoryContext, directExecutor());
    OutputBufferMemoryManager memoryManager = buffer.getMemoryManager();
    memoryManager.setNoBlockOnFull();
    // even if setNoBlockOnFull() is called the buffer should block on memory when we add the first page
    // as no memory is available (MockMemoryReservationHandler will return a future that is not done)
    enqueuePage(buffer, page);
    // more memory is available
    blockedFuture.set(null);
    memoryManager.onMemoryAvailable();
    assertTrue(memoryManager.getBufferBlockedFuture().isDone(), "buffer shouldn't be blocked");
    // we should be able to add one more page after more memory is available
    addPage(buffer, page);
    // the buffer is full now, but setNoBlockOnFull() is called so the buffer shouldn't block
    addPage(buffer, page);
}

Example 4 with AggregatedMemoryContext

use of io.trino.memory.context.AggregatedMemoryContext in project trino by trinodb.

the class TestBroadcastOutputBuffer method testSharedBufferBlocking2.

@Test
public void testSharedBufferBlocking2() {
    // start with a complete future
    SettableFuture<Void> blockedFuture = SettableFuture.create();
    blockedFuture.set(null);
    MockMemoryReservationHandler reservationHandler = new MockMemoryReservationHandler(blockedFuture);
    AggregatedMemoryContext memoryContext = newRootAggregatedMemoryContext(reservationHandler, 0L);
    Page page = createPage(1);
    long pageSize = serializePage(page).getRetainedSize();
    // create a buffer that can only hold two pages
    BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), DataSize.ofBytes(pageSize * 2), memoryContext, directExecutor());
    OutputBufferMemoryManager memoryManager = buffer.getMemoryManager();
    // add two pages to fill up the buffer (memory is available)
    addPage(buffer, page);
    addPage(buffer, page);
    // fill up the memory pool
    blockedFuture = SettableFuture.create();
    reservationHandler.updateBlockedFuture(blockedFuture);
    // allocate one more byte to make the buffer full
    memoryManager.updateMemoryUsage(1L);
    // more memory is available
    blockedFuture.set(null);
    memoryManager.onMemoryAvailable();
    // memoryManager should still return a blocked future as the buffer is still full
    assertFalse(memoryManager.getBufferBlockedFuture().isDone(), "buffer should be blocked");
    // remove all pages from the memory manager and the 1 byte that we added above
    memoryManager.updateMemoryUsage(-pageSize * 2 - 1);
    // now we have both buffer space and memory available, so memoryManager shouldn't be blocked
    assertTrue(memoryManager.getBufferBlockedFuture().isDone(), "buffer shouldn't be blocked");
    // we should be able to add two pages after more memory is available
    addPage(buffer, page);
    addPage(buffer, page);
    // the buffer is full now
    enqueuePage(buffer, page);
}

Example 5 with AggregatedMemoryContext

use of io.trino.memory.context.AggregatedMemoryContext in project trino by trinodb.

the class TestSliceDictionaryColumnReader method testDictionaryReaderUpdatesRetainedSize.

@Test
public void testDictionaryReaderUpdatesRetainedSize() throws Exception {
    // create orc file
    List<String> values = createValues();
    File temporaryDirectory = createTempDir();
    File orcFile = new File(temporaryDirectory, randomUUID().toString());
    writeOrcColumnTrino(orcFile, NONE, VARCHAR, values.iterator(), new OrcWriterStats());
    // prepare for read
    OrcDataSource dataSource = new MemoryOrcDataSource(new OrcDataSourceId(orcFile.getPath()), Slices.wrappedBuffer(readAllBytes(orcFile.toPath())));
    OrcReader orcReader = OrcReader.createOrcReader(dataSource, new OrcReaderOptions()).orElseThrow(() -> new RuntimeException("File is empty"));
    Footer footer = orcReader.getFooter();
    List<OrcColumn> columns = orcReader.getRootColumn().getNestedColumns();
    assertTrue(columns.size() == 1);
    StripeReader stripeReader = new StripeReader(dataSource, UTC, Optional.empty(), footer.getTypes(), ImmutableSet.copyOf(columns), footer.getRowsInRowGroup(), OrcPredicate.TRUE, ORIGINAL, new OrcMetadataReader(), Optional.empty());
    AggregatedMemoryContext memoryContext = newSimpleAggregatedMemoryContext();
    SliceDictionaryColumnReader columnReader = new SliceDictionaryColumnReader(columns.get(0), memoryContext.newLocalMemoryContext(TestSliceDictionaryColumnReader.class.getSimpleName()), -1, false);
    List<StripeInformation> stripeInformations = footer.getStripes();
    for (StripeInformation stripeInformation : stripeInformations) {
        Stripe stripe = stripeReader.readStripe(stripeInformation, newSimpleAggregatedMemoryContext());
        List<RowGroup> rowGroups = stripe.getRowGroups();
        columnReader.startStripe(stripe.getFileTimeZone(), stripe.getDictionaryStreamSources(), stripe.getColumnEncodings());
        for (RowGroup rowGroup : rowGroups) {
            columnReader.startRowGroup(rowGroup.getStreamSources());
            columnReader.prepareNextRead(1000);
            columnReader.readBlock();
            // memory usage check
            assertEquals(memoryContext.getBytes(), columnReader.getRetainedSizeInBytes());
        }
    }
    columnReader.close();
    assertTrue(memoryContext.getBytes() == 0);
}