%SqliteKeyValueSink
% An implementation of AbstractKeyValueSink, which writes to one of a selection
% of tables in a single database file based upon the Partitioner.

%   Copyright 2014-2020 The MathWorks, Inc.

classdef (Sealed, Hidden) SqliteKeyValueSink < parallel.internal.mapreduce.AbstractKeyValueSink
    properties (GetAccess = private, SetAccess = immutable)
        % General properties for this sink
        NumReduceTasks; % The number of expected reduce tasks.
        Partitioner; % The Partitioner associated with this sink.
        SqliteStore; % An SQLite database instance.
    end
    
    properties (Access = private)
        % Properties for task metadata output
        ReduceTaskSizes; % A vector of sizes for how many values have been written against each reduce task.
        KeyType; % The key type that has been written to this sink.
        ValueType; % The value type that has been written to this sink.
        
        % Properties regarding the commit process.
        TransactionInProgress = false;
        TablesExist = false;
    end
    
    methods
        % Construct a sink that writes to multiple tables in a target
        % SQLite database file.
        function this = SqliteKeyValueSink(targetFile, numReduceTasks, partitioner)
            import parallel.internal.mapreduce.SqliteKeyValueSink;
            
            this.SqliteStore = matlab.mapreduce.internal.KeyValueStore(targetFile, 'createIfNecessary', 'doNotDelete');
            
            this.NumReduceTasks = numReduceTasks;
            this.Partitioner = partitioner;
            
            this.ReduceTaskSizes = zeros(numReduceTasks, 1);
            
            this.TablesExist = this.SqliteStore.storeExists('KeyValueStore1');
        end
        
        % Add the given <key, value> to the underlying data storage.
        % The keys and values input must be column vectors of the same length.
        % Both inputs must be vertcat-able with the input of other calls to
        % add.
        function add(this, keys, values)
            if isempty(keys)
                return;
            end
            
            reduceIndices = this.Partitioner.partition(keys);
            
            % Transaction initialization
            if ~this.TransactionInProgress
                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);
                
                if ~this.TablesExist
                    this.pCreateAllStores(exampleKey);
                end
                
                this.pBeginTransaction('Map Task Intermediate Output');
            end
            
            % Flush of buffers
            for reduceTaskID = 1:this.NumReduceTasks
                indices = (reduceIndices == reduceTaskID);
                if ~any(indices)
                    continue;
                end
                
                selectedKeys = keys(indices);
                selectedValues = values(indices);
                this.ReduceTaskSizes(reduceTaskID) = this.ReduceTaskSizes(reduceTaskID) + numel(selectedValues);
                
                this.SqliteStore.selectStoreToWrite(sprintf('KeyValueStore%i', reduceTaskID));
                this.SqliteStore.addKeyValue(selectedKeys, selectedValues);
            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)
            if this.TransactionInProgress
                this.pCommitTransaction('Map Task Intermediate Output');
            end
        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.MapTaskResult;
            results = MapTaskResult(this.ReduceTaskSizes, this.KeyType, this.ValueType);
        end
        
        % Throw away all side-effects of this class.
        function delete(this)
            if this.TransactionInProgress
                try
                    this.pRollbackTransaction('Map Task Intermediate Output');
                catch err
                    dctSchedulerMessage( 1, ...
                        'Failed to rollback transaction, original error', err);
                end
            end
        end
    end
    
    methods (Static)
        % Helper function that adds a key index to the Key-Value table for
        % each store corresponding reduce task IDs in the given map output database.
        function addKeyIndex(targetFile, reduceTaskIds)
            sqliteStore = matlab.mapreduce.internal.KeyValueStore(targetFile, 'createIfNecessary', 'doNotDelete');
            if ~sqliteStore.storeExists(sprintf('KeyValueStore%i', reduceTaskIds(1)))
                return;
            end
            
            detail = 'Adding Key Index';
            sqliteStore.beginTransaction(detail);
            try
                for id = reduceTaskIds
                    name = sprintf('KeyValueStore%i', id);
                    sqliteStore.createStoreKeyIndex(name);
                end
                sqliteStore.commitTransaction(detail);
            catch err
                sqliteStore.rollbackTransaction(detail);
                err.rethrow();
            end
        end
    end
    
    methods (Access = private)
        % Create all output stores in the target database file.
        function pCreateAllStores(this, exampleKey)
            detail = 'Intermediate storage creation';
            sqliteStore = this.SqliteStore;
            sqliteStore.beginTransaction(detail);
            try
                for ii = 1:this.NumReduceTasks
                    name = sprintf('KeyValueStore%i', ii);
                    sqliteStore.createStore(name);
                    sqliteStore.setStoreKeyType(exampleKey);
                end
                sqliteStore.commitTransaction(detail);
            catch err
                sqliteStore.rollbackTransaction(detail);
                err.rethrow();
            end
            this.TablesExist = true;
        end
        
        % Helper function that begins a transaction
        function pBeginTransaction(this, detail)
            this.SqliteStore.beginTransaction(detail);
            this.TransactionInProgress = true;
        end
        
        % Helper function that commits a transaction
        function pCommitTransaction(this, detail)
            this.SqliteStore.commitTransaction(detail);
            this.TransactionInProgress = false;
        end
        
        % Helper function that rolls back a transaction
        function pRollbackTransaction(this, detail)
            this.SqliteStore.rollbackTransaction(detail);
            this.TransactionInProgress = false;
        end
    end
end
