%ParallelPoolExecutor
% Base class for Parallel pool implementations of the
% PartitionedArrayExecutor interface.
%
% This class has two responsibilities:
% 1. This provides the high level logic for evaluating a tall array
%    expression using StageTask intermediate representation. Child classes
%    only need to provide the setup, teardown and execution of running a
%    StageTask remotely.
% 2. This provides the implementation for metadata questions of
%    PartitionedArrayExecutor in terms of pool (E.G. is this datastore
%    supported, how to resolve a partitioning based on the pool).
%
% Child classes need to implement the following methods:
%            buildStageTasks: Build the stage tasks to evaluate from the
%                             given taskGraph.
%   executeStageTaskRemotely: Execute one stage task remotely.
%                    getName: Return the name of the execution environment.


%   Copyright 2016-2020 The MathWorks, Inc.

classdef (Abstract) ParallelPoolExecutor < matlab.bigdata.internal.executor.PartitionedArrayExecutor
    properties (SetAccess = immutable)
        % The parallel pool to use for execution.
        Pool
    end
    
    % Methods to be overridden by child classes.
    methods (Access = protected, Abstract)
        % Convert the input task graph into an array of independent tasks
        % that can be executed one by one.
        %
        % This must return:
        %  - stageTasks: An array of independent ExecutionTask instances.
        %  - broadcastMap: The map that will hold broadcasts visible to the
        %                  client after evaluation.
        %  - additionalArguments: Cell array of additional arguments to
        %                         pass on to executeStageTaskRemotely.
        [stageTasks, broadcastMap, additionalArguments] ...
            = buildStageTasks(obj, taskGraph, outputHandlers, progressReporter);
        
        % Execute one stage task using the pool. This performs all
        % operations that require the pool itself.
        readFailureSummary = executeStageTaskRemotely(obj, ...
            stageTask, broadcastMap, progressReporter, varargin)
        
        % Get a user visible string that represents this execution
        % environment.
        name = getName(obj);
    end
    
    % Methods overridden in the PartitionedArrayExecutor interface.
    methods
        %EXECUTEWITHHANDLER Execute the provided graph of tasks redirecting
        %  all output to the given output handler.
        function readFailureSummary = executeWithHandler(obj, taskGraph, outputHandlers)
            import matlab.bigdata.internal.executor.ProgressReporter;
            progressReporter = ProgressReporter.getCurrent();
            
            [stageTasks, broadcastMap, additionalArguments] ...
                = obj.buildStageTasks(taskGraph, outputHandlers, progressReporter);
            
            for stageTask = stageTasks(:)'
                executionStrategy = stageTask.ExecutionPartitionStrategy;
                if executionStrategy.IsDatastorePartitioning
                    obj.errorIfDatastoreNotSupported(executionStrategy.Datastore, obj.getName, obj.Pool);
                end
            end
            
            import matlab.bigdata.internal.executor.ReadFailureAccumulator
            readFailureAccumulator = ReadFailureAccumulator();
            
            numTasks = numel(stageTasks) - iCountNumPureBroadcastStages(stageTasks);
            numPasses = iCountNumPassStages(stageTasks);
            
            progressReporter.startOfExecution(obj.getName(), numTasks, numPasses);
            for ii = 1:numel(stageTasks)
                isFullPass = stageTasks(ii).IsPass;
                
                if iIsPureBroadcast(stageTasks(ii))
                    % Data that is in broadcast state both exists in the
                    % local MATLAB and is small. The overhead is too high
                    % to execute this on the pool so we do serial
                    % execution.
                    obj.executeStageTaskLocally(stageTasks(ii), broadcastMap);
                else
                    progressReporter.startOfNextTask(isFullPass);
                    taskReadFailureSummary = obj.executeStageTaskRemotely(...
                        stageTasks(ii), broadcastMap, progressReporter, additionalArguments{:});
                    readFailureAccumulator.append(taskReadFailureSummary);
                    progressReporter.endOfTask();
                    if isFullPass
                        obj.incrementTotalNumPasses();
                    end
                end
                
                % If this stage task generated broadcast output to be sent
                % to the client, we can perform that action now.
                outputBroadcasts = stageTasks(ii).OutputBroadcasts;
                outputBroadcasts = outputBroadcasts(ismember(outputBroadcasts, taskGraph.OutputTasks));
                for outputBroadcast = outputBroadcasts(:)'
                    outputHandlers.handleBroadcastOutput(...
                        outputBroadcast.Id, broadcastMap.get(outputBroadcast.Id));
                end
            end
            progressReporter.endOfExecution();
            readFailureSummary = readFailureAccumulator.Summary;
        end
        
        % Count the number of passes required to execute the provided graph
        % of tasks.
        function numPasses = countNumPasses(obj, taskGraph) %#ok<STOUT,INUSD>
            assert(false, "Assertion Failed: countNumPasses is obsolete");
        end
        
        %NUMPARTITIONS Retrieve the number of partitions for the given
        %  partition strategy.
        function n = numPartitions(obj, partitionStrategy)
            partitionStrategy = obj.resolvePartitionStrategy(partitionStrategy);
            n = numpartitions(partitionStrategy);
        end
        
        % Check whether this executor is valid.
        function tf = checkIsValid(obj)
            % This means that a given tall array attached to this executor
            % will no longer be valid once the pool is deleted, or closed
            % due to idle timeout.
            tf = isvalid(obj.Pool) && obj.Pool.Connected;
        end
        
        %CHECKDATASTORESUPPORT Check whether the provided datastore is supported.
        function checkDatastoreSupport(obj, inputds)
            parallel.internal.bigdata.ParallelPoolExecutor.errorIfDatastoreNotSupported(inputds, obj.getName(), obj.Pool);
        end
        
        %KEEPALIVE Notify to the executor that operations have just been
        %performed and it should reset any idle timeouts.
        function keepAlive(obj)
            obj.Pool.IdleTimeout = obj.Pool.IdleTimeout;
        end
    end
    
    % Helper methods for checking if a datastore is supported before
    % attempting to open an execution environment.
    methods (Static, Hidden)
        %ERRORIFDATASTORENOTSUPPORT Issue an error if the provided
        % datastore is not supported by parallel pool.
        function errorIfDatastoreNotSupported(inputds, envName, poolObj)
            if nargin < 2 || isempty(envName)
                envName = getString(message('parallel:bigdata:ParallelPoolName'));
            end
            matlab.io.datastore.internal.validateParallelBigdata(...
                inputds, "Partitionable", envName);
            
            % Check for selection of supported datastores for
            % ThreadPool.
            if nargin > 2 && isa(poolObj, "parallel.ThreadPool")
                % Unwrap the underlying datastore if required.
                if isa(inputds, "matlab.io.datastore.internal.FrameworkDatastore")
                    inputds = inputds.Datastore;
                end
                if isa(inputds, "matlab.io.datastore.TransformedDatastore")
                    inputds = inputds.UnderlyingDatastore;
                end
                % Supported datastores: MemoryDatastore,
                % TabularTextDatastore, and custom implemented datastores.
                if ~any([isa(inputds, "matlab.bigdata.internal.MemoryDatastore"), ...
                        isa(inputds, "matlab.io.datastore.TabularTextDatastore"), ...
                        ~startsWith(which(class(inputds)), toolboxdir(''))])
                    datastoreFullName = class(inputds);
                    datastoreSimpleName = strsplit(datastoreFullName, '.');
                    datastoreSimpleName = datastoreSimpleName{end};
                    error(message("parallel:lang:pool:DisallowedParallelFeature", datastoreSimpleName));
                end
            end
        end
        
        % Check if pool auto-open is enabled.
        function tf = isAutoCreateEnabled()
            ps = parallel.Settings;
            tf = ps.Pool.AutoCreate;
        end
    end
    
    % Methods available for child classes to invoke.
    methods (Access = protected)
        function obj = ParallelPoolExecutor(pool)
            % Constructor for child classes.
            obj.Pool = pool;
        end
        
        function partitionStrategy = resolvePartitionStrategy(obj, partitionStrategy)
            % Resolve a partition strategy to a version that knows the
            % exact number of partitions.
            
            % This balances performance against load balancing:-
            %  * In general case, limit number of partitions to 3 * pool
            %    size to provide some load balancing by oversubscribing,
            %    but not so much as to add excessive overheads.
            %  * Except in case of Any-to-Any communication, where we only
            %    use 1 * pool size as the overheads are greater in
            %    comparison to the likelihood of imbalanced partitions.
            defaultNumPartitionsHint = obj.Pool.NumWorkers;
            maxNumPartitionsHint = 3 * defaultNumPartitionsHint;
            % For process based pools, enable the option to do
            % prepartitioning to avoid sending full copies of datastores.
            doPartitioningFirst = isa(obj.Pool, 'parallel.ProcessPool');
            partitionStrategy = resolve(partitionStrategy, ...
                defaultNumPartitionsHint, maxNumPartitionsHint, doPartitioningFirst);
        end
        
        function [shuffleSortStrategy, writerFactory, readerFactory] ...
                = createMatArrayShuffleStorage(~, outputCommunicationType, outputPartitionStrategy, useSharedStorage)
            % Create a ShuffleSortStrategy that communicates the output of
            % a task into that task's OutputPartitionStrategy.
            
            import parallel.internal.bigdata.MatArrayShuffleSortStrategy;
            import parallel.internal.bigdata.RemoteShuffleSortStrategy;
            numOutputPartitions = numpartitions(outputPartitionStrategy);
            shuffleSortStrategy = MatArrayShuffleSortStrategy(numOutputPartitions, useSharedStorage);
            if ~useSharedStorage
                shuffleSortStrategy = RemoteShuffleSortStrategy(shuffleSortStrategy);
            end
            
            writerFactory = @(partition) shuffleSortStrategy.createMapOutputSink(...
                partition.PartitionIndex, iGetDefaultKey(partition, outputCommunicationType));
            
            if outputPartitionStrategy.IsBroadcast
                % When a given piece of intermediate data is in broadcast
                % state, it has only 1 partition. Every data processor that
                % requires this data must read from partition 1.
                broadcastPartitionIndex = 1;
                readerFactory = @(partition) shuffleSortStrategy.createReduceInputSource(broadcastPartitionIndex);
            else
                readerFactory = @(partition) shuffleSortStrategy.createReduceInputSource(partition.PartitionIndex);
            end
        end
        
        % Implementation of createShuffleStorage that uses Sqlite files as
        % intermediate storage type.
        function [shuffleSortStrategy, writerFactory, readerFactory] ...
                = createSqliteShuffleStorage(~, outputCommunicationType, outputPartitionStrategy, useSharedStorage)
            import matlab.bigdata.internal.executor.OutputCommunicationType;
            import parallel.internal.bigdata.RemoteShuffleSortStrategy;
            import parallel.internal.bigdata.SinkAdapterWriter;
            import parallel.internal.bigdata.SourceAdapterReader;
            import parallel.internal.bigdata.TrivialPartitioner;
            
            % We disable usage of worker ID by the shuffle sort strategy as
            % this option introduces non-deterministic ordering of data.
            % There are ways to use worker ID without introducing
            % non-deterministic ordering, this is a potential optimization.
            import parallel.internal.mapreduce.SqliteRemoteCopyShuffleSortStrategy;
            import parallel.internal.mapreduce.SqliteShuffleSortStrategy;
            import parallel.internal.bigdata.RemoteShuffleSortStrategy;
            useWorkerId = false;
            numOutputPartitions = numpartitions(outputPartitionStrategy);
            if useSharedStorage
                shuffleSortStrategy = SqliteShuffleSortStrategy(numOutputPartitions, TrivialPartitioner(), useWorkerId);
            else
                shuffleSortStrategy = RemoteShuffleSortStrategy(...
                    SqliteRemoteCopyShuffleSortStrategy(numOutputPartitions, TrivialPartitioner(), useWorkerId));
            end
            
            writerFactory = @(partition) SinkAdapterWriter(...
                shuffleSortStrategy.createMapOutputSink(partition.PartitionIndex), ...
                iGetDefaultKey(partition, outputCommunicationType));
            
            if outputPartitionStrategy.IsBroadcast
                % When a given piece of intermediate data is in broadcast
                % state, it has only 1 partition. Every data processor that
                % requires this data must read from partition 1.
                broadcastPartitionIndex = 1;
                readerFactory = @(partition) SourceAdapterReader(...
                    shuffleSortStrategy.createReduceInputSource(broadcastPartitionIndex));
            else
                readerFactory = @(partition) SourceAdapterReader(...
                    shuffleSortStrategy.createReduceInputSource(partition.PartitionIndex));
            end
        end
    end
    
    methods (Access = private)
        function executeStageTaskLocally(obj, task, broadcastMap) %#ok<INUSL>
            % Execute the provided independent stage serially.
            assert(iIsPureBroadcast(task));
            
            partitionStrategy = task.ExecutionPartitionStrategy;
            partitionIndex = 1;
            numExecutorPartitions = 1;
            partition = partitionStrategy.createPartition(partitionIndex, numExecutorPartitions);
            partitionContext = matlab.bigdata.internal.executor.PartitionContext(...
                partition, broadcastMap);
            
            dataProcessor = feval(task.DataProcessorFactory, partitionContext);
            while ~dataProcessor.IsFinished
                process(dataProcessor, false(0));
            end
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function defaultKey = iGetDefaultKey(partition, outputCommunicationType)
% For the provided partition and output communication type, get the default
% key that should be written to the intermediate data store.
import matlab.bigdata.internal.executor.OutputCommunicationType;

switch outputCommunicationType
    case OutputCommunicationType.Simple
        defaultKey = partition.PartitionIndex;
    case OutputCommunicationType.Broadcast
        defaultKey = 1;
    case OutputCommunicationType.AllToOne
        defaultKey = 1;
    case OutputCommunicationType.AnyToAny
        % In AnyToAny communication, the partition indices should be
        % specified by the data processor implementation underlying the
        % task.
        defaultKey = [];
end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Helper function that returns true if a given stage task executes in
% broadcast mode with broadcast output.
function tf = iIsPureBroadcast(stageTask)
tf = stageTask.ExecutionPartitionStrategy.IsBroadcast && isempty(stageTask.OutputShuffles);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Count the number of stages that have execution across the full
% data.
function numStages = iCountNumPassStages(stageTasks)
numStages = sum([stageTasks.IsPass]);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Count the number of stages that are pure broadcast.
function numStages = iCountNumPureBroadcastStages(stageTasks)
numStages = 0;
for ii = 1:numel(stageTasks)
    numStages = numStages + iIsPureBroadcast(stageTasks(ii));
end
end
