From 4ee3543ea43db0fce63bbaa8baca8953362dcb5f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bar=C4=B1=C5=9F=20Soner=20U=C5=9Fakl=C4=B1?= Date: Thu, 21 May 2020 22:12:58 -0400 Subject: [PATCH] feat: tweak intersection code, add tests --- src/database/mongo/sorted/intersect.js | 82 ++++++++++++++------------ test/database/sorted.js | 52 ++++++++++++++++ 2 files changed, 96 insertions(+), 38 deletions(-) diff --git a/src/database/mongo/sorted/intersect.js b/src/database/mongo/sorted/intersect.js index 242dfd5c70..32be949ae7 100644 --- a/src/database/mongo/sorted/intersect.js +++ b/src/database/mongo/sorted/intersect.js @@ -14,29 +14,16 @@ module.exports = function (module) { projection: { _id: 0, value: 1 }, }).toArray(); - items = items.map(i => i.value); const otherSets = keys.filter(s => s !== counts.smallestSet); - if (otherSets.length === 1) { - return await objects.countDocuments({ - _key: otherSets[0], value: { $in: items }, - }); - } - items = await intersectValuesWithSets(items, otherSets); - return items.length; - }; - - async function intersectValuesWithSets(items, sets) { - for (let i = 0; i < sets.length; i++) { + for (let i = 0; i < otherSets.length; i++) { /* eslint-disable no-await-in-loop */ - items = await module.client.collection('objects').find({ - _key: sets[i], value: { $in: items }, - }, { - projection: { _id: 0, value: 1 }, - }).toArray(); - items = items.map(i => i.value); + const query = { _key: otherSets[i], value: { $in: items.map(i => i.value) } }; + if (i === otherSets.length - 1) { + return await objects.countDocuments(query); + } + items = await objects.find(query, { projection: { _id: 0, value: 1 } }).toArray(); } - return items; - } + }; async function countSets(sets, limit) { const objects = module.client.collection('objects'); @@ -92,30 +79,49 @@ module.exports = function (module) { let items = await objects.find({ _key: params.counts.smallestSet }, { projection: { _id: 0, value: 1 }, }).toArray(); - if (!items.length) { - return []; - } - + const sortSet = params.sets[params.weights.indexOf(1)]; const otherSets = params.sets.filter(s => s !== params.counts.smallestSet); - items = await intersectValuesWithSets(items.map(i => i.value), otherSets); - if (!items.length) { - return []; - } const project = { _id: 0, value: 1 }; if (params.withScores) { project.score = 1; } - const sortSet = params.sets[params.weights.indexOf(1)]; - let res = await objects - .find({ _key: sortSet, value: { $in: items } }, { projection: project }) - .sort({ score: params.sort }) - .skip(params.start) - .limit(params.limit) - .toArray(); + + if (sortSet !== params.counts.smallestSet) { + // move sortSet to the end of array + otherSets.push(otherSets.splice(otherSets.indexOf(sortSet), 1)[0]); + for (let i = 0; i < otherSets.length; i++) { + /* eslint-disable no-await-in-loop */ + const cursor = objects.find({ _key: otherSets[i], value: { $in: items.map(i => i.value) } }); + // at the last step sort by sortSet + if (i === otherSets.length - 1) { + cursor.project(project).sort({ score: params.sort }).skip(params.start).limit(params.limit); + } else { + cursor.project({ _id: 0, value: 1 }); + } + items = await cursor.toArray(); + } + } else { + for (let i = 0; i < otherSets.length; i++) { + /* eslint-disable no-await-in-loop */ + items = await module.client.collection('objects').find({ + _key: otherSets[i], value: { $in: items.map(i => i.value) }, + }, { projection: { _id: 0, value: 1 } }).toArray(); + } + if (!items.length) { + return []; + } + items = await objects.find({ _key: sortSet, value: { $in: items.map(i => i.value) } }) + .project(project) + .sort({ score: params.sort }) + .skip(params.start) + .limit(params.limit) + .toArray(); + } + if (!params.withScores) { - res = res.map(i => i.value); + items = items.map(i => i.value); } - return res; + return items; } async function intersectBatch(params) { @@ -159,7 +165,7 @@ module.exports = function (module) { } async function intersectAggregate(params) { - var aggregate = {}; + const aggregate = {}; if (params.aggregate) { aggregate['$' + params.aggregate.toLowerCase()] = '$score'; diff --git a/test/database/sorted.js b/test/database/sorted.js index c8d3d4c5cf..c8bc399f2e 100644 --- a/test/database/sorted.js +++ b/test/database/sorted.js @@ -1217,6 +1217,58 @@ describe('Sorted Set methods', function () { done(); }); }); + + it('should return correct results if sorting by different zset', async function () { + await db.sortedSetAdd('bigzset', [1, 2, 3, 4, 5, 6], ['a', 'b', 'c', 'd', 'e', 'f']); + await db.sortedSetAdd('smallzset', [3, 2, 1], ['b', 'e', 'g']); + const data = await db.getSortedSetRevIntersect({ + sets: ['bigzset', 'smallzset'], + start: 0, + stop: 19, + weights: [1, 0], + withScores: true, + }); + assert.deepStrictEqual(data, [{ value: 'e', score: 5 }, { value: 'b', score: 2 }]); + const data2 = await db.getSortedSetRevIntersect({ + sets: ['bigzset', 'smallzset'], + start: 0, + stop: 19, + weights: [0, 1], + withScores: true, + }); + assert.deepStrictEqual(data2, [{ value: 'b', score: 3 }, { value: 'e', score: 2 }]); + }); + + it('should return correct results when intersecting big zsets', async function () { + const scores = []; + const values = []; + for (let i = 0; i < 30000; i++) { + scores.push((i + 1) * 1000); + values.push(String(i + 1)); + } + await db.sortedSetAdd('verybigzset', scores, values); + + scores.length = 0; + values.length = 0; + for (let i = 15000; i < 45000; i++) { + scores.push((i + 1) * 1000); + values.push(String(i + 1)); + } + await db.sortedSetAdd('anotherbigzset', scores, values); + const data = await db.getSortedSetRevIntersect({ + sets: ['verybigzset', 'anotherbigzset'], + start: 0, + stop: 3, + weights: [1, 0], + withScores: true, + }); + assert.deepStrictEqual(data, [ + { value: '30000', score: 30000000 }, + { value: '29999', score: 29999000 }, + { value: '29998', score: 29998000 }, + { value: '29997', score: 29997000 }, + ]); + }); }); describe('sortedSetIntersectCard', function () {