Skip to content

Commit 8e44776

Browse files
spinscaleAlexander Reelsen
spinscale
and
Alexander Reelsen
authored
feat: Add sparse vector query (#214)
Sparse vector query is supported in Elasticsearch v8.15+. --------- Co-authored-by: Alexander Reelsen <[email protected]>
1 parent 7f3ce49 commit 8e44776

File tree

7 files changed

+371
-1
lines changed

7 files changed

+371
-1
lines changed

‎src/index.d.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3768,6 +3768,96 @@ declare namespace esb {
37683768
spanQry?: SpanQueryBase
37693769
): SpanFieldMaskingQuery;
37703770

3771+
/**
3772+
* The sparse vector query executes a query consisting of sparse vectors, such as built by a learned sparse retrieval model,
3773+
*
3774+
* NOTE: Only available in Elasticsearch v8.15+
3775+
*/
3776+
export class SparseVectorQuery extends Query {
3777+
constructor(field?: string);
3778+
3779+
/**
3780+
* Sets the field to query
3781+
*
3782+
* @param {string} field the field for the query
3783+
* @returns {SparseVectorQuery}
3784+
*/
3785+
field(field : string) : SparseVectorQuery;
3786+
3787+
/**
3788+
* Set model inference id
3789+
*
3790+
* @param {string} inferenceId The model inference ID
3791+
* @returns {SparseVectorQuery}
3792+
*/
3793+
inferenceId(inferenceId : string) : SparseVectorQuery;
3794+
3795+
/**
3796+
* Sets the input query
3797+
*
3798+
* @param {string} query The input query
3799+
* @returns {SparseVectorQuery}
3800+
*/
3801+
query(query : string) : SparseVectorQuery;
3802+
3803+
/**
3804+
* Set a query vector to the query to run. if you don't use inference
3805+
*
3806+
* @param {Object} queryVector
3807+
* @returns {SparseVectorQuery}
3808+
*/
3809+
queryVector(queryVector : object) : SparseVectorQuery;
3810+
3811+
/**
3812+
* Enable pruning
3813+
*
3814+
* NOTE: Only available in Elasticsearch v9.0+
3815+
*
3816+
* @param {boolean} prune
3817+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
3818+
*/
3819+
prune(prune: boolean): SparseVectorQuery;
3820+
3821+
/**
3822+
* Set pruning config tokens_freq_ratio_threshold
3823+
*
3824+
* NOTE: Only available in Elasticsearch v9.0+
3825+
*
3826+
* @param {number} tokensFreqRatioThreshold
3827+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
3828+
*/
3829+
tokensFreqRatioThreshold(tokensFreqRatioThreshold : number) : SparseVectorQuery;
3830+
3831+
/**
3832+
* Set pruning config tokens_weight_threshold
3833+
*
3834+
* NOTE: Only available in Elasticsearch v9.0+
3835+
*
3836+
* @param {number} tokensWeightThreshold
3837+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
3838+
*/
3839+
tokensWeightThreshold(tokensWeightThreshold : number) : SparseVectorQuery;
3840+
3841+
/**
3842+
* Set pruning config only_score_pruned_tokens
3843+
*
3844+
* NOTE: Only available in Elasticsearch v9.0+
3845+
*
3846+
* @param {boolean} onlyScorePrunedTokens
3847+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
3848+
*/
3849+
onlyScorePrunedTokens(onlyScorePrunedTokens : boolean) : SparseVectorQuery;
3850+
}
3851+
3852+
/**
3853+
* Factory function to instantiate a new SparseVectorQuery object.
3854+
*
3855+
* @returns {SparseVectorQuery}
3856+
*/
3857+
export function sparseVectorQuery(
3858+
field? : string
3859+
) : SparseVectorQuery;
3860+
37713861
/**
37723862
* Knn performs k-nearest neighbor (KNN) searches.
37733863
* This class allows configuring the KNN search with various parameters such as field, query vector,

‎src/index.js‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,8 @@ const {
8989
SpanContainingQuery,
9090
SpanWithinQuery,
9191
SpanFieldMaskingQuery
92-
}
92+
},
93+
vectorQueries: { SparseVectorQuery }
9394
} = require('./queries');
9495

9596
const {
@@ -345,6 +346,9 @@ exports.spanWithinQuery = constructorWrapper(SpanWithinQuery);
345346
exports.SpanFieldMaskingQuery = SpanFieldMaskingQuery;
346347
exports.spanFieldMaskingQuery = constructorWrapper(SpanFieldMaskingQuery);
347348

349+
exports.SparseVectorQuery = SparseVectorQuery;
350+
exports.sparseVectorQuery = constructorWrapper(SparseVectorQuery);
351+
348352
/* ============ ============ ============ */
349353
/* ======== KNN ======== */
350354
/* ============ ============ ============ */

‎src/queries/index.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,5 @@ exports.geoQueries = require('./geo-queries');
1616
exports.specializedQueries = require('./specialized-queries');
1717

1818
exports.spanQueries = require('./span-queries');
19+
20+
exports.vectorQueries = require('./vector-queries');

‎src/queries/vector-queries/index.js‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
'use strict';
2+
3+
exports.SparseVectorQuery = require('./sparse-vector-query');

‎src/queries/vector-queries/sparse-vector-query.js‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
'use strict';
2+
3+
const { Query } = require('../../core');
4+
5+
/**
6+
* The sparse vector query executes a query consisting of sparse vectors, such as built by a learned sparse retrieval model,
7+
* NOTE: Only available in Elasticsearch v8.15+
8+
*
9+
* [Elasticsearch reference](https://www.elastic.co/docs/reference/query-languages/query-dsl/query-dsl-sparse-vector-query)
10+
*
11+
* @example
12+
* const qry = esb.sparseVector().field('ml_tokens').inferenceId('model_id').query('my query');
13+
*
14+
* @extends Query
15+
*/
16+
class SparseVectorQuery extends Query {
17+
// eslint-disable-next-line require-jsdoc
18+
constructor() {
19+
super('sparse_vector');
20+
}
21+
22+
/**
23+
* Sets the field to query
24+
*
25+
* @param {string} field the field for the query
26+
* @returns {SparseVectorQuery}
27+
*/
28+
field(field) {
29+
this._queryOpts.field = field;
30+
return this;
31+
}
32+
33+
/**
34+
* Set model inference id
35+
*
36+
* @param {string} inferenceId The model inference ID
37+
* @returns {SparseVectorQuery}
38+
*/
39+
inferenceId(inferenceId) {
40+
this._queryOpts.inference_id = inferenceId;
41+
return this;
42+
}
43+
44+
/**
45+
* Sets the input query.
46+
* You should set either query or query vector, but not both
47+
*
48+
* @param {string} query The input query
49+
* @returns {SparseVectorQuery}
50+
*/
51+
query(query) {
52+
this._queryOpts.query = query;
53+
return this;
54+
}
55+
56+
/**
57+
* Set a query vector to the query to run. if you don't use inference
58+
* You should set either query or query vector, but not both
59+
*
60+
* @param {Object} queryVector
61+
* @returns {SparseVectorQuery}
62+
*/
63+
queryVector(queryVector) {
64+
this._queryOpts.query_vector = queryVector;
65+
return this;
66+
}
67+
68+
/**
69+
* Enable pruning
70+
*
71+
* NOTE: Only available in Elasticsearch v9.0+
72+
*
73+
* @param {boolean} prune
74+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
75+
*/
76+
prune(prune) {
77+
this._queryOpts.prune = prune;
78+
return this;
79+
}
80+
81+
/**
82+
* Set pruning config tokens_freq_ratio_threshold
83+
*
84+
* NOTE: Only available in Elasticsearch v9.0+
85+
*
86+
* @param {number} tokensFreqRatioThreshold
87+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
88+
*/
89+
tokensFreqRatioThreshold(tokensFreqRatioThreshold) {
90+
if (!this._queryOpts.pruning_config) {
91+
this._queryOpts.pruning_config = {};
92+
}
93+
this._queryOpts.pruning_config.tokens_freq_ratio_threshold = tokensFreqRatioThreshold;
94+
return this;
95+
}
96+
97+
/**
98+
* Set pruning config tokens_weight_threshold
99+
*
100+
* NOTE: Only available in Elasticsearch v9.0+
101+
*
102+
* @param {number} tokensWeightThreshold
103+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
104+
*/
105+
tokensWeightThreshold(tokensWeightThreshold) {
106+
if (!this._queryOpts.pruning_config) {
107+
this._queryOpts.pruning_config = {};
108+
}
109+
this._queryOpts.pruning_config.tokens_weight_threshold = tokensWeightThreshold;
110+
return this;
111+
}
112+
113+
/**
114+
* Set pruning config only_score_pruned_tokens
115+
*
116+
* NOTE: Only available in Elasticsearch v9.0+
117+
*
118+
* @param {boolean} onlyScorePrunedTokens
119+
* @returns {SparseVectorQuery} returns `this` so that calls can be chained.
120+
*/
121+
onlyScorePrunedTokens(onlyScorePrunedTokens) {
122+
if (!this._queryOpts.pruning_config) {
123+
this._queryOpts.pruning_config = {};
124+
}
125+
this._queryOpts.pruning_config.only_score_pruned_tokens = onlyScorePrunedTokens;
126+
return this;
127+
}
128+
}
129+
130+
module.exports = SparseVectorQuery;

‎test/index.test.js‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,9 @@ test('queries are exported', t => {
172172

173173
t.truthy(esb.SpanFieldMaskingQuery);
174174
t.truthy(esb.spanFieldMaskingQuery);
175+
176+
t.truthy(esb.sparseVectorQuery());
177+
t.truthy(esb.SparseVectorQuery);
175178
});
176179

177180
test('aggregations are exported', t => {

0 commit comments

Comments
 (0)