mirror of
https://github.com/NodeBB/NodeBB.git
synced 2025-10-26 08:36:12 +01:00
feat: tweak intersection code, add tests
This commit is contained in:
@@ -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);
|
||||
}
|
||||
return res;
|
||||
|
||||
if (!params.withScores) {
|
||||
items = items.map(i => i.value);
|
||||
}
|
||||
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';
|
||||
|
||||
@@ -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 () {
|
||||
|
||||
Reference in New Issue
Block a user