Java Code Examples for io.airlift.units.DataSize#toBytes()

public static String formatDataRate(DataSize dataSize, Duration duration, boolean longForm)
{
    double rate = dataSize.toBytes() / duration.getValue(SECONDS);
    if (Double.isNaN(rate) || Double.isInfinite(rate)) {
        rate = 0;
    }

    String rateString = formatDataSize(DataSize.ofBytes(Math.round(rate)), false);
    if (longForm) {
        if (!rateString.endsWith("B")) {
            rateString += "B";
        }
        rateString += "/s";
    }
    return rateString;
}
 

@Override
public synchronized List<SerializedPageReference> getPages(DataSize maxSize)
{
    long maxBytes = maxSize.toBytes();
    List<SerializedPageReference> pages = new ArrayList<>();
    long bytesRemoved = 0;

    while (true) {
        SerializedPageReference page = masterBuffer.peek();
        if (page == null) {
            break;
        }
        bytesRemoved += page.getRetainedSizeInBytes();
        // break (and don't add) if this page would exceed the limit
        if (!pages.isEmpty() && bytesRemoved > maxBytes) {
            break;
        }
        // this should not happen since we have a lock
        checkState(masterBuffer.poll() == page, "Master buffer corrupted");
        pages.add(page);
    }

    bufferedPages.set(masterBuffer.size());

    return ImmutableList.copyOf(pages);
}
 

@Override
public synchronized List<SerializedPageReference> getPages(DataSize maxSize)
{
    long maxBytes = maxSize.toBytes();
    List<SerializedPageReference> pages = new ArrayList<>();
    long bytesRemoved = 0;

    while (true) {
        SerializedPageReference page = buffer.peek();
        if (page == null) {
            break;
        }
        bytesRemoved += page.getRetainedSizeInBytes();
        // break (and don't add) if this page would exceed the limit
        if (!pages.isEmpty() && bytesRemoved > maxBytes) {
            break;
        }
        // this should not happen since we have a lock
        checkState(buffer.poll() == page, "Buffer corrupted");
        pages.add(page);
    }

    return ImmutableList.copyOf(pages);
}
 

@VisibleForTesting
static OrcDataSource wrapWithCacheIfTinyStripes(OrcDataSource dataSource, List<StripeInformation> stripes, DataSize maxMergeDistance, DataSize maxReadSize)
{
    if (dataSource instanceof CachingOrcDataSource) {
        return dataSource;
    }
    for (StripeInformation stripe : stripes) {
        if (stripe.getTotalLength() > maxReadSize.toBytes()) {
            return dataSource;
        }
    }
    return new CachingOrcDataSource(dataSource, createTinyStripesRangeFinder(stripes, maxMergeDistance, maxReadSize));
}
 

private boolean canReplicate(SemiJoinNode node, Context context)
{
    DataSize joinMaxBroadcastTableSize = getJoinMaxBroadcastTableSize(context.getSession());

    PlanNode buildSide = node.getFilteringSource();
    PlanNodeStatsEstimate buildSideStatsEstimate = context.getStatsProvider().getStats(buildSide);
    double buildSideSizeInBytes = buildSideStatsEstimate.getOutputSizeInBytes(buildSide.getOutputSymbols(), context.getSymbolAllocator().getTypes());
    return buildSideSizeInBytes <= joinMaxBroadcastTableSize.toBytes();
}
 

public static String formatDataSize(DataSize size, boolean longForm)
{
    double fractional = size.toBytes();
    String unit = null;
    if (fractional >= 1024) {
        fractional /= 1024;
        unit = "K";
    }
    if (fractional >= 1024) {
        fractional /= 1024;
        unit = "M";
    }
    if (fractional >= 1024) {
        fractional /= 1024;
        unit = "G";
    }
    if (fractional >= 1024) {
        fractional /= 1024;
        unit = "T";
    }
    if (fractional >= 1024) {
        fractional /= 1024;
        unit = "P";
    }

    if (unit == null) {
        unit = "B";
    }
    else if (longForm) {
        unit += "B";
    }

    return format("%s%s", getFormat(fractional).format(fractional), unit);
}
 

/**
 * Merge disk ranges that are closer than {@code maxMergeDistance}.
 */
public static List<DiskRange> mergeAdjacentDiskRanges(Collection<DiskRange> diskRanges, DataSize maxMergeDistance, DataSize maxReadSize)
{
    // sort ranges by start offset
    List<DiskRange> ranges = new ArrayList<>(diskRanges);
    Collections.sort(ranges, new Comparator<DiskRange>()
    {
        @Override
        public int compare(DiskRange o1, DiskRange o2)
        {
            return Long.compare(o1.getOffset(), o2.getOffset());
        }
    });

    // merge overlapping ranges
    long maxReadSizeBytes = maxReadSize.toBytes();
    long maxMergeDistanceBytes = maxMergeDistance.toBytes();
    ImmutableList.Builder<DiskRange> result = ImmutableList.builder();
    DiskRange last = ranges.get(0);
    for (int i = 1; i < ranges.size(); i++) {
        DiskRange current = ranges.get(i);
        DiskRange merged = last.span(current);
        if (merged.getLength() <= maxReadSizeBytes && last.getEnd() + maxMergeDistanceBytes >= current.getOffset()) {
            last = merged;
        }
        else {
            result.add(last);
            last = current;
        }
    }
    result.add(last);

    return result.build();
}
 

public SortBuffer(
        DataSize maxMemory,
        List<Type> types,
        List<Integer> sortFields,
        List<SortOrder> sortOrders,
        PageSorter pageSorter)
{
    checkArgument(maxMemory.toBytes() > 0, "maxMemory is zero");
    this.maxMemoryBytes = maxMemory.toBytes();
    this.types = requireNonNull(types, "types is null");
    this.sortFields = ImmutableList.copyOf(requireNonNull(sortFields, "sortFields is null"));
    this.sortOrders = ImmutableList.copyOf(requireNonNull(sortOrders, "sortOrders is null"));
    this.pageSorter = requireNonNull(pageSorter, "pageSorter is null");
    this.pageBuilder = new PageBuilder(types);
}
 

public static List<DiskRange> mergeAdjacentDiskRanges(Collection<DiskRange> diskRanges, DataSize maxMergeDistance, DataSize maxReadSize)
{
    // sort ranges by start offset
    List<DiskRange> ranges = new ArrayList<>(diskRanges);
    ranges.sort(comparingLong(DiskRange::getOffset));

    long maxReadSizeBytes = maxReadSize.toBytes();
    long maxMergeDistanceBytes = maxMergeDistance.toBytes();

    // merge overlapping ranges
    ImmutableList.Builder<DiskRange> result = ImmutableList.builder();
    DiskRange last = ranges.get(0);
    for (int i = 1; i < ranges.size(); i++) {
        DiskRange current = ranges.get(i);
        DiskRange merged = null;
        boolean blockTooLong = false;
        try {
            merged = last.span(current);
        }
        catch (ArithmeticException e) {
            blockTooLong = true;
        }
        if (!blockTooLong && merged.getLength() <= maxReadSizeBytes && last.getEnd() + maxMergeDistanceBytes >= current.getOffset()) {
            last = merged;
        }
        else {
            result.add(last);
            last = current;
        }
    }
    result.add(last);

    return result.build();
}
 

public SpillableHashAggregationBuilder(
        List<AccumulatorFactory> accumulatorFactories,
        AggregationNode.Step step,
        int expectedGroups,
        List<Type> groupByTypes,
        List<Integer> groupByChannels,
        Optional<Integer> hashChannel,
        OperatorContext operatorContext,
        DataSize memoryLimitForMerge,
        DataSize memoryLimitForMergeWithMemory,
        SpillerFactory spillerFactory,
        JoinCompiler joinCompiler)
{
    this.accumulatorFactories = accumulatorFactories;
    this.step = step;
    this.expectedGroups = expectedGroups;
    this.groupByTypes = groupByTypes;
    this.groupByChannels = groupByChannels;
    this.hashChannel = hashChannel;
    this.operatorContext = operatorContext;
    this.localUserMemoryContext = operatorContext.localUserMemoryContext();
    this.localRevocableMemoryContext = operatorContext.localRevocableMemoryContext();
    this.memoryLimitForMerge = memoryLimitForMerge.toBytes();
    this.memoryLimitForMergeWithMemory = memoryLimitForMergeWithMemory.toBytes();
    this.spillerFactory = spillerFactory;
    this.joinCompiler = joinCompiler;

    rebuildHashAggregationBuilder();
}
 

private DynamicFilterSourceOperator(
        OperatorContext context,
        Consumer<TupleDomain<DynamicFilterId>> dynamicPredicateConsumer,
        List<Channel> channels,
        PlanNodeId planNodeId,
        int maxFilterPositionsCount,
        DataSize maxFilterSize)
{
    this.context = requireNonNull(context, "context is null");
    this.maxFilterPositionsCount = maxFilterPositionsCount;
    this.maxFilterSizeInBytes = maxFilterSize.toBytes();

    this.dynamicPredicateConsumer = requireNonNull(dynamicPredicateConsumer, "dynamicPredicateConsumer is null");
    this.channels = requireNonNull(channels, "channels is null");

    this.blockBuilders = new BlockBuilder[channels.size()];
    this.valueSets = new TypedSet[channels.size()];
    for (int channelIndex = 0; channelIndex < channels.size(); ++channelIndex) {
        Type type = channels.get(channelIndex).type;
        this.blockBuilders[channelIndex] = type.createBlockBuilder(null, EXPECTED_BLOCK_BUILDER_SIZE);
        this.valueSets[channelIndex] = new TypedSet(
                type,
                Optional.empty(),
                blockBuilders[channelIndex],
                EXPECTED_BLOCK_BUILDER_SIZE,
                String.format("DynamicFilterSourceOperator_%s_%d", planNodeId, channelIndex),
                Optional.empty() /* maxBlockMemory */);
    }
}
 

private static Long toBytes(DataSize dataSize)
{
    if (dataSize == null) {
        return null;
    }
    return dataSize.toBytes();
}
 

@VisibleForTesting
static OrcDataSource wrapWithCacheIfTinyStripes(OrcDataSource dataSource, List<StripeInformation> stripes, DataSize maxMergeDistance, DataSize tinyStripeThreshold)
{
    if (dataSource instanceof CachingOrcDataSource) {
        return dataSource;
    }
    for (StripeInformation stripe : stripes) {
        if (stripe.getTotalLength() > tinyStripeThreshold.toBytes()) {
            return dataSource;
        }
    }
    return new CachingOrcDataSource(dataSource, createTinyStripesRangeFinder(stripes, maxMergeDistance, tinyStripeThreshold));
}
 

/**
 * Merge disk ranges that are closer than {@code maxMergeDistance}.
 */
public static List<DiskRange> mergeAdjacentDiskRanges(Collection<DiskRange> diskRanges, DataSize maxMergeDistance, DataSize maxReadSize)
{
    // sort ranges by start offset
    List<DiskRange> ranges = new ArrayList<>(diskRanges);
    ranges.sort(comparingLong(DiskRange::getOffset));

    // merge overlapping ranges
    long maxReadSizeBytes = maxReadSize.toBytes();
    long maxMergeDistanceBytes = maxMergeDistance.toBytes();
    ImmutableList.Builder<DiskRange> result = ImmutableList.builder();
    DiskRange last = ranges.get(0);
    for (int i = 1; i < ranges.size(); i++) {
        DiskRange current = ranges.get(i);
        DiskRange merged = last.span(current);
        if (merged.getLength() <= maxReadSizeBytes && last.getEnd() + maxMergeDistanceBytes >= current.getOffset()) {
            last = merged;
        }
        else {
            result.add(last);
            last = current;
        }
    }
    result.add(last);

    return result.build();
}
 

private static OrcDataSource wrapWithCacheIfTiny(OrcDataSource dataSource, DataSize maxCacheSize)
{
    if (dataSource instanceof CachingOrcDataSource) {
        return dataSource;
    }
    if (dataSource.getSize() > maxCacheSize.toBytes()) {
        return dataSource;
    }
    DiskRange diskRange = new DiskRange(0, toIntExact(dataSource.getSize()));
    return new CachingOrcDataSource(dataSource, desiredOffset -> diskRange);
}
 

public RecordFileWriter(
        Path path,
        List<String> inputColumnNames,
        StorageFormat storageFormat,
        Properties schema,
        DataSize estimatedWriterSystemMemoryUsage,
        JobConf conf,
        TypeManager typeManager,
        ConnectorSession session)
{
    this.path = requireNonNull(path, "path is null");
    this.conf = requireNonNull(conf, "conf is null");

    // existing tables may have columns in a different order
    List<String> fileColumnNames = getColumnNames(schema);
    List<Type> fileColumnTypes = getColumnTypes(schema).stream()
            .map(hiveType -> hiveType.getType(typeManager))
            .collect(toList());

    fieldCount = fileColumnNames.size();

    String serDe = storageFormat.getSerDe();
    serializer = initializeSerializer(conf, schema, serDe);
    recordWriter = createRecordWriter(path, conf, schema, storageFormat.getOutputFormat(), session);

    List<ObjectInspector> objectInspectors = getRowColumnInspectors(fileColumnTypes);
    tableInspector = getStandardStructObjectInspector(fileColumnNames, objectInspectors);

    // reorder (and possibly reduce) struct fields to match input
    structFields = ImmutableList.copyOf(inputColumnNames.stream()
            .map(tableInspector::getStructFieldRef)
            .collect(toImmutableList()));

    row = tableInspector.create();

    setters = new FieldSetter[structFields.size()];
    for (int i = 0; i < setters.length; i++) {
        setters[i] = createFieldSetter(tableInspector, row, structFields.get(i), fileColumnTypes.get(structFields.get(i).getFieldID()));
    }

    this.estimatedWriterSystemMemoryUsage = estimatedWriterSystemMemoryUsage.toBytes();
}
 

@Inject
public SqlTaskManager(
        LocalExecutionPlanner planner,
        LocationFactory locationFactory,
        TaskExecutor taskExecutor,
        SplitMonitor splitMonitor,
        NodeInfo nodeInfo,
        LocalMemoryManager localMemoryManager,
        TaskManagementExecutor taskManagementExecutor,
        TaskManagerConfig config,
        NodeMemoryConfig nodeMemoryConfig,
        LocalSpillManager localSpillManager,
        NodeSpillConfig nodeSpillConfig,
        GcMonitor gcMonitor)
{
    requireNonNull(nodeInfo, "nodeInfo is null");
    requireNonNull(config, "config is null");
    infoCacheTime = config.getInfoMaxAge();
    clientTimeout = config.getClientTimeout();

    DataSize maxBufferSize = config.getSinkMaxBufferSize();

    taskNotificationExecutor = newFixedThreadPool(config.getTaskNotificationThreads(), threadsNamed("task-notification-%s"));
    taskNotificationExecutorMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) taskNotificationExecutor);

    this.taskManagementExecutor = requireNonNull(taskManagementExecutor, "taskManagementExecutor cannot be null").getExecutor();
    this.driverYieldExecutor = newScheduledThreadPool(config.getTaskYieldThreads(), threadsNamed("task-yield-%s"));

    SqlTaskExecutionFactory sqlTaskExecutionFactory = new SqlTaskExecutionFactory(taskNotificationExecutor, taskExecutor, planner, splitMonitor, config);

    this.localMemoryManager = requireNonNull(localMemoryManager, "localMemoryManager is null");
    DataSize maxQueryUserMemoryPerNode = nodeMemoryConfig.getMaxQueryMemoryPerNode();
    DataSize maxQueryTotalMemoryPerNode = nodeMemoryConfig.getMaxQueryTotalMemoryPerNode();
    DataSize maxQuerySpillPerNode = nodeSpillConfig.getQueryMaxSpillPerNode();

    DataSize maxQueryMemoryPerNode = nodeMemoryConfig.getMaxQueryMemoryPerNode();
    queryMaxMemoryPerNode = maxQueryMemoryPerNode.toBytes();
    queryMaxTotalMemoryPerNode = maxQueryMemoryPerNode.toBytes();

    queryContexts = CacheBuilder.newBuilder().weakValues().build(CacheLoader.from(
            queryId -> createQueryContext(queryId, localMemoryManager, localSpillManager, gcMonitor, maxQueryUserMemoryPerNode, maxQueryTotalMemoryPerNode, maxQuerySpillPerNode)));

    tasks = CacheBuilder.newBuilder().build(CacheLoader.from(
            taskId -> createSqlTask(
                    taskId,
                    locationFactory.createLocalTaskLocation(taskId),
                    nodeInfo.getNodeId(),
                    queryContexts.getUnchecked(taskId.getQueryId()),
                    sqlTaskExecutionFactory,
                    taskNotificationExecutor,
                    sqlTask -> {
                        finishedTaskStats.merge(sqlTask.getIoStats());
                        return null;
                    },
                    maxBufferSize,
                    failedTasks)));
}
 

public SpillSpaceTracker(DataSize maxSize)
{
    requireNonNull(maxSize, "maxSize is null");
    maxBytes = maxSize.toBytes();
    currentBytes = 0;
}
 

/**
 * @return a result with at least one page if we have pages in buffer, empty result otherwise
 */
private synchronized BufferResult processRead(long sequenceId, DataSize maxSize)
{
    // When pages are added to the partition buffer they are effectively
    // assigned an id starting from zero. When a read is processed, the
    // "token" is the id of the page to start the read from, so the first
    // step of the read is to acknowledge, and drop all pages up to the
    // provided sequenceId.  Then pages starting from the sequenceId are
    // returned with the sequenceId of the next page to read.
    //
    // Since the buffer API is asynchronous there are a number of problems
    // that can occur our of order request (typically from retries due to
    // request failures):
    // - Request to read pages that have already been acknowledged.
    //   Simply, send an result with no pages and the requested sequenceId,
    //   and since the client has already acknowledge the pages, it will
    //   ignore the out of order response.
    // - Request to read after the buffer has been destroyed.  When the
    //   buffer is destroyed all pages are dropped, so the read sequenceId
    //   appears to be off the end of the queue.  Normally a read past the
    //   end of the queue would be be an error, but this specific case is
    //   detected and handled.  The client is sent an empty response with
    //   the finished flag set and next token is the max acknowledged page
    //   when the buffer is destroyed.
    //

    // if request is for pages before the current position, just return an empty result
    if (sequenceId < currentSequenceId.get()) {
        return emptyResults(taskInstanceId, sequenceId, false);
    }

    // if this buffer is finished, notify the client of this, so the client
    // will destroy this buffer
    if (pages.isEmpty() && noMorePages) {
        return emptyResults(taskInstanceId, currentSequenceId.get(), true);
    }

    // if request is for pages after the current position, there is a bug somewhere
    // a read call is always proceeded by acknowledge pages, which
    // will advance the sequence id to at least the request position, unless
    // the buffer is destroyed, and in that case the buffer will be empty with
    // no more pages set, which is checked above
    verify(sequenceId == currentSequenceId.get(), "Invalid sequence id");

    // read the new pages
    long maxBytes = maxSize.toBytes();
    List<SerializedPage> result = new ArrayList<>();
    long bytes = 0;

    for (SerializedPageReference page : pages) {
        bytes += page.getRetainedSizeInBytes();
        // break (and don't add) if this page would exceed the limit
        if (!result.isEmpty() && bytes > maxBytes) {
            break;
        }
        result.add(page.getSerializedPage());
    }
    return new BufferResult(taskInstanceId, sequenceId, sequenceId + result.size(), false, result);
}
 

public BufferResult getPages(long sequenceId, DataSize maxSize)
{
    // if location is complete return GONE
    if (completed.get() && serializedPages.isEmpty()) {
        return BufferResult.emptyResults(TASK_INSTANCE_ID, token.get(), true);
    }

    assertEquals(sequenceId, token.get(), "token");

    // wait for a single page to arrive
    SerializedPage serializedPage = null;
    try {
        serializedPage = serializedPages.poll(10, TimeUnit.MILLISECONDS);
    }
    catch (InterruptedException e) {
        Thread.currentThread().interrupt();
    }

    // if no page, return NO CONTENT
    if (serializedPage == null) {
        return BufferResult.emptyResults(TASK_INSTANCE_ID, token.get(), false);
    }

    // add serializedPages up to the size limit
    List<SerializedPage> responsePages = new ArrayList<>();
    responsePages.add(serializedPage);
    long responseSize = serializedPage.getSizeInBytes();
    while (responseSize < maxSize.toBytes()) {
        serializedPage = serializedPages.poll();
        if (serializedPage == null) {
            break;
        }
        responsePages.add(serializedPage);
        responseSize += serializedPage.getSizeInBytes();
    }

    // update sequence id
    long nextToken = token.get() + responsePages.size();

    BufferResult bufferResult = new BufferResult(TASK_INSTANCE_ID, token.get(), nextToken, false, responsePages);
    token.set(nextToken);

    return bufferResult;
}