diff --git a/src/core/index.js b/src/core/index.js index dcf0305..669c334 100644 --- a/src/core/index.js +++ b/src/core/index.js @@ -8,6 +8,8 @@ exports.Aggregation = require('./aggregation'); exports.Query = require('./query'); +exports.KNN = require('./knn'); + exports.Suggester = require('./suggester'); exports.Script = require('./script'); diff --git a/src/core/knn.js b/src/core/knn.js new file mode 100644 index 0000000..c5970e4 --- /dev/null +++ b/src/core/knn.js @@ -0,0 +1,133 @@ +'use strict'; + +const { recursiveToJSON, checkType } = require('./util'); +const Query = require('./query'); + +/** + * Class representing a k-Nearest Neighbors (k-NN) query. + * This class extends the Query class to support the specifics of k-NN search, including setting up the field, + * query vector, number of neighbors (k), and number of candidates. + * + * NOTE: kNN search was added to Elasticsearch in v8.0 + * + * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html) + */ +class KNN { + /** + * Creates an instance of KNN. + */ + constructor(field, k, numCandidates) { + if (k > numCandidates) + throw new Error('KNN numCandidates cannot be less than k'); + this._body = {}; + this._body.field = field; + this._body.k = k; + this._body.filter = []; + this._body.num_candidates = numCandidates; + } + + /** + * Sets the query vector for the k-NN search. + * @param {Array} vector - The query vector. + * @returns {KNN} Returns the instance of KNN for method chaining. + */ + queryVector(vector) { + if (this._body.query_vector_builder) + throw new Error( + 'cannot provide both query_vector_builder and query_vector' + ); + this._body.query_vector = vector; + return this; + } + + /** + * Sets the query vector builder for the k-NN search. + * This method configures a query vector builder using a specified model ID and model text. + * It's important to note that either a direct query vector or a query vector builder can be + * provided, but not both. + * + * @param {string} modelId - The ID of the model to be used for generating the query vector. + * @param {string} modelText - The text input based on which the query vector is generated. + * @returns {KNN} Returns the instance of KNN for method chaining. + * @throws {Error} Throws an error if both query_vector_builder and query_vector are provided. + * + * Usage example: + * let knn = new KNN(); + * knn.queryVectorBuilder('model_123', 'Sample model text'); + */ + queryVectorBuilder(modelId, modelText) { + if (this._body.query_vector) + throw new Error( + 'cannot provide both query_vector_builder and query_vector' + ); + this._body.query_vector_builder = { + text_embeddings: { + model_id: modelId, + model_text: modelText + } + }; + return this; + } + + /** + * Adds one or more filter queries to the k-NN search. + * + * This method is designed to apply filters to the k-NN search. It accepts either a single + * query or an array of queries. Each query acts as a filter, refining the search results + * according to the specified conditions. These queries must be instances of the `Query` class. + * If any provided query is not an instance of `Query`, a TypeError is thrown. + * + * @param {Query|Query[]} queries - A single `Query` instance or an array of `Query` instances for filtering. + * @returns {KNN} Returns `this` to allow method chaining. + * @throws {TypeError} If any of the provided queries is not an instance of `Query`. + * + * Usage example: + * let knn = new KNN(); + * knn.filter(new TermQuery('field', 'value')); // Applying a single filter query + * knn.filter([new TermQuery('field1', 'value1'), new TermQuery('field2', 'value2')]); // Applying multiple filter queries + */ + filter(queries) { + const queryArray = Array.isArray(queries) ? queries : [queries]; + queryArray.forEach(query => { + checkType(query, Query); + this._body.filter.push(query); + }); + return this; + } + + /** + * Sets the field to perform the k-NN search on. + * @param {number} boost - The number of the boost + * @returns {KNN} Returns the instance of KNN for method chaining. + */ + boost(boost) { + this._body.boost = boost; + return this; + } + + /** + * Sets the field to perform the k-NN search on. + * @param {number} similarity - The number of the similarity + * @returns {KNN} Returns the instance of KNN for method chaining. + */ + similarity(similarity) { + this._body.similarity = similarity; + return this; + } + + /** + * Override default `toJSON` to return DSL representation for the `query` + * + * @override + * @returns {Object} returns an Object which maps to the elasticsearch query DSL + */ + toJSON() { + if (!this._body.query_vector && !this._body.query_vector_builder) + throw new Error( + 'either query_vector_builder or query_vector must be provided' + ); + return recursiveToJSON(this._body); + } +} + +module.exports = KNN; diff --git a/src/core/request-body-search.js b/src/core/request-body-search.js index c1a90d1..79ecb87 100644 --- a/src/core/request-body-search.js +++ b/src/core/request-body-search.js @@ -10,7 +10,8 @@ const Query = require('./query'), Rescore = require('./rescore'), Sort = require('./sort'), Highlight = require('./highlight'), - InnerHits = require('./inner-hits'); + InnerHits = require('./inner-hits'), + KNN = require('./knn'); const { checkType, setDefault, recursiveToJSON } = require('./util'); @@ -69,6 +70,7 @@ class RequestBodySearch { constructor() { // Maybe accept some optional parameter? this._body = {}; + this._knn = []; this._aggs = []; this._suggests = []; this._suggestText = null; @@ -87,6 +89,21 @@ class RequestBodySearch { return this; } + /** + * Sets knn on the search request body. + * + * @param {Knn|Knn[]} knn + * @returns {RequestBodySearch} returns `this` so that calls can be chained. + */ + kNN(knn) { + const knns = Array.isArray(knn) ? knn : [knn]; + knns.forEach(_knn => { + checkType(_knn, KNN); + this._knn.push(_knn); + }); + return this; + } + /** * Sets aggregation on the request body. * Alias for method `aggregation` @@ -785,6 +802,8 @@ class RequestBodySearch { toJSON() { const dsl = recursiveToJSON(this._body); + if (!isEmpty(this._knn)) dsl.knn = this._knn; + if (!isEmpty(this._aggs)) dsl.aggs = recMerge(this._aggs); if (!isEmpty(this._suggests) || !isNil(this._suggestText)) { diff --git a/src/index.d.ts b/src/index.d.ts index 9a01a76..9ccc120 100644 --- a/src/index.d.ts +++ b/src/index.d.ts @@ -18,6 +18,13 @@ declare namespace esb { */ query(query: Query): this; + /** + * Sets knn on the request body. + * + * @param {KNN|KNN[]} knn + */ + kNN(knn: KNN | KNN[]): this; + /** * Sets aggregation on the request body. * Alias for method `aggregation` @@ -3074,7 +3081,7 @@ declare namespace esb { /** * Sets the script used to compute the score of documents returned by the query. - * + * * @param {Script} script A valid `Script` object */ script(script: Script): this; @@ -3268,6 +3275,86 @@ declare namespace esb { */ export function distanceFeatureQuery(field?: string): DistanceFeatureQuery; + /** + * The `rank_feature` query boosts the relevance score on the numeric value of + * document with a rank_feature/rank_features field. + * + * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-rank-feature-query.html) + * + * NOTE: This query was added in elasticsearch v7.0. + * + * @example + * const query = new RankFeatureQuery('rank_feature_field'); + * query + * .linear() + * .toJSON(); + * @param {string} field The field inside the document to be used in the query + * @return {RankFeatureQuery} + */ + export class RankFeatureQuery extends Query { + constructor(field?: string); + + /** + * Sets the field for the `distance_feature` query + * @param {string} fieldName Name of the field inside the document + * @returns {DistanceFeatureQuery} Instance of the query + */ + field(fieldName: string): RankFeatureQuery; + + /** + * Linear function to boost relevance scores based on the value of the rank feature field + * @returns {RankFeatureQuery} + */ + linear(): RankFeatureQuery; + + /** + * Saturation function to boost relevance scores based on the value of the rank feature field. + * Uses a default pivot value computed by Elasticsearch. + * @returns {RankFeatureQuery} + */ + saturation(): RankFeatureQuery; + + /** + * Saturation function to boost relevance scores based on the value of the rank feature field. + * @param {number} pivot + * @returns {RankFeatureQuery} + */ + saturationPivot(pivot: number): RankFeatureQuery; + + /** + * The log function gives a score equal to log(scaling_factor + S), where S + * is the value of the rank feature field and scaling_factor is a configurable + * scaling factor. + * @param {number} scaling_factor + * @returns {RankFeatureQuery} + */ + log(scalingFactor: number): RankFeatureQuery; + + /** + * The sigmoid function extends the saturation function with a configurable exponent. + * @param {number} pivot + * @param {number} exponent + * @returns {RankFeatureQuery} + */ + sigmoid(pivot: number, exponent: number): RankFeatureQuery; + } + + /** + * The `rank_feature` query boosts the relevance score on the numeric value of + * document with a rank_feature/rank_features field. + * + * [Elasticsearch reference](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-rank-feature-query.html) + * + * @example + * const query = new RankFeatureQuery('rank_feature_field'); + * query + * .linear() + * .toJSON(); + * @param {string} field The field inside the document to be used in the query + * @return {RankFeatureQuery} + */ + export function rankFeatureQuery(field?: string): RankFeatureQuery; + /** * Interface-like class used to group and identify various implementations of Span queries. * @@ -3614,6 +3701,84 @@ declare namespace esb { spanQry?: SpanQueryBase ): SpanFieldMaskingQuery; + /** + * Knn performs k-nearest neighbor (KNN) searches. + * This class allows configuring the KNN search with various parameters such as field, query vector, + * number of nearest neighbors (k), number of candidates, boost factor, and similarity metric. + * + * NOTE: Only available in Elasticsearch v8.0+ + */ + export class KNN { + /** + * Creates an instance of Knn, initializing the internal state for the k-NN search. + * + * @param {string} field - (Optional) The field against which to perform the k-NN search. + * @param {number} k - (Optional) The number of nearest neighbors to retrieve. + * @param {number} numCandidates - (Optional) The number of candidate neighbors to consider during the search. + * @throws {Error} If the number of candidates (numCandidates) is less than the number of neighbors (k). + */ + constructor(field: string, k: number, numCandidates: number); + + /** + * Sets the query vector for the KNN search, an array of numbers representing the reference point. + * + * @param {number[]} vector + */ + queryVector(vector: number[]): this; + + /** + * Sets the query vector builder for the k-NN search. + * This method configures a query vector builder using a specified model ID and model text. + * Note that either a direct query vector or a query vector builder can be provided, but not both. + * + * @param {string} modelId - The ID of the model used for generating the query vector. + * @param {string} modelText - The text input based on which the query vector is generated. + * @returns {KNN} Returns the instance of Knn for method chaining. + * @throws {Error} If both query_vector_builder and query_vector are provided. + */ + queryVectorBuilder(modelId: string, modelText: string): this; + + /** + * Adds one or more filter queries to the k-NN search. + * This method is designed to apply filters to the k-NN search. It accepts either a single + * query or an array of queries. Each query acts as a filter, refining the search results + * according to the specified conditions. These queries must be instances of the `Query` class. + * + * @param {Query|Query[]} queries - A single `Query` instance or an array of `Query` instances for filtering. + * @returns {KNN} Returns `this` to allow method chaining. + * @throws {TypeError} If any of the provided queries is not an instance of `Query`. + */ + filter(queries: Query | Query[]): this; + + /** + * Applies a boost factor to the query to influence the relevance score of returned documents. + * + * @param {number} boost + */ + boost(boost: number): this; + + /** + * Sets the similarity metric used in the KNN algorithm to calculate similarity. + * + * @param {number} similarity + */ + similarity(similarity: number): this; + + /** + * Override default `toJSON` to return DSL representation for the `query` + * + * @override + */ + toJSON(): object; + } + + /** + * Factory function to instantiate a new Knn object. + * + * @returns {KNN} + */ + export function kNN(field: string, k: number, numCandidates: number): KNN; + /** * Base class implementation for all aggregation types. * diff --git a/src/index.js b/src/index.js index af7f821..a3bb1c0 100644 --- a/src/index.js +++ b/src/index.js @@ -14,6 +14,7 @@ const { InnerHits, SearchTemplate, Query, + KNN, util: { constructorWrapper } } = require('./core'); @@ -333,6 +334,13 @@ exports.spanWithinQuery = constructorWrapper(SpanWithinQuery); exports.SpanFieldMaskingQuery = SpanFieldMaskingQuery; exports.spanFieldMaskingQuery = constructorWrapper(SpanFieldMaskingQuery); + +/* ============ ============ ============ */ +/* ======== KNN ======== */ +/* ============ ============ ============ */ +exports.KNN = KNN; +exports.kNN = constructorWrapper(KNN); + /* ============ ============ ============ */ /* ======== Metrics Aggregations ======== */ /* ============ ============ ============ */ diff --git a/test/core-test/knn.test.js b/test/core-test/knn.test.js new file mode 100644 index 0000000..6f717a0 --- /dev/null +++ b/test/core-test/knn.test.js @@ -0,0 +1,87 @@ +import test from 'ava'; +import { KNN, TermQuery } from '../../src'; + +test('knn can be instantiated', t => { + const knn = new KNN('my_field', 5, 10).queryVector([1, 2, 3]); + const json = knn.toJSON(); + t.truthy(json); +}); + +test('knn throws error if numCandidates is less than k', t => { + const error = t.throws(() => + new KNN('my_field', 10, 5).queryVector([1, 2, 3]) + ); + t.is(error.message, 'KNN numCandidates cannot be less than k'); +}); + +test('knn queryVector sets correctly', t => { + const vector = [1, 2, 3]; + const knn = new KNN('my_field', 5, 10).queryVector(vector); + const json = knn.toJSON(); + t.deepEqual(json.query_vector, vector); +}); + +test('knn queryVectorBuilder sets correctly', t => { + const modelId = 'model_123'; + const modelText = 'Sample model text'; + const knn = new KNN('my_field', 5, 10).queryVectorBuilder( + modelId, + modelText + ); + const json = knn.toJSON(); + t.deepEqual(json.query_vector_builder.text_embeddings, { + model_id: modelId, + model_text: modelText + }); +}); + +test('knn filter method adds queries correctly', t => { + const knn = new KNN('my_field', 5, 10).queryVector([1, 2, 3]); + const query = new TermQuery('field', 'value'); + knn.filter(query); + const json = knn.toJSON(); + t.deepEqual(json.filter, [query.toJSON()]); +}); + +test('knn boost method sets correctly', t => { + const boostValue = 1.5; + const knn = new KNN('my_field', 5, 10) + .boost(boostValue) + .queryVector([1, 2, 3]); + const json = knn.toJSON(); + t.is(json.boost, boostValue); +}); + +test('knn similarity method sets correctly', t => { + const similarityValue = 0.8; + const knn = new KNN('my_field', 5, 10) + .similarity(similarityValue) + .queryVector([1, 2, 3]); + const json = knn.toJSON(); + t.is(json.similarity, similarityValue); +}); + +test('knn toJSON method returns correct DSL', t => { + const knn = new KNN('my_field', 5, 10) + .queryVector([1, 2, 3]) + .filter(new TermQuery('field', 'value')); + + const expectedDSL = { + field: 'my_field', + k: 5, + num_candidates: 10, + query_vector: [1, 2, 3], + filter: [{ term: { field: 'value' } }] + }; + + t.deepEqual(knn.toJSON(), expectedDSL); +}); + +test('knn toJSON throws error if neither query_vector nor query_vector_builder is provided', t => { + const knn = new KNN('my_field', 5, 10); + const error = t.throws(() => knn.toJSON()); + t.is( + error.message, + 'either query_vector_builder or query_vector must be provided' + ); +});