%PartitionedArrayTask
% Implementation of parallel.internal.mapreduce.AbstractTask that executes
% a task generated by matlab.bigdata.internal.convertToIndependentTasks for
% either a single partition or a collection of tasks that have been
% explicitly assigned to workers.
%

%   Copyright 2016-2019 The MathWorks, Inc.

classdef (Sealed) PartitionedArrayTask < parallel.internal.mapreduce.AbstractTask
    properties (SetAccess = immutable)
        % The factory to use to generate the underlying DataProcessor
        % instance. This is allowed to be a parallel.pool.Constant, in
        % which case the factory is stored in the Value property.
        DataProcessorFactory (1,1)
        
        % A function handle that maps worker ID to the partition indices
        % to be evaluated by that worker.
        PartitionIndexFunctor (1,1)
        
        % The number of partitions over the entire ExecutionTask.
        NumPartitions (1,1) double
        
        % The underlying partitioning strategy.
        PartitionStrategy
        
        % The underlying partition that will be taken if there's no
        % caching.
        Partition
        
        % Whether to use parfevalOnAll to complete this task.
        UseParfevalOnAll (1,1) logical = false
        
        % A function handle that is called on completion of each partition
        % with the partition index of that partition.
        PartitionCompleteCallback (1,1) = @(partitionIndex, numPartitions) []
        
        % A ParallelPoolBroadcastMap to be used by PartitionContext /
        % broadcast mechanism.
        ParallelBroadcastMap (1,1)
        
        % The maximum number of read failures allowed in the underlying
        % partition strategy.
        MaxNumReadFailures
    end
    
    methods
        % The main constructor.
        %
        % Syntax:
        %   obj = PartitionedArrayTask(factory, partitionIndex, numPartitions, partitionStrategy)
        %   generates a task for the single partition of index
        %   partitionIndex. This will be run by the first available worker.
        %
        %   obj = PartitionedArrayTask(factory, partitionIndexFunctor, numPartitions, partitionStrategy)
        %   generates a task that is run by all workers, where each worker
        %   runs all of the partitions of index partitionIndexFunctor(workerIndex).
        %
        %   obj = PartitionedArrayTask(..., partitionCompleteCallback) calls
        %   the given function handle on the worker on completion of every
        %   partition.
        %
        function obj = PartitionedArrayTask(dataProcessorFactory, partitionIndexFunctor, ...
                partitionStrategy, maxNumReadFailures, partitionCompleteCallback, broadcastMap)
            obj = obj@parallel.internal.mapreduce.AbstractTask([]);
            obj.DataProcessorFactory = dataProcessorFactory;
            if isnumeric(partitionIndexFunctor)
                obj.PartitionIndexFunctor = @(~)partitionIndexFunctor;
            else
                obj.PartitionIndexFunctor = partitionIndexFunctor;
                obj.UseParfevalOnAll = true;
            end
            
            if isa(partitionStrategy, 'matlab.bigdata.internal.executor.Partition')
                obj.Partition = partitionStrategy;
                obj.NumPartitions = partitionStrategy.NumPartitions;
                obj.PartitionStrategy = [];
            else
                obj.PartitionStrategy = partitionStrategy;
                obj.NumPartitions = numpartitions(partitionStrategy);
                obj.Partition = [];
            end
            
            obj.MaxNumReadFailures = maxNumReadFailures;
            obj.PartitionCompleteCallback = partitionCompleteCallback;
            obj.ParallelBroadcastMap = broadcastMap;
        end
    end
    
    % Overrides of parallel.internal.mapreduce.AbstractTask methods.
    methods
        % Execute the current task on the current worker.
        function readFailureSummary = run(obj)
            try
                readFailureAccumulator = matlab.bigdata.internal.executor.ReadFailureAccumulator(...
                    obj.MaxNumReadFailures);
                partitionIndices = feval(obj.PartitionIndexFunctor, iGetWorkerId());
                for partitionIndex = partitionIndices(:)'
                    obj.doRun(partitionIndex, readFailureAccumulator);
                    feval(obj.PartitionCompleteCallback, partitionIndex, obj.NumPartitions);
                end
                readFailureSummary = readFailureAccumulator.Summary;
            catch err
                % ParallelException shortens the stack trace such that only
                % frames from the map function are sent back to the client.
                % However, we do not want to hide UndefinedFunction as this
                % allows pool auto-attach to work.
                if strcmp(err.identifier, 'MATLAB:UndefinedFunction')
                    err.rethrow();
                else
                    pErr = ParallelException.hBuildFromRemoteException(err, 'MapTask');
                    pErr.throw();
                end
            end
        end
    end
    
    methods (Access = private)
        % Execute the provided partition index on the current worker.
        function doRun(obj, partitionIndex, readFailureAccumulator)
            if ~isempty(obj.PartitionStrategy)
                partition = obj.PartitionStrategy.createPartition(partitionIndex, obj.NumPartitions);
            else
                partition = obj.Partition;
            end
            partitionContext = matlab.bigdata.internal.executor.PartitionContext(...
                partition, obj.ParallelBroadcastMap);
            dataProcessorFactory = obj.DataProcessorFactory;
            if isa(dataProcessorFactory, 'parallel.pool.Constant')
                dataProcessorFactory = dataProcessorFactory.Value;
            end
            dataProcessor = feval(dataProcessorFactory, partitionContext);
            
            while ~dataProcessor.IsFinished
                process(dataProcessor, false(0));
            end
            readFailureAccumulator.append(partitionContext.getReadFailureSummary());
        end
    end
    
    % Overrides of parallel.internal.mapreduce.AbstractTask methods.
    methods (Hidden)
        % Launch the task on the pool using the provided visitor.
        function future = hLaunchTask(obj, visitor)
            % Each instance of this class is intended to run for a single
            % partition on a single worker.
            if obj.UseParfevalOnAll
                future = visitor.runOnAll(obj);
            else
                future = visitor.runOnSingle(obj);
            end
        end
    end
end

% Get the ID of the current worker.
function id = iGetWorkerId()
sessionObj = parallel.internal.pool.workerSession();
if isa(sessionObj, 'parallel.internal.pool.AbstractWorker')
    id = sessionObj.WorkerIndex;
else
    id = 1;
end
end
