io.trino.operator.join.SpillingJoinProcessor Maven / Gradle / Ivy

Go to download
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.operator.join;

import com.google.common.io.Closer;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import io.trino.operator.WorkProcessor;
import io.trino.operator.join.DefaultPageJoiner.SavedRow;
import io.trino.operator.join.PageJoiner.PageJoinerFactory;
import io.trino.spi.Page;
import io.trino.spiller.PartitioningSpillerFactory;
import jakarta.annotation.Nullable;

import java.io.IOException;
import java.util.Iterator;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.function.Supplier;

import static com.google.common.collect.Iterators.singletonIterator;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static io.airlift.concurrent.MoreFutures.getDone;
import static java.util.Collections.emptyIterator;
import static java.util.Objects.requireNonNull;

public class SpillingJoinProcessor
        implements WorkProcessor.Process>
{
    private final Runnable afterClose;
    private final OptionalInt lookupJoinsCount;
    private final boolean waitForBuild;
    private final LookupSourceFactory lookupSourceFactory;
    private final ListenableFuture lookupSourceProvider;
    private final PageJoinerFactory pageJoinerFactory;
    private final PageJoiner sourcePagesJoiner;
    private final WorkProcessor joinedSourcePages;

    private boolean closed;

    @Nullable
    private ListenableFuture>> partitionedConsumption;
    @Nullable
    private Iterator>> lookupPartitions;
    @Nullable
    private PartitionedConsumption.Partition> previousPartition;
    @Nullable
    private ListenableFuture> previousPartitionLookupSource;

    public SpillingJoinProcessor(
            Runnable afterClose,
            OptionalInt lookupJoinsCount,
            boolean waitForBuild,
            LookupSourceFactory lookupSourceFactory,
            ListenableFuture lookupSourceProvider,
            PartitioningSpillerFactory partitioningSpillerFactory,
            PageJoinerFactory pageJoinerFactory,
            WorkProcessor sourcePages)
    {
        this.afterClose = requireNonNull(afterClose, "afterClose is null");
        this.lookupJoinsCount = requireNonNull(lookupJoinsCount, "lookupJoinsCount is null");
        this.waitForBuild = waitForBuild;
        this.lookupSourceFactory = requireNonNull(lookupSourceFactory, "lookupSourceFactory is null");
        this.lookupSourceProvider = requireNonNull(lookupSourceProvider, "lookupSourceProvider is null");
        this.pageJoinerFactory = requireNonNull(pageJoinerFactory, "pageJoinerFactory is null");
        sourcePagesJoiner = pageJoinerFactory.getPageJoiner(
                lookupSourceProvider,
                Optional.of(partitioningSpillerFactory),
                emptyIterator());
        joinedSourcePages = sourcePages.transform(sourcePagesJoiner);
    }

    public void close()
    {
        if (closed) {
            return;
        }
        closed = true;

        try (Closer closer = Closer.create()) {
            // `afterClose` must be run last.
            // Closer is documented to mimic try-with-resource, which implies close will happen in reverse order.
            closer.register(afterClose::run);

            closer.register(sourcePagesJoiner);
            sourcePagesJoiner.getSpiller().ifPresent(closer::register);
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public WorkProcessor.ProcessState> process()
    {
        // wait for build side to be completed before fetching any probe data
        // TODO: fix support for probe short-circuit: https://github.com/trinodb/trino/issues/3957
        if (waitForBuild && !lookupSourceProvider.isDone()) {
            return WorkProcessor.ProcessState.blocked(asVoid(lookupSourceProvider));
        }

        if (!joinedSourcePages.isFinished()) {
            return WorkProcessor.ProcessState.ofResult(joinedSourcePages);
        }

        if (partitionedConsumption == null) {
            partitionedConsumption = lookupSourceFactory.finishProbeOperator(lookupJoinsCount);
            return WorkProcessor.ProcessState.blocked(asVoid(partitionedConsumption));
        }

        if (lookupPartitions == null) {
            lookupPartitions = getDone(partitionedConsumption).beginConsumption();
        }

        if (previousPartition != null) {
            // If we had no rows for the previous spill partition, we would finish before it is unspilled.
            // Partition must be loaded before it can be released. // TODO remove this constraint
            if (!previousPartitionLookupSource.isDone()) {
                return WorkProcessor.ProcessState.blocked(asVoid(previousPartitionLookupSource));
            }

            previousPartition.release();
            previousPartition = null;
            previousPartitionLookupSource = null;
        }

        if (!lookupPartitions.hasNext()) {
            close();
            return WorkProcessor.ProcessState.finished();
        }

        PartitionedConsumption.Partition> partition = lookupPartitions.next();
        previousPartition = partition;
        previousPartitionLookupSource = partition.load();

        return WorkProcessor.ProcessState.ofResult(joinUnspilledPages(partition));
    }

    private static  ListenableFuture asVoid(ListenableFuture future)
    {
        return Futures.transform(future, v -> null, directExecutor());
    }

    private WorkProcessor joinUnspilledPages(PartitionedConsumption.Partition> partition)
    {
        int partitionNumber = partition.number();
        WorkProcessor unspilledInputPages = WorkProcessor.fromIterator(sourcePagesJoiner.getSpiller()
                .map(spiller -> spiller.getSpilledPages(partitionNumber))
                .orElse(emptyIterator()));
        Iterator savedRow = Optional.ofNullable(sourcePagesJoiner.getSpilledRows().remove(partitionNumber))
                .map(row -> (Iterator) singletonIterator(row))
                .orElse(emptyIterator());

        ListenableFuture unspilledLookupSourceProvider = Futures.transform(
                partition.load(),
                supplier -> new StaticLookupSourceProvider(supplier.get()),
                directExecutor());

        return unspilledInputPages.transform(pageJoinerFactory.getPageJoiner(
                unspilledLookupSourceProvider,
                Optional.empty(),
                savedRow));
    }
}