%ReduceTask
% Class for holding all the details necessary for performing a Reduce Task.
% This should be both serializable and immutable.

%   Copyright 2014-2019 The MathWorks, Inc.

classdef (Sealed, Hidden) ReduceTask < parallel.internal.mapreduce.AbstractTask
    properties (GetAccess = private, SetAccess = immutable)
        ReduceTaskIDFunctor; % Number representing the current reduce task index.
        
        ReduceFunction; % Function handle to the customers Reduce function.
        
        OutputKeyValueWriterFactory; % A function handle of the form fcn(reduceTaskId) that creates an output writer.
    end
    
    properties (Transient, GetAccess = private, SetAccess = immutable)
        RunOnAllWorkers = false;
    end
    
    methods (Static)
        % Creates a ReduceTask that will run on one worker
        % Syntax:
        % ReduceTask.createRunOnSingle(reduceTaskIDFunctor, reduceFunction, outputKeyValueWriterFactory, shuffleSortStrategy)
        % ReduceTask.createRunOnSingle(reduceTaskIDFunctor, outputKeyValueWriterFactory, remoteMapReduceDetails)
        function obj = createRunOnSingle(reduceTaskIDFunctor, varargin)
            import parallel.internal.mapreduce.ReduceTask;
            narginchk(3,4);
            if isnumeric(reduceTaskIDFunctor)
                reduceTaskIDFunctor = @(~) reduceTaskIDFunctor;
            end
            obj = ReduceTask(false, reduceTaskIDFunctor, varargin{:});
        end
        
        % Creates a ReduceTask that will run on all workers        
        % Syntax:
        % ReduceTask.createRunOnAll(reduceTaskIDFunctor, reduceFunction, outputKeyValueWriterFactory, shuffleSortStrategy)
        % ReduceTask.createRunOnAll(reduceTaskIDFunctor, outputKeyValueWriterFactory, remoteMapReduceDetails)
        function obj = createRunOnAll(reduceTaskIDFunctor, varargin)
            import parallel.internal.mapreduce.ReduceTask;
            narginchk(3,4);
            if isnumeric(reduceTaskIDFunctor)
                reduceTaskIDFunctor = @(~) reduceTaskIDFunctor;
            end
            obj = ReduceTask(true, reduceTaskIDFunctor, varargin{:});
        end
    end
    
    methods (Access = private)
        % Use one of the static create methods to build a ReduceTask
        % Syntax:
        % ReduceTask(runOnAllWorkers, reduceTaskIDFunctor, reduceFunction, outputKeyValueWriterFactory, shuffleSortStrategy)
        % ReduceTask(runOnAllWorkers, reduceTaskIDFunctor, outputKeyValueWriterFactory, remoteMapReduceDetails, runOnAllWorkers)
        function obj = ReduceTask(runOnAllWorkers, reduceTaskIDFunctor, varargin)
            if nargin == 4
                reduceFunction = [];
                outputKeyValueWriterFactory = varargin{1};
                baseInputs = varargin(2);
            elseif nargin == 5
                reduceFunction = varargin{1};
                outputKeyValueWriterFactory = varargin{2};
                baseInputs = varargin(3);
            else
                assert(false, 'ReduceTask requires 4 or 5 inputs');
            end
            
            obj = obj@parallel.internal.mapreduce.AbstractTask(baseInputs{:});
            obj.RunOnAllWorkers = runOnAllWorkers;
            obj.ReduceTaskIDFunctor = reduceTaskIDFunctor;
            obj.ReduceFunction = reduceFunction;
            obj.OutputKeyValueWriterFactory = outputKeyValueWriterFactory;
        end
    end
    
    methods
        function results = run(obj)
            import parallel.internal.mapreduce.ReduceTaskResult;
            reduceTaskIds = obj.ReduceTaskIDFunctor(iGetWorkerId());
            results = cell(size(reduceTaskIds));
            for ii = 1:numel(reduceTaskIds)
                results{ii} = obj.reduce(reduceTaskIds(ii));
            end
            results = vertcat(results{:});
            if isempty(results)
                results = ReduceTaskResult();
            end
        end
    end
    
    methods (Access = private)
        % Setup a sink that represents the output folder.
        function [keyValueOutputStore, reduceOutputSink] = pSetupOutputSink(obj, reduceTaskID)
            import parallel.mapreduce.KeyValueOutputStore;
            try
                reduceOutputSink = obj.OutputKeyValueWriterFactory(reduceTaskID);
                keyValueOutputStore = KeyValueOutputStore.hBuild(reduceOutputSink);
            catch err
                oErr = MException(message('parallel:mapreduce:ParallelOutputError'));
                % Adding the original error as a cause will only show the
                % first frame, which should contain a user visible error as
                % to the reason why this failed.
                oErr = oErr.addCause(err);
                throw(oErr);
            end
        end
        
        % Runs the given reduce task.
        function results = reduce(obj, reduceTaskID)
            import parallel.mapreduce.ValueIterator;
            
            shuffleSortStrategy = obj.getShuffleSortStrategy();
            reduceInputSource = shuffleSortStrategy.createReduceInputSource(reduceTaskID);
            
            reduceFunction = obj.getReduceFunction();
            try
                [keyValueOutputStore, reduceOutputSink] = obj.pSetupOutputSink(reduceTaskID);
                
                while reduceInputSource.hasData()
                    key = reduceInputSource.nextKey();
                    iterator = ValueIterator.hBuild(reduceInputSource);
                    
                    reduceFunction(key, iterator, keyValueOutputStore);
                    
                    reduceInputSource.discardRemainingDataForCurrentKey();
                end
                obj.pCommitOutputSink(reduceOutputSink);
                
                results = reduceOutputSink.getTaskOutputMetadata();
            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, 'ReduceTask');
                    pErr.throw();
                end
            end
        end        
    end
    
    methods (Access = private)
        % Helper function that returns the reduce function.
        function reduceFunction = getReduceFunction(obj)
            if isempty(obj.MapReduceRemoteCache)
                reduceFunction = obj.ReduceFunction;
            else
                reduceFunction = obj.MapReduceRemoteCache.getReduceFunction();
            end
            
            % The reduce function handle is empty only if we could not load it.
            if isempty(reduceFunction)
                error(message('parallel:mapreduce:InvalidReduceFunction'));
            end
        end
    end
    
    methods (Access = private, Static)
        % Commit the output sink.
        function pCommitOutputSink(reduceOutputSink)
            try
                reduceOutputSink.commit();
            catch err
                oErr = MException(message('parallel:mapreduce:ParallelOutputError'));
                % Adding the original error as a cause will only show the
                % first frame, which should contain a user visible error as
                % to the reason why this failed.
                oErr = oErr.addCause(err);
                throw(oErr);
            end
        end
    end
    
    methods (Hidden)
        % Launch a task on the given pool.
        % Visitor will be an object/struct with both a runOnSingle and a 
        % runOnAll method. This should call one of those and return the
        % result.
        function future = hLaunchTask(obj, visitor)
            if obj.RunOnAllWorkers
                future = visitor.runOnAll(obj);
            else
                future = visitor.runOnSingle(obj);
            end
        end
    end
end

function id = iGetWorkerId()
sessionObj = parallel.internal.pool.workerSession();
if isa(sessionObj, 'parallel.internal.pool.AbstractWorker')
    id = sessionObj.WorkerIndex;
else
    id = 1;
end
end
