Skip to content

Commit dbd0c95

Browse files
committed
MLE-19601 : Update annTopK in Node API
1 parent f57dc15 commit dbd0c95

File tree

3 files changed

+85
-12
lines changed

3 files changed

+85
-12
lines changed

lib/plan-builder-base.js

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,18 @@ function castArg(arg, funcName, paramName, argPos, paramTypes) {
401401
}
402402
});
403403
return true;
404+
case 'PlanAnnTopKOptions':
405+
const planAnnTopKOptionsSet = new Set(['maxDistance', 'max-distance', 'searchFactor','search-factor']);
406+
if(Object.getPrototypeOf(arg) === Map.prototype){
407+
arg.forEach((value, key) => {
408+
if(!planAnnTopKOptionsSet.has(key)) {
409+
throw new Error(
410+
`${argLabel(funcName, paramName, argPos)} has invalid key- ${key}`
411+
);
412+
}
413+
});
414+
}
415+
return true;
404416
default:
405417
return false;
406418
}
@@ -428,7 +440,7 @@ function castArg(arg, funcName, paramName, argPos, paramTypes) {
428440
break;
429441
case 'string':
430442
if (isProtoChained(paramTypes, [types.XsAnyAtomicType, types.TextNode, types.JsonContentNode, types.XmlContentNode,
431-
types.XsString])) {
443+
types.XsString, types.ServerType])) {
432444
return arg;
433445
}
434446
break;

lib/plan-builder-generated.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8089,7 +8089,7 @@ shortestPath(...args) {
80898089
*/
80908090
annTopK(...args) {
80918091
const namer = bldrbase.getNamer(args, 'inputK');
8092-
const paramdefs = [['inputK', [types.XsInteger, PlanParam], true, false], ['vectorColumn', [PlanExprCol, PlanColumn, types.XsString], true, false], ['queryVector', [PlanParam, types.VecVector], true, false], ['distanceColumn', [PlanExprCol, PlanColumn, types.XsString], false, false], ['options', [PlanAnnTopKOptions], false, false]];
8092+
const paramdefs = [['inputK', [types.XsInteger, PlanParam], true, false], ['vectorColumn', [PlanExprCol, PlanColumn, types.XsString], true, false], ['queryVector', [PlanParam, types.VecVector], true, false], ['distanceColumn', [PlanExprCol, PlanColumn, types.XsString], false, false], ['options', [PlanAnnTopKOptions], false, true]];
80938093
const checkedArgs = (namer !== null) ?
80948094
bldrbase.makeNamedArgs(namer, 'PlanModifyPlan.annTopK', 3, new Set(['inputK', 'vectorColumn', 'queryVector', 'distanceColumn', 'options']), paramdefs, args) :
80958095
bldrbase.makePositionalArgs('PlanModifyPlan.annTopK', 3, false, paramdefs, args);

test-basic/annTopK.js

Lines changed: 71 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,10 @@ const assert = require('assert');
1111
const testlib = require('../etc/test-lib');
1212
let serverConfiguration = {};
1313
const execPlan = pbb.execPlan;
14+
const planAnnTopKOptionsMap = new Map();
1415

1516
describe('tests for annTopK', function () {
17+
this.timeout(5000)
1618
before(function (done) {
1719
try {
1820
testlib.findServerConfiguration(serverConfiguration);
@@ -27,22 +29,81 @@ describe('tests for annTopK', function () {
2729
}
2830
});
2931

30-
it('happy path', function (done) {
32+
it('annTopK without PlanAnnTopKOptions', function (done) {
3133
execPlan(p
3234
.fromView('vectors', 'persons', '')
33-
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'), 0.5)
35+
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'))
3436
.orderBy(p.col('name'))
3537
)
3638
.then(function (response) {
37-
const rows = response.rows;
38-
assert(rows.length === 2, 'Expecting both rows in the view to be returned.');
39-
assert(rows[0].name.value === 'Alice');
40-
assert(rows[0].distance.type === 'xs:float', 'Verifying that the distance column was populated.');
41-
assert(rows[1].name.value === 'Bob');
42-
assert(rows[1].distance.type === 'xs:float', 'Verifying that the distance column was populated.');
43-
done();
39+
verifyResults(response.rows, done);
40+
})
41+
.catch(error => done(error));
42+
});
43+
44+
it('annTopK with PlanAnnTopKOptions as a single string', function (done) {
45+
execPlan(p
46+
.fromView('vectors', 'persons', '')
47+
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'), 'onlyIndex')
48+
.orderBy(p.col('name'))
49+
)
50+
.then(function (response) {
51+
verifyResults(response.rows, done);
4452
})
45-
.catch(done);
53+
.catch(error => done(error));
54+
});
55+
56+
it('annTopK with PlanAnnTopKOptions as an array of string', function (done) {
57+
execPlan(p
58+
.fromView('vectors', 'persons', '')
59+
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'),
60+
['onlyIndex', "maxDistance=0.15", "searchFactor=1.0"])
61+
.orderBy(p.col('name'))
62+
).then(function (response) {
63+
verifyResults(response.rows, done);
64+
}).catch(error => done(error));
4665
});
4766

67+
it('annTopK with PlanAnnTopKOptions as a map', function (done) {
68+
planAnnTopKOptionsMap.set("maxDistance", 0.158454656600952);
69+
planAnnTopKOptionsMap.set("searchFactor", 10.0);
70+
execPlan(p
71+
.fromView('vectors', 'persons', '')
72+
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'),
73+
planAnnTopKOptionsMap)
74+
.orderBy(p.col('name'))
75+
)
76+
.then(function (response) {
77+
verifyResults(response.rows, done);
78+
})
79+
.catch(error => done(error));
80+
});
81+
82+
it('annTopK with invalid PlanAnnTopKOptions', function (done) {
83+
planAnnTopKOptionsMap.set('invalid', 10.0);
84+
try{
85+
execPlan(p
86+
.fromView('vectors', 'persons', '')
87+
.annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'),
88+
planAnnTopKOptionsMap)
89+
.orderBy(p.col('name'))
90+
);
91+
} catch(error){
92+
assert(error.message.toString().includes('options argument at 4 of PlanModifyPlan.annTopK() has invalid key- invalid'))
93+
done();
94+
}
95+
});
96+
97+
function verifyResults(rows, done){
98+
try {
99+
assert(rows.length === 2, 'Expecting both rows in the view to be returned.');
100+
assert(rows[0].name.value === 'Alice');
101+
assert(rows[0].distance.type === 'xs:float', 'Verifying that the distance column was populated.');
102+
assert(rows[1].name.value === 'Bob');
103+
assert(rows[1].distance.type === 'xs:float', 'Verifying that the distance column was populated.');
104+
done();
105+
} catch (error){
106+
throw error;
107+
}
108+
}
48109
});

0 commit comments

Comments
 (0)