%LocalFileKeyValueSink
% A helper class that wraps an instance of matlab.mapreduce.internal.KVSerializer
% as a sink.

%   Copyright 2014-2020 The MathWorks, Inc.

classdef (Sealed, Hidden) LocalFileKeyValueSink < parallel.internal.mapreduce.AbstractKeyValueSink
    
    properties (GetAccess = private, SetAccess = immutable)
        FinalPath; % Path of the output after commit.
        TempPath; % Temporary path of the data before commit.
        Buffer; % Buffer for vectorizing writes.
        Serializer; % Core MATLAB object that gives the implementation of how to write the output file.
    end
    
    properties (Access = private, Constant)
        BinaryMaxNumAdds = 1024; % MaxNumAdds for Binary Output.
        BinaryMaxContentSize = 32 * 1024 * 1024; % MaxContentSize for Binary Output.
    end
    
    properties (Access = private)
        KeyType = []; % The key type that has been written to this sink.
        ValueType = []; % The value type that has been written to this sink.
        Committed = false; % Flag specifying whether the output has been committed.
        CommittedOutputFiles = {}; % A list of committed output files.
    end

    methods
        % Create a Sink that writes to <basePath>/part<id>.<ext>
        % The files in a datastore do not need to conform to a given naming
        % scheme, so this is purely to ensure different workers do not
        % write to the same file.
        function this = LocalFileKeyValueSink(outputType, basePath, id, suffix)
            import matlab.mapreduce.internal.ReduceKVSerializer;
            import matlab.mapreduce.internal.ReduceTextKVSerializer;
            import parallel.internal.mapreduce.KeyValueBuffer;
            import parallel.internal.mapreduce.LocalFileKeyValueSink;
            
            
            this.FinalPath = basePath;
            this.TempPath = fullfile(basePath, sprintf('.%i', id));
            this.Buffer = KeyValueBuffer();
            
            [success, msg, msgid] = mkdir(this.TempPath);
            if ~success
                error(msgid, msg);
            end
            
            prefix = sprintf('result_r%i', id);
            switch (outputType)
                case 'Binary'
                    this.Serializer = ReduceKVSerializer(prefix, this.TempPath, suffix);
                    this.Buffer = KeyValueBuffer(LocalFileKeyValueSink.BinaryMaxNumAdds,...
                        LocalFileKeyValueSink.BinaryMaxContentSize);
                case 'TabularText'
                    this.Serializer = ReduceTextKVSerializer(prefix, this.TempPath, suffix);
                otherwise
                    assert(false)
            end
        end
        
        % ReduceTextKVSerializer is paired with it's own Processor for
        % checking data. This creates an instance of the correct Processor
        % class for the serializer.
        function processor = createKeyValueProcessor(this)
            import matlab.mapreduce.internal.KeyValueProcessor;
            import matlab.mapreduce.internal.TextKeyValueProcessor;
            if isa(this.Serializer, 'matlab.mapreduce.internal.ReduceTextKVSerializer')
                processor = TextKeyValueProcessor();
            else
                processor = KeyValueProcessor();
            end
        end
        
        % Add the given <key, value> pairs to the underlying data storage.
        % The keys and values input must be vectors of the same length.
        function add(this, keys, values)
            needsFlush = this.Buffer.add(keys, values);
            if needsFlush
                this.pFlushBuffer();
            end
        end
        
        % Commit any output to the final location. This is only called if
        % the task is successful. If a task is not successful, the sink
        % is expected to cleanup all side-effects when destroyed.
        function commit(this)
            assert(~this.Committed);
            if this.Buffer.NumAdds > 0
                this.pFlushBuffer();
            end
            
            ds = this.Serializer.constructDatastore();
            files = ds.Files;
            
            delete(ds);
            delete(this.Serializer);
            
            outputFiles = cell(size(files));
            for ii = 1:numel(files)
                [~, filename, ext] = fileparts(files{ii});
                filename = [filename ext]; %#ok<AGROW>
                outputFiles{ii} = fullfile(this.FinalPath, filename);
                
                % We use 'f' here to cope with output folders on Windows platform
                % that are attributed as read-only. This attribute on folders
                % is typically ignored by Windows. Having reached this
                % point, we have already made writes to the folder using
                % mkdir.
                movefile(files{ii}, outputFiles{ii}, 'f');
            end            
            this.Committed = true;
            this.CommittedOutputFiles = outputFiles;
        end
        
        % Cleanup any temporary resources
        function delete(this)
            if ~this.Committed
                delete(this.Serializer);
            end
            if exist(this.TempPath, 'dir') == 7
                [success, msg, ~] = rmdir(this.TempPath, 's');
                if ~success
                    dctSchedulerMessage( 1, ...
                        'Failed to rollback output file, original error: %s', msg);
                end
            end
        end
        
        % Flush the buffer, writing to the temporary file.
        function pFlushBuffer(this)
            [keys, values] = this.Buffer.get();
            this.Buffer.reset();
            
            % Transaction initialization
            if isempty(this.KeyType) && ~isempty(keys)
                exampleKey = keys(1);
                if iscellstr(exampleKey) || isstring(exampleKey)
                    exampleKey = exampleKey{1};
                end
                
                % We want to store the key/value class type for future error
                % checking. We are guaranteed that all data is the same type by
                % the checking that has happened in KeyValueOutputStore. This
                % is for comparing the data of two different workers.
                this.KeyType = class(exampleKey);
                this.ValueType = class(values);
            end
            
            this.Serializer.serialize(keys, values, Inf);
        end
        
        % Provide any data that is to be returned to the driver code for
        % MapReduce. This can be an arbitrary serializable value and will
        % be called once at the end of a Map or Reduce task to obtain the
        % return value of the run method.
        function results = getTaskOutputMetadata(this)
            import parallel.internal.mapreduce.ReduceTaskResult;
            results = ReduceTaskResult(this.CommittedOutputFiles, this.KeyType, this.ValueType);
        end
    end
end
