com.facebook.presto.spark.classloader_interface.PrestoSparkTaskProcessor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of presto-spark-launcher Show documentation
There is a newer version: 0.290
Show newest version
/*
 * 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 com.facebook.presto.spark.classloader_interface;

import org.apache.spark.TaskContext;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.util.CollectionAccumulator;
import scala.Tuple2;
import scala.collection.Iterator;

import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import static java.util.Collections.emptyMap;
import static java.util.Objects.requireNonNull;

public class PrestoSparkTaskProcessor
        implements Serializable
{
    private final PrestoSparkTaskExecutorFactoryProvider taskExecutorFactoryProvider;
    private final SerializedPrestoSparkTaskDescriptor serializedTaskDescriptor;
    private final CollectionAccumulator taskInfoCollector;
    private final CollectionAccumulator shuffleStatsCollector;
    // fragmentId -> Broadcast
    private final Map> broadcastInputs;
    private final Class outputType;

    public PrestoSparkTaskProcessor(
            PrestoSparkTaskExecutorFactoryProvider taskExecutorFactoryProvider,
            SerializedPrestoSparkTaskDescriptor serializedTaskDescriptor,
            CollectionAccumulator taskInfoCollector,
            CollectionAccumulator shuffleStatsCollector,
            Map> broadcastInputs,
            Class outputType)
    {
        this.taskExecutorFactoryProvider = requireNonNull(taskExecutorFactoryProvider, "taskExecutorFactoryProvider is null");
        this.serializedTaskDescriptor = requireNonNull(serializedTaskDescriptor, "serializedTaskDescriptor is null");
        this.taskInfoCollector = requireNonNull(taskInfoCollector, "taskInfoCollector is null");
        this.shuffleStatsCollector = requireNonNull(shuffleStatsCollector, "shuffleStatsCollector is null");
        this.broadcastInputs = new HashMap<>(requireNonNull(broadcastInputs, "broadcastInputs is null"));
        this.outputType = requireNonNull(outputType, "outputType is null");
    }

    public Iterator> process(
            Iterator serializedTaskSources,
            // fragmentId -> Iterator<[partitionId, page]>
            Map>> shuffleInputs)
    {
        return taskExecutorFactoryProvider.get().create(
                TaskContext.get().partitionId(),
                TaskContext.get().attemptNumber(),
                serializedTaskDescriptor,
                serializedTaskSources,
                new PrestoSparkJavaExecutionTaskInputs(shuffleInputs, broadcastInputs, emptyMap()),
                taskInfoCollector,
                shuffleStatsCollector,
                outputType);
    }

    /**
     * Overloaded member method that processes a native task.
     */
    public Iterator> process(
            Iterator serializedTaskSources,
            // fragmentId -> Iterator<[partitionId, page]>
            Map shuffleReadDescriptors,
            Optional shuffleWriteDescriptor)
    {
        return taskExecutorFactoryProvider.getNative().create(
                TaskContext.get().partitionId(),
                TaskContext.get().attemptNumber(),
                serializedTaskDescriptor,
                serializedTaskSources,
                new PrestoSparkNativeTaskInputs(shuffleReadDescriptors, shuffleWriteDescriptor, broadcastInputs),
                taskInfoCollector,
                shuffleStatsCollector,
                outputType);
    }
}