feat: tweak intersection code, add tests

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);
 			}
-		return items;
+			items = await objects.find(query, { projection: { _id: 0, value: 1 } }).toArray();
 		}
+	};
 
 	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 })
+
+		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';

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 () {