%PartitionedSink
%  Class that partitions added key-value pairs and delegates to a set of sinks.

%   Copyright 2014-2016 The MathWorks, Inc.

classdef (Sealed, Hidden) PartitionedSink < parallel.internal.mapreduce.AbstractKeyValueSink
    properties (SetAccess = immutable)
        Sinks       % Cell array of sink object for each reduce task ID
        Partitioner % The mapping between keys and reduce task IDs
    end
    
    methods
        function obj = PartitionedSink(sinks, partitioner)
            obj.Sinks = sinks;
            obj.Partitioner = partitioner;
        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(obj, keys, values)
            reduceIndices = obj.Partitioner.partition(keys);
            for reduceTaskID = 1:numel(obj.Sinks)
                indices = (reduceIndices == reduceTaskID);
                if ~any(indices)
                    continue;
                end
                
                selectedKeys = keys(indices);
                selectedValues = values(indices);
                
                obj.Sinks{reduceTaskID}.add(selectedKeys, selectedValues);
            end
        end
        
        % Commit any output to the final location. obj 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(obj)
            cellfun(@(x) x.commit(), obj.Sinks);       
        end
        
        % Provide any data that is to be returned to the driver code for
        % MapReduce. obj 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(obj)
            innerResults = cellfun(@(x){x.getTaskOutputMetadata()}, obj.Sinks);
            results = vertcat(innerResults{:});
        end
    end
end
