From 788b6c8c9f7d4cf470e1bf6e48970f21467d73bb Mon Sep 17 00:00:00 2001 From: Andrew Carbonetto Date: Fri, 1 Nov 2024 16:05:48 -0700 Subject: [PATCH] Node: Add `FT.SEARCH` command (#2551) * Node: add FT.SEARCH Signed-off-by: Andrew Carbonetto --------- Signed-off-by: Andrew Carbonetto --- CHANGELOG.md | 1 + node/npm/glide/index.ts | 4 + node/src/server-modules/GlideFt.ts | 144 +++++++++++++-- node/src/server-modules/GlideFtOptions.ts | 45 ++++- node/tests/ServerModules.test.ts | 210 +++++++++++++++++++++- node/tests/TestUtilities.ts | 1 - 6 files changed, 391 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 478aa96339..a1740d9e75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,7 @@ * Python: Add `JSON.ARRAPPEND` command ([#2382](https://github.com/valkey-io/valkey-glide/pull/2382)) * Python: Add `JSON.RESP` command ([#2451](https://github.com/valkey-io/valkey-glide/pull/2451)) * Node: Add `JSON.STRLEN` and `JSON.STRAPPEND` command ([#2537](https://github.com/valkey-io/valkey-glide/pull/2537)) +* Node: Add `FT.SEARCH` ([#2551](https://github.com/valkey-io/valkey-glide/pull/2551)) * Python: Fix example ([#2556](https://github.com/valkey-io/valkey-glide/issues/2556)) #### Breaking Changes diff --git a/node/npm/glide/index.ts b/node/npm/glide/index.ts index 0dc40fd055..781fd26594 100644 --- a/node/npm/glide/index.ts +++ b/node/npm/glide/index.ts @@ -126,7 +126,9 @@ function initialize() { VectorFieldAttributesFlat, VectorFieldAttributesHnsw, FtCreateOptions, + FtSearchOptions, FtInfoReturnType, + FtSearchReturnType, GlideRecord, GlideString, JsonGetOptions, @@ -248,7 +250,9 @@ function initialize() { VectorFieldAttributesFlat, VectorFieldAttributesHnsw, FtCreateOptions, + FtSearchOptions, FtInfoReturnType, + FtSearchReturnType, GlideRecord, GlideJson, GlideString, diff --git a/node/src/server-modules/GlideFt.ts b/node/src/server-modules/GlideFt.ts index 0d58cbfeb0..60c7f72459 100644 --- a/node/src/server-modules/GlideFt.ts +++ b/node/src/server-modules/GlideFt.ts @@ -12,10 +12,10 @@ import { } from "../BaseClient"; import { GlideClient } from "../GlideClient"; import { GlideClusterClient } from "../GlideClusterClient"; -import { Field, FtCreateOptions } from "./GlideFtOptions"; +import { Field, FtCreateOptions, FtSearchOptions } from "./GlideFtOptions"; -/** Data type of {@link GlideFt.info | info} command response. */ -type FtInfoReturnType = Record< +/** Response type of {@link GlideFt.info | ft.info} command. */ +export type FtInfoReturnType = Record< string, | GlideString | number @@ -23,15 +23,23 @@ type FtInfoReturnType = Record< | Record[]> >; +/** + * Response type for the {@link GlideFt.search | ft.search} command. + */ +export type FtSearchReturnType = [ + number, + GlideRecord>, +]; + /** Module for Vector Search commands. */ export class GlideFt { /** * Creates an index and initiates a backfill of that index. * - * @param client The client to execute the command. - * @param indexName The index name for the index to be created. - * @param schema The fields of the index schema, specifying the fields and their types. - * @param options Optional arguments for the `FT.CREATE` command. See {@link FtCreateOptions}. + * @param client - The client to execute the command. + * @param indexName - The index name for the index to be created. + * @param schema - The fields of the index schema, specifying the fields and their types. + * @param options - (Optional) Options for the `FT.CREATE` command. See {@link FtCreateOptions}. * * @returns If the index is successfully created, returns "OK". * @@ -182,8 +190,8 @@ export class GlideFt { /** * Deletes an index and associated content. Indexed document keys are unaffected. * - * @param client The client to execute the command. - * @param indexName The index name. + * @param client - The client to execute the command. + * @param indexName - The index name. * * @returns "OK" * @@ -269,6 +277,122 @@ export class GlideFt { > ).then(convertGlideRecordToRecord); } + + /** + * Uses the provided query expression to locate keys within an index. Once located, the count + * and/or content of indexed fields within those keys can be returned. + * + * @param client - The client to execute the command. + * @param indexName - The index name to search into. + * @param query - The text query to search. + * @param options - (Optional) See {@link FtSearchOptions} and {@link DecoderOption}. + * + * @returns A two-element array, where the first element is the number of documents in the result set, and the + * second element has the format: `GlideRecord>`: + * a mapping between document names and a map of their attributes. + * + * If `count` or `limit` with values `{offset: 0, count: 0}` is + * set, the command returns array with only one element: the number of documents. + * + * @example + * ```typescript + * // + * const vector = Buffer.alloc(24); + * const result = await GlideFt.search(client, "json_idx1", "*=>[KNN 2 @VEC $query_vec]", {params: [{key: "query_vec", value: vector}]}); + * console.log(result); // Output: + * // [ + * // 2, + * // [ + * // { + * // key: "json:2", + * // value: [ + * // { + * // key: "$", + * // value: '{"vec":[1.1,1.2,1.3,1.4,1.5,1.6]}', + * // }, + * // { + * // key: "__VEC_score", + * // value: "11.1100006104", + * // }, + * // ], + * // }, + * // { + * // key: "json:0", + * // value: [ + * // { + * // key: "$", + * // value: '{"vec":[1,2,3,4,5,6]}', + * // }, + * // { + * // key: "__VEC_score", + * // value: "91", + * // }, + * // ], + * // }, + * // ], + * // ] + * ``` + */ + static async search( + client: GlideClient | GlideClusterClient, + indexName: GlideString, + query: GlideString, + options?: FtSearchOptions & DecoderOption, + ): Promise { + const args: GlideString[] = ["FT.SEARCH", indexName, query]; + + if (options) { + // RETURN + if (options.returnFields) { + const returnFields: GlideString[] = []; + options.returnFields.forEach((returnField) => + returnField.alias + ? returnFields.push( + returnField.fieldIdentifier, + "AS", + returnField.alias, + ) + : returnFields.push(returnField.fieldIdentifier), + ); + args.push( + "RETURN", + returnFields.length.toString(), + ...returnFields, + ); + } + + // TIMEOUT + if (options.timeout) { + args.push("TIMEOUT", options.timeout.toString()); + } + + // PARAMS + if (options.params) { + args.push("PARAMS", (options.params.length * 2).toString()); + options.params.forEach((param) => + args.push(param.key, param.value), + ); + } + + // LIMIT + if (options.limit) { + args.push( + "LIMIT", + options.limit.offset.toString(), + options.limit.count.toString(), + ); + } + + // COUNT + if (options.count) { + args.push("COUNT"); + } + } + + return _handleCustomCommand(client, args, options) as Promise< + [number, GlideRecord>] + >; + } } /** @@ -277,7 +401,7 @@ export class GlideFt { async function _handleCustomCommand( client: GlideClient | GlideClusterClient, args: GlideString[], - decoderOption?: DecoderOption, + decoderOption: DecoderOption = {}, ): Promise { return client instanceof GlideClient ? (client as GlideClient).customCommand(args, decoderOption) diff --git a/node/src/server-modules/GlideFtOptions.ts b/node/src/server-modules/GlideFtOptions.ts index 24846da6d2..fffccd11c1 100644 --- a/node/src/server-modules/GlideFtOptions.ts +++ b/node/src/server-modules/GlideFtOptions.ts @@ -2,7 +2,7 @@ * Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0 */ -import { GlideString } from "../BaseClient"; +import { GlideRecord, GlideString } from "../BaseClient"; interface BaseField { /** The name of the field. */ @@ -118,3 +118,46 @@ export interface FtCreateOptions { /** The prefix of the key to be indexed. */ prefixes?: GlideString[]; } + +/** + * Represents the input options to be used in the FT.SEARCH command. + * All fields in this class are optional inputs for FT.SEARCH. + */ +export type FtSearchOptions = { + /** + * Add a field to be returned. + * @param fieldIdentifier field name to return. + * @param alias optional alias for the field name to return. + */ + returnFields?: { fieldIdentifier: GlideString; alias?: GlideString }[]; + + /** Query timeout in milliseconds. */ + timeout?: number; + + /** + * Query parameters, which could be referenced in the query by `$` sign, followed by + * the parameter name. + */ + params?: GlideRecord; +} & ( + | { + /** + * Configure query pagination. By default only first 10 documents are returned. + * + * @param offset Zero-based offset. + * @param count Number of elements to return. + */ + limit?: { offset: number; count: number }; + /** `limit` and `count` are mutually exclusive. */ + count?: never; + } + | { + /** + * Once set, the query will return only the number of documents in the result set without actually + * returning them. + */ + count?: boolean; + /** `limit` and `count` are mutually exclusive. */ + limit?: never; + } +); diff --git a/node/tests/ServerModules.test.ts b/node/tests/ServerModules.test.ts index 314158426f..24ecfd9435 100644 --- a/node/tests/ServerModules.test.ts +++ b/node/tests/ServerModules.test.ts @@ -13,6 +13,7 @@ import { v4 as uuidv4 } from "uuid"; import { ConditionalChange, Decoder, + FtSearchReturnType, GlideClusterClient, GlideFt, GlideJson, @@ -32,6 +33,7 @@ import { } from "./TestUtilities"; const TIMEOUT = 50000; +const DATA_PROCESSING_TIMEOUT = 1000; describe("Server Module Tests", () => { let cluster: ValkeyCluster; @@ -1217,7 +1219,7 @@ describe("Server Module Tests", () => { expect(info).toContain("# search_index_stats"); }); - it("Ft.Create test", async () => { + it("FT.CREATE test", async () => { client = await GlideClusterClient.createClient( getClientConfigurationOption( cluster.getAddresses(), @@ -1371,7 +1373,7 @@ describe("Server Module Tests", () => { } }); - it("Ft.DROPINDEX test", async () => { + it("FT.DROPINDEX test", async () => { client = await GlideClusterClient.createClient( getClientConfigurationOption( cluster.getAddresses(), @@ -1488,5 +1490,209 @@ describe("Server Module Tests", () => { ); }, ); + + it("FT.SEARCH binary test", async () => { + client = await GlideClusterClient.createClient( + getClientConfigurationOption( + cluster.getAddresses(), + ProtocolVersion.RESP3, + ), + ); + + const prefix = "{" + uuidv4() + "}:"; + const index = prefix + "index"; + + // setup a hash index: + expect( + await GlideFt.create( + client, + index, + [ + { + type: "VECTOR", + name: "vec", + alias: "VEC", + attributes: { + algorithm: "HNSW", + distanceMetric: "L2", + dimensions: 2, + }, + }, + ], + { + dataType: "HASH", + prefixes: [prefix], + }, + ), + ).toEqual("OK"); + + const binaryValue1 = Buffer.alloc(8); + expect( + await client.hset(Buffer.from(prefix + "0"), [ + // value of + { field: "vec", value: binaryValue1 }, + ]), + ).toEqual(1); + + const binaryValue2: Buffer = Buffer.alloc(8); + binaryValue2[6] = 0x80; + binaryValue2[7] = 0xbf; + expect( + await client.hset(Buffer.from(prefix + "1"), [ + // value of + { field: "vec", value: binaryValue2 }, + ]), + ).toEqual(1); + + // let server digest the data and update index + const sleep = new Promise((resolve) => + setTimeout(resolve, DATA_PROCESSING_TIMEOUT), + ); + await sleep; + + // With the `COUNT` parameters - returns only the count + const binaryResultCount: FtSearchReturnType = await GlideFt.search( + client, + index, + "*=>[KNN 2 @VEC $query_vec]", + { + params: [{ key: "query_vec", value: binaryValue1 }], + timeout: 10000, + count: true, + decoder: Decoder.Bytes, + }, + ); + expect(binaryResultCount).toEqual([2]); + + const binaryResult: FtSearchReturnType = await GlideFt.search( + client, + index, + "*=>[KNN 2 @VEC $query_vec]", + { + params: [{ key: "query_vec", value: binaryValue1 }], + timeout: 10000, + decoder: Decoder.Bytes, + }, + ); + + const expectedBinaryResult: FtSearchReturnType = [ + 2, + [ + { + key: Buffer.from(prefix + "1"), + value: [ + { + key: Buffer.from("vec"), + value: binaryValue2, + }, + { + key: Buffer.from("__VEC_score"), + value: Buffer.from("1"), + }, + ], + }, + { + key: Buffer.from(prefix + "0"), + value: [ + { + key: Buffer.from("vec"), + value: binaryValue1, + }, + { + key: Buffer.from("__VEC_score"), + value: Buffer.from("0"), + }, + ], + }, + ], + ]; + expect(binaryResult).toEqual(expectedBinaryResult); + }); + + it("FT.SEARCH string test", async () => { + client = await GlideClusterClient.createClient( + getClientConfigurationOption( + cluster.getAddresses(), + ProtocolVersion.RESP3, + ), + ); + + const prefix = "{" + uuidv4() + "}:"; + const index = prefix + "index"; + + // set string values + expect( + await GlideJson.set( + client, + prefix + "1", + "$", + '[{"arr": 42}, {"val": "hello"}, {"val": "world"}]', + ), + ).toEqual("OK"); + + // setup a json index: + expect( + await GlideFt.create( + client, + index, + [ + { + type: "NUMERIC", + name: "$..arr", + alias: "arr", + }, + { + type: "TEXT", + name: "$..val", + alias: "val", + }, + ], + { + dataType: "JSON", + prefixes: [prefix], + }, + ), + ).toEqual("OK"); + + // let server digest the data and update index + const sleep = new Promise((resolve) => + setTimeout(resolve, DATA_PROCESSING_TIMEOUT), + ); + await sleep; + + const stringResult: FtSearchReturnType = await GlideFt.search( + client, + index, + "*", + { + returnFields: [ + { fieldIdentifier: "$..arr", alias: "myarr" }, + { fieldIdentifier: "$..val", alias: "myval" }, + ], + timeout: 10000, + decoder: Decoder.String, + limit: { offset: 0, count: 2 }, + }, + ); + const expectedStringResult: FtSearchReturnType = [ + 1, + [ + { + key: prefix + "1", + value: [ + { + key: "myarr", + value: "42", + }, + { + key: "myval", + value: "hello", + }, + ], + }, + ], + ]; + expect(stringResult).toEqual(expectedStringResult); + }); }); }); diff --git a/node/tests/TestUtilities.ts b/node/tests/TestUtilities.ts index c3fac91e09..0b64b31a04 100644 --- a/node/tests/TestUtilities.ts +++ b/node/tests/TestUtilities.ts @@ -177,7 +177,6 @@ export function flushallOnPort(port: number): Promise { */ export const parseEndpoints = (endpointsStr: string): [string, number][] => { try { - console.log(endpointsStr); const endpoints: string[][] = endpointsStr .split(",") .map((endpoint) => endpoint.split(":"));