diff --git a/src/database/mongo/sorted.js b/src/database/mongo/sorted.js index 33e6d4e79b..4969fd8030 100644 --- a/src/database/mongo/sorted.js +++ b/src/database/mongo/sorted.js @@ -474,48 +474,64 @@ module.exports = function(db, module) { }); }; - module.getSortedSetUnion = function(sets, start, stop, callback) { - getSortedSetUnion(sets, 1, start, stop, callback); + module.getSortedSetUnion = function(params, callback) { + params.sort = 1; + getSortedSetUnion(params, callback); }; - module.getSortedSetRevUnion = function(sets, start, stop, callback) { - getSortedSetUnion(sets, -1, start, stop, callback); + module.getSortedSetRevUnion = function(params, callback) { + params.sort = -1; + getSortedSetUnion(params, callback); }; - function getSortedSetUnion(sets, sort, start, stop, callback) { - if (!Array.isArray(sets) || !sets.length) { + function getSortedSetUnion(params, callback) { + if (!Array.isArray(params.sets) || !params.sets.length) { return callback(); } - var limit = stop - start + 1; + var limit = params.stop - params.start + 1; if (limit <= 0) { limit = 0; } + var aggregate = {}; + if (params.aggregate) { + aggregate['$' + params.aggregate.toLowerCase()] = '$score'; + } else { + aggregate.$sum = '$score'; + } + var pipeline = [ - { $match: { _key: {$in: sets}} }, - { $group: { _id: {value: '$value'}, totalScore: {$sum : "$score"}} }, - { $sort: { totalScore: sort} } + { $match: { _key: {$in: params.sets}} }, + { $group: { _id: {value: '$value'}, totalScore: aggregate} }, + { $sort: { totalScore: params.sort} } ]; - if (start) { - pipeline.push({ $skip: start }); + if (params.start) { + pipeline.push({ $skip: params.start }); } if (limit > 0) { pipeline.push({ $limit: limit }); } - pipeline.push({ $project: { _id: 0, value: '$_id.value' }}); + var project = { _id: 0, value: '$_id.value' }; + if (params.withScores) { + project.score = '$totalScore'; + } + pipeline.push({ $project: project }); db.collection('objects').aggregate(pipeline, function(err, data) { if (err || !data) { return callback(err); } - data = data.map(function(item) { - return item.value; - }); + if (!params.withScores) { + data = data.map(function(item) { + return item.value; + }); + } + callback(null, data); }); } diff --git a/src/database/redis/sorted.js b/src/database/redis/sorted.js index eaa3b2430f..20a3d7b0ff 100644 --- a/src/database/redis/sorted.js +++ b/src/database/redis/sorted.js @@ -232,32 +232,34 @@ module.exports = function(redisClient, module) { multi.exec(callback); }; - module.getSortedSetUnion = function(sets, start, stop, callback) { - sortedSetUnion('zrange', sets, start, stop, false, callback); + module.getSortedSetUnion = function(params, callback) { + params.method = 'zrange'; + sortedSetUnion(params, callback); }; - module.getSortedSetRevUnion = function(sets, start, stop, callback) { - sortedSetUnion('zrevrange', sets, start, stop, false, callback); + module.getSortedSetRevUnion = function(params, callback) { + params.method = 'zrevrange'; + sortedSetUnion(params, callback); }; - function sortedSetUnion(method, sets, start, stop, withScores, callback) { + function sortedSetUnion(params, callback) { var tempSetName = 'temp_' + Date.now(); - var params = [tempSetName, start, stop]; - if (withScores) { + var rangeParams = [tempSetName, params.start, params.stop]; + if (params.withScores) { params.push('WITHSCORES'); } var multi = redisClient.multi(); - multi.zunionstore([tempSetName, sets.length].concat(sets)); - multi[method](params); + multi.zunionstore([tempSetName, params.sets.length].concat(params.sets)); + multi[params.method](rangeParams); multi.del(tempSetName); multi.exec(function(err, results) { if (err) { return callback(err); } - if (!withScores) { + if (!params.withScores) { return callback(null, results ? results[1] : null); } results = results[1] || []; diff --git a/tests/database/sorted.js b/tests/database/sorted.js index c1066f5ff4..9f9d72f580 100644 --- a/tests/database/sorted.js +++ b/tests/database/sorted.js @@ -438,7 +438,7 @@ describe('Sorted Set methods', function() { describe('getSortedSetUnion()', function() { it('should return an array of values from both sorted sets sorted by scores lowest to highest', function(done) { - db.getSortedSetUnion(['sortedSetTest2', 'sortedSetTest3'], 0, -1, function(err, values) { + db.getSortedSetUnion({sets: ['sortedSetTest2', 'sortedSetTest3'], start: 0, stop: -1}, function(err, values) { assert.equal(err, null); assert.equal(arguments.length, 2); assert.deepEqual(values, ['value1', 'value2', 'value4']); @@ -449,7 +449,7 @@ describe('Sorted Set methods', function() { describe('getSortedSetRevUnion()', function() { it('should return an array of values from both sorted sets sorted by scores highest to lowest', function(done) { - db.getSortedSetRevUnion(['sortedSetTest2', 'sortedSetTest3'], 0, -1, function(err, values) { + db.getSortedSetRevUnion({sets: ['sortedSetTest2', 'sortedSetTest3'], start: 0, stop: -1}, function(err, values) { assert.equal(err, null); assert.equal(arguments.length, 2); assert.deepEqual(values, ['value4', 'value2', 'value1']);