From b5ca84f076ca0668e90ca3fc7dd878f732ccd956 Mon Sep 17 00:00:00 2001 From: Mark Duckworth <1124037+MarkDuckworth@users.noreply.github.com> Date: Thu, 5 Sep 2024 12:58:24 -0600 Subject: [PATCH 1/2] feat: Return computed distance and set distance thresholds on VectorQueries (#2090) * Return computed distance and set distance thresholds on VectorQueries --- api-report/firestore.api.md | 170 ++-- dev/src/index.ts | 1 + dev/src/reference/query.ts | 80 +- dev/src/reference/vector-query-options.ts | 59 +- dev/src/reference/vector-query.ts | 55 +- dev/system-test/firestore.ts | 973 ++++++++++++++++------ dev/test/vector-query.ts | 347 +++++++- package.json | 4 +- types/firestore.d.ts | 94 ++- 9 files changed, 1365 insertions(+), 418 deletions(-) diff --git a/api-report/firestore.api.md b/api-report/firestore.api.md index 5f9978015..3a146aef2 100644 --- a/api-report/firestore.api.md +++ b/api-report/firestore.api.md @@ -4,24 +4,23 @@ ```ts -/// - import * as $protobuf from 'protobufjs'; import { DocumentData } from '@google-cloud/firestore'; import { Duplex } from 'stream'; import * as firestore from '@google-cloud/firestore'; import { GoogleError } from 'google-gax'; import { Readable } from 'stream'; +import { Span as Span_2 } from '@opentelemetry/api'; // @public export class Aggregate { - constructor(alias: string, aggregateType: AggregateType, fieldPath?: string | FieldPath | undefined); + constructor(alias: string, aggregateType: AggregateType, fieldPath?: (string | FieldPath) | undefined); // (undocumented) readonly aggregateType: AggregateType; // (undocumented) readonly alias: string; // (undocumented) - readonly fieldPath?: string | FieldPath | undefined; + readonly fieldPath?: (string | FieldPath) | undefined; // Warning: (ae-forgotten-export) The symbol "google" needs to be exported by the entry point index.d.ts // // @internal @@ -920,6 +919,11 @@ class Firestore implements firestore.Firestore { // Warning: (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration toJSON(): object; // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // Warning: (ae-forgotten-export) The symbol "TraceUtil" needs to be exported by the entry point index.d.ts + // + // @internal + _traceUtil: TraceUtil; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // // @internal unregisterListener(): void; @@ -1039,10 +1043,14 @@ export class Query, options: { limit: number; distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; }): VectorQuery; + findNearest(options: VectorQueryOptions): VectorQuery; + // (undocumented) + _findNearest(options: VectorQueryOptions): VectorQuery; // Warning: (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag // Warning: (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" // Warning: (tsdoc-undefined-tag) The TSDoc tag "@type" is not defined in this configuration @@ -1449,15 +1457,18 @@ export class Transaction implements firestore.Transaction { // @public export class VectorQuery implements firestore.VectorQuery { // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration - // Warning: (ae-forgotten-export) The symbol "VectorQueryOptions" needs to be exported by the entry point index.d.ts // // @internal - constructor(_query: Query, vectorField: string | firestore.FieldPath, queryVector: firestore.VectorValue | Array, options: VectorQueryOptions); + constructor(_query: Query, _options: VectorQueryOptions); // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // // @internal _createSnapshot(readTime: Timestamp, size: number, docs: () => Array>, changes: () => Array>): VectorQuerySnapshot; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration + explain(options?: firestore.ExplainOptions): Promise>>; get(): Promise>; + // (undocumented) + _getResponse(explainOptions?: firestore.ExplainOptions): Promise>>; isEqual(other: firestore.VectorQuery): boolean; get query(): Query; // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration @@ -1475,7 +1486,17 @@ export class VectorQuery; + vectorField: string | firestore.FieldPath; } // @public @@ -1559,6 +1580,8 @@ export class WriteBatch implements firestore.WriteBatch { // Warning: (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag // Warning: (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" delete(documentRef: firestore.DocumentReference, precondition?: firestore.Precondition): WriteBatch; + // (undocumented) + protected readonly _firestore: Firestore; // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // // @internal @@ -1609,81 +1632,82 @@ export class WriteResult implements firestore.WriteResult { // Warnings were encountered during analysis: // +// build/src/aggregate.d.ts:48:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen // build/src/aggregate.d.ts:49:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/aggregate.d.ts:50:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/bulk-writer.d.ts:50:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:84:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:147:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:154:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:161:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:169:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:184:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:191:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:200:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:218:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:230:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:237:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:488:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/bulk-writer.d.ts:491:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:497:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:504:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:510:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bulk-writer.d.ts:517:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/bundle.d.ts:20:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/filter.d.ts:121:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/filter.d.ts:156:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:49:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:83:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:146:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:153:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:160:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:168:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:183:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:190:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:199:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:217:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:229:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:236:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:487:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/bulk-writer.d.ts:490:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:496:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:503:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:509:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bulk-writer.d.ts:516:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/bundle.d.ts:19:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/filter.d.ts:120:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/filter.d.ts:155:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/index.d.ts:292:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/index.d.ts:312:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/index.d.ts:319:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:334:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:341:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:350:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:358:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:365:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:374:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:857:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:876:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:878:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:880:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:881:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:891:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/index.d.ts:893:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:894:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:896:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/index.d.ts:897:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/path.d.ts:30:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/path.d.ts:32:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration -// build/src/path.d.ts:120:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/path.d.ts:312:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:340:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:347:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:356:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:364:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:371:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:380:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:864:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:883:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:885:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:887:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:888:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:898:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/index.d.ts:900:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:901:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:903:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/index.d.ts:904:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/path.d.ts:29:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/path.d.ts:31:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration +// build/src/path.d.ts:119:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/path.d.ts:311:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/rate-limiter.d.ts:13:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/aggregate-query.d.ts:87:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/aggregate-query.d.ts:85:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/reference/field-filter-internal.d.ts:24:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/reference/field-filter-internal.d.ts:26:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration // build/src/reference/field-order.d.ts:22:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/reference/field-order.d.ts:24:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration -// build/src/reference/query-options.d.ts:28:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/query.d.ts:393:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration -// build/src/reference/query.d.ts:399:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/query.d.ts:401:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:409:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/query.d.ts:411:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:411:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' -// build/src/reference/query.d.ts:413:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:413:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' -// build/src/reference/query.d.ts:415:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:417:24 - (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag -// build/src/reference/query.d.ts:417:17 - (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" -// build/src/reference/query.d.ts:426:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:428:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration -// build/src/reference/query.d.ts:430:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/query.d.ts:616:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference/query.d.ts:617:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/vector-query.d.ts:52:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference/vector-query.d.ts:57:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/serializer.d.ts:26:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/serializer.d.ts:36:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/write-batch.d.ts:86:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/write-batch.d.ts:109:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/query-options.d.ts:27:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/query.d.ts:425:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration +// build/src/reference/query.d.ts:431:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/query.d.ts:433:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:441:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/query.d.ts:443:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:443:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' +// build/src/reference/query.d.ts:445:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:445:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' +// build/src/reference/query.d.ts:447:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:449:24 - (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag +// build/src/reference/query.d.ts:449:17 - (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" +// build/src/reference/query.d.ts:458:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:460:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration +// build/src/reference/query.d.ts:462:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/query.d.ts:648:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference/query.d.ts:649:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/vector-query.d.ts:50:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/vector-query.d.ts:55:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference/vector-query.d.ts:60:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/serializer.d.ts:25:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/serializer.d.ts:35:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/write-batch.d.ts:85:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/write-batch.d.ts:108:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // (No @packageDocumentation comment for this package) diff --git a/dev/src/index.ts b/dev/src/index.ts index a55aa96d5..0b2f195af 100644 --- a/dev/src/index.ts +++ b/dev/src/index.ts @@ -103,6 +103,7 @@ export type {AggregateQuery} from './reference/aggregate-query'; export type {AggregateQuerySnapshot} from './reference/aggregate-query-snapshot'; export type {VectorQuery} from './reference/vector-query'; export type {VectorQuerySnapshot} from './reference/vector-query-snapshot'; +export type {VectorQueryOptions} from './reference/vector-query-options'; export {BulkWriter} from './bulk-writer'; export type {BulkWriterError} from './bulk-writer'; export type {BundleBuilder} from './bundle'; diff --git a/dev/src/reference/query.ts b/dev/src/reference/query.ts index e2811de3b..558acfe86 100644 --- a/dev/src/reference/query.ts +++ b/dev/src/reference/query.ts @@ -629,6 +629,9 @@ export class Query< * @param options - Options control the vector query. `limit` specifies the upper bound of documents to return, must * be a positive integer with a maximum value of 1000. `distanceMeasure` specifies what type of distance is calculated * when performing the query. + * + * @deprecated Use the new {@link findNearest} implementation + * accepting a single `options` param. */ findNearest( vectorField: string | firestore.FieldPath, @@ -637,17 +640,77 @@ export class Query< limit: number; distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; } + ): VectorQuery; + + /** + * Returns a query that can perform vector distance (similarity) search with given parameters. + * + * The returned query, when executed, performs a distance (similarity) search on the specified + * `vectorField` against the given `queryVector` and returns the top documents that are closest + * to the `queryVector`. + * + * Only documents whose `vectorField` field is a {@link VectorValue} of the same dimension as `queryVector` + * participate in the query, all other documents are ignored. + * + * @example + * ``` + * // Returns the closest 10 documents whose Euclidean distance from their 'embedding' fields are closed to [41, 42]. + * const vectorQuery = col.findNearest({ + * vectorField: 'embedding', + * queryVector: [41, 42], + * limit: 10, + * distanceMeasure: 'EUCLIDEAN', + * distanceResultField: 'distance', + * distanceThreshold: 0.125 + * }); + * + * const querySnapshot = await aggregateQuery.get(); + * querySnapshot.forEach(...); + * ``` + * @param options - An argument specifying the behavior of the {@link VectorQuery} returned by this function. + * See {@link VectorQueryOptions}. + */ + findNearest( + options: VectorQueryOptions + ): VectorQuery; + + findNearest( + vectorFieldOrOptions: string | firestore.FieldPath | VectorQueryOptions, + queryVector?: firestore.VectorValue | Array, + options?: { + limit?: number; + distanceMeasure?: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + } ): VectorQuery { - validateFieldPath('vectorField', vectorField); + if ( + typeof vectorFieldOrOptions === 'string' || + vectorFieldOrOptions instanceof FieldPath + ) { + const vqOptions: VectorQueryOptions = { + distanceMeasure: options!.distanceMeasure!, + limit: options!.limit!, + queryVector: queryVector!, + vectorField: vectorFieldOrOptions, + }; + return this._findNearest(vqOptions); + } else { + return this._findNearest(vectorFieldOrOptions as VectorQueryOptions); + } + } + + _findNearest( + options: VectorQueryOptions + ): VectorQuery { + validateFieldPath('vectorField', options.vectorField); if (options.limit <= 0) { - throw invalidArgumentMessage('options.limit', 'positive limit number'); + throw invalidArgumentMessage('limit', 'positive limit number'); } if ( - (Array.isArray(queryVector) - ? queryVector.length - : queryVector.toArray().length) === 0 + (Array.isArray(options.queryVector) + ? options.queryVector.length + : options.queryVector.toArray().length) === 0 ) { throw invalidArgumentMessage( 'queryVector', @@ -655,12 +718,7 @@ export class Query< ); } - return new VectorQuery( - this, - vectorField, - queryVector, - new VectorQueryOptions(options.limit, options.distanceMeasure) - ); + return new VectorQuery(this, options); } /** diff --git a/dev/src/reference/vector-query-options.ts b/dev/src/reference/vector-query-options.ts index cc083aa62..10a851a96 100644 --- a/dev/src/reference/vector-query-options.ts +++ b/dev/src/reference/vector-query-options.ts @@ -14,23 +14,48 @@ * limitations under the License. */ -export class VectorQueryOptions { - constructor( - readonly limit: number, - readonly distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT' - ) {} +import * as firestore from '@google-cloud/firestore'; - isEqual(other: VectorQueryOptions): boolean { - if (this === other) { - return true; - } - if (!(other instanceof VectorQueryOptions)) { - return false; - } +/** + * Specifies the behavior of the {@link VectorQuery} generated by a call to {@link Query.findNearest}. + */ +export interface VectorQueryOptions { + /** + * A string or {@link FieldPath} specifying the vector field to search on. + */ + vectorField: string | firestore.FieldPath; + + /** + * The {@link VectorValue} used to measure the distance from `vectorField` values in the documents. + */ + queryVector: firestore.VectorValue | Array; + + /** + * Specifies the upper bound of documents to return, must be a positive integer with a maximum value of 1000. + */ + limit: number; + + /** + * Specifies what type of distance is calculated when performing the query. + */ + distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + + /** + * Optionally specifies the name of a field that will be set on each returned DocumentSnapshot, + * which will contain the computed distance for the document. + */ + distanceResultField?: string | firestore.FieldPath; - return ( - this.limit === other.limit && - this.distanceMeasure === other.distanceMeasure - ); - } + /** + * Specifies a threshold for which no less similar documents will be returned. The behavior + * of the specified `distanceMeasure` will affect the meaning of the distance threshold. + * + * - For `distanceMeasure: "EUCLIDEAN"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE euclidean_distance <= distanceThreshold + * - For `distanceMeasure: "COSINE"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE cosine_distance <= distanceThreshold + * - For `distanceMeasure: "DOT_PRODUCT"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE dot_product_distance >= distanceThreshold + */ + distanceThreshold?: number; } diff --git a/dev/src/reference/vector-query.ts b/dev/src/reference/vector-query.ts index 4df93e60f..3fe36194f 100644 --- a/dev/src/reference/vector-query.ts +++ b/dev/src/reference/vector-query.ts @@ -56,9 +56,7 @@ export class VectorQuery< */ constructor( private readonly _query: Query, - private readonly vectorField: string | firestore.FieldPath, - private readonly queryVector: firestore.VectorValue | Array, - private readonly options: VectorQueryOptions + private readonly _options: VectorQueryOptions ) { this._queryUtil = new QueryUtil< AppModelType, @@ -79,9 +77,21 @@ export class VectorQuery< * @internal */ private get _rawVectorField(): string { - return typeof this.vectorField === 'string' - ? this.vectorField - : this.vectorField.toString(); + return typeof this._options.vectorField === 'string' + ? this._options.vectorField + : this._options.vectorField.toString(); + } + + /** + * @private + * @internal + */ + private get _rawDistanceResultField(): string | undefined { + if (typeof this._options.distanceResultField === 'undefined') return; + + return typeof this._options.distanceResultField === 'string' + ? this._options.distanceResultField + : this._options.distanceResultField.toString(); } /** @@ -89,9 +99,9 @@ export class VectorQuery< * @internal */ private get _rawQueryVector(): Array { - return Array.isArray(this.queryVector) - ? this.queryVector - : this.queryVector.toArray(); + return Array.isArray(this._options.queryVector) + ? this._options.queryVector + : this._options.queryVector.toArray(); } /** @@ -157,7 +167,7 @@ export class VectorQuery< } /** - * Internal method for serializing a query to its RunAggregationQuery proto + * Internal method for serializing a query to its proto * representation with an optional transaction id. * * @private @@ -170,17 +180,25 @@ export class VectorQuery< ): api.IRunQueryRequest { const queryProto = this._query.toProto(transactionOrReadTime); - const queryVector = Array.isArray(this.queryVector) - ? new VectorValue(this.queryVector) - : (this.queryVector as VectorValue); + const queryVector = Array.isArray(this._options.queryVector) + ? new VectorValue(this._options.queryVector) + : (this._options.queryVector as VectorValue); queryProto.structuredQuery!.findNearest = { - limit: {value: this.options.limit}, - distanceMeasure: this.options.distanceMeasure, + limit: {value: this._options.limit}, + distanceMeasure: this._options.distanceMeasure, vectorField: { - fieldPath: FieldPath.fromArgument(this.vectorField).formattedName, + fieldPath: FieldPath.fromArgument(this._options.vectorField) + .formattedName, }, queryVector: queryVector._toProto(this._query._serializer), + distanceResultField: this._options?.distanceResultField + ? FieldPath.fromArgument(this._options.distanceResultField!) + .formattedName + : undefined, + distanceThreshold: this._options?.distanceThreshold + ? {value: this._options?.distanceThreshold} + : undefined, }; if (explainOptions) { @@ -253,7 +271,10 @@ export class VectorQuery< return ( this._rawVectorField === other._rawVectorField && isPrimitiveArrayEqual(this._rawQueryVector, other._rawQueryVector) && - this.options.isEqual(other.options) + this._options.limit === other._options.limit && + this._options.distanceMeasure === other._options.distanceMeasure && + this._options.distanceThreshold === other._options.distanceThreshold && + this._rawDistanceResultField === other._rawDistanceResultField ); } } diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts index 5f7d683d7..f0299b13b 100644 --- a/dev/system-test/firestore.ts +++ b/dev/system-test/firestore.ts @@ -443,7 +443,9 @@ describe('Firestore class', () => { const explainResults = await indexTestHelper .query(collectionReference) - .findNearest('embedding', FieldValue.vector([1, 3]), { + .findNearest({ + vectorField: 'embedding', + queryVector: FieldValue.vector([1, 3]), limit: 10, distanceMeasure: 'COSINE', }) @@ -473,7 +475,9 @@ describe('Firestore class', () => { const explainResults = await indexTestHelper .query(collectionReference) - .findNearest('embedding', FieldValue.vector([1, 3]), { + .findNearest({ + vectorField: 'embedding', + queryVector: FieldValue.vector([1, 3]), limit: 10, distanceMeasure: 'COSINE', }) @@ -1966,326 +1970,785 @@ describe('Query class', () => { }); }); - it('supports findNearest by EUCLIDEAN distance', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + describe('vector search', () => { + it('supports findNearest by EUCLIDEAN distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - const collectionReference = await indexTestHelper.createTestDocs([ - {foo: 'bar'}, - {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, - {foo: 'bar', embedding: FieldValue.vector([1, 1])}, - {foo: 'bar', embedding: FieldValue.vector([10, 0])}, - {foo: 'bar', embedding: FieldValue.vector([20, 0])}, - {foo: 'bar', embedding: FieldValue.vector([100, 100])}, - ]); + const collectionReference = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, + {foo: 'bar', embedding: FieldValue.vector([1, 1])}, + {foo: 'bar', embedding: FieldValue.vector([10, 0])}, + {foo: 'bar', embedding: FieldValue.vector([20, 0])}, + {foo: 'bar', embedding: FieldValue.vector([100, 100])}, + ]); - const vectorQuery = indexTestHelper - .query(collectionReference) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'EUCLIDEAN', + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([10, 0]))) + .to.be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1]))).to + .be.true; + expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0]))) + .to.be.true; + }); + + it('supports findNearest by COSINE distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 1])}, + '4': {foo: 'bar', embedding: FieldValue.vector([20, 0])}, + '5': {foo: 'bar', embedding: FieldValue.vector([10, 0])}, + '6': {foo: 'bar', embedding: FieldValue.vector([100, 100])}, }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(3); - expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([10, 0]))).to - .be.true; - expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1]))).to - .be.true; - expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0]))).to - .be.true; - }); + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'COSINE', + }); - it('supports findNearest by COSINE distance', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const res = await vectorQuery.get(); + + expect(res.size).to.equal(3); - const collectionReference = await indexTestHelper.setTestDocs({ - '1': {foo: 'bar'}, - '2': {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, - '3': {foo: 'bar', embedding: FieldValue.vector([1, 1])}, - '4': {foo: 'bar', embedding: FieldValue.vector([20, 0])}, - '5': {foo: 'bar', embedding: FieldValue.vector([10, 0])}, - '6': {foo: 'bar', embedding: FieldValue.vector([100, 100])}, + if (res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 1]))) { + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([100, 100])) + ).to.be.true; + } else { + expect( + res.docs[0].get('embedding').isEqual(FieldValue.vector([100, 100])) + ).to.be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1]))) + .to.be.true; + } + + expect( + res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0])) || + res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0])) + ).to.be.true; }); - const vectorQuery = indexTestHelper - .query(collectionReference) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'COSINE', - }); + it('supports findNearest by DOT_PRODUCT distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - const res = await vectorQuery.get(); + const collectionReference = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, + {foo: 'bar', embedding: FieldValue.vector([1, 1])}, + {foo: 'bar', embedding: FieldValue.vector([10, 0])}, + {foo: 'bar', embedding: FieldValue.vector([20, 0])}, + {foo: 'bar', embedding: FieldValue.vector([100, 100])}, + ]); - expect(res.size).to.equal(3); + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'DOT_PRODUCT', + }); - if (res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 1]))) { + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); expect( - res.docs[1].get('embedding').isEqual(FieldValue.vector([100, 100])) + res.docs[0].get('embedding').isEqual(FieldValue.vector([100, 100])) ).to.be.true; - } else { + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([20, 0]))) + .to.be.true; + expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([10, 0]))) + .to.be.true; + }); + + it('findNearest works with converters', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + class FooDistance { + constructor( + readonly foo: string, + readonly embedding: Array + ) {} + } + + const fooConverter = { + toFirestore(d: FooDistance): DocumentData { + return {title: d.foo, embedding: FieldValue.vector(d.embedding)}; + }, + fromFirestore(snapshot: QueryDocumentSnapshot): FooDistance { + const data = snapshot.data(); + return new FooDistance(data.foo, data.embedding.toArray()); + }, + }; + + const collectionRef = await indexTestHelper.createTestDocs([ + {foo: 'bar', embedding: FieldValue.vector([5, 5])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionRef) + .withConverter(fooConverter) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + + expect(res.size).to.equal(1); + expect(res.docs[0].data().foo).to.equal('bar'); + expect(res.docs[0].data().embedding).to.deep.equal([5, 5]); + }); + + it('supports findNearest skipping fields of wrong types', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionRef = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + + // These documents are skipped because it is not really a vector value + {foo: 'bar', embedding: [10, 10]}, + {foo: 'bar', embedding: 'not actually a vector'}, + {foo: 'bar', embedding: null}, + + // Actual vector values + {foo: 'bar', embedding: FieldValue.vector([9, 9])}, + {foo: 'bar', embedding: FieldValue.vector([50, 50])}, + {foo: 'bar', embedding: FieldValue.vector([100, 100])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionRef) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 100, // Intentionally large to get all matches. + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([9, 9]))).to + .be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([50, 50]))) + .to.be.true; expect( - res.docs[0].get('embedding').isEqual(FieldValue.vector([100, 100])) + res.docs[2].get('embedding').isEqual(FieldValue.vector([100, 100])) ).to.be.true; - expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1]))).to + }); + + it('findNearest ignores mismatching dimensions', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionRef = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + + // Vectors with dimension mismatch + {foo: 'bar', embedding: FieldValue.vector([10])}, + + // Vectors with dimension match + {foo: 'bar', embedding: FieldValue.vector([9, 9])}, + {foo: 'bar', embedding: FieldValue.vector([50, 50])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionRef) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(2); + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([9, 9]))).to .be.true; - } + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([50, 50]))) + .to.be.true; + }); - expect( - res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0])) || - res.docs[2].get('embedding').isEqual(FieldValue.vector([20, 0])) - ).to.be.true; - }); + it('supports findNearest on non-existent field', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - it('supports findNearest by DOT_PRODUCT distance', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const collectionRef = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + {foo: 'bar', otherField: [10, 10]}, + {foo: 'bar', otherField: 'not actually a vector'}, + {foo: 'bar', otherField: null}, + ]); - const collectionReference = await indexTestHelper.createTestDocs([ - {foo: 'bar'}, - {foo: 'xxx', embedding: FieldValue.vector([10, 10])}, - {foo: 'bar', embedding: FieldValue.vector([1, 1])}, - {foo: 'bar', embedding: FieldValue.vector([10, 0])}, - {foo: 'bar', embedding: FieldValue.vector([20, 0])}, - {foo: 'bar', embedding: FieldValue.vector([100, 100])}, - ]); + const vectorQuery = indexTestHelper + .query(collectionRef) + .where('foo', '==', 'bar') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); - const vectorQuery = indexTestHelper - .query(collectionReference) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'DOT_PRODUCT', + const res = await vectorQuery.get(); + + expect(res.size).to.equal(0); + }); + + it('supports findNearest on vector nested in a map', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.createTestDocs([ + {nested: {foo: 'bar'}}, + {nested: {foo: 'xxx', embedding: FieldValue.vector([10, 10])}}, + {nested: {foo: 'bar', embedding: FieldValue.vector([1, 1])}}, + {nested: {foo: 'bar', embedding: FieldValue.vector([10, 0])}}, + {nested: {foo: 'bar', embedding: FieldValue.vector([20, 0])}}, + {nested: {foo: 'bar', embedding: FieldValue.vector([100, 100])}}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'nested.embedding', + queryVector: [10, 10], + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); + expect( + res.docs[0].get('nested.embedding').isEqual(FieldValue.vector([10, 10])) + ).to.be.true; + expect( + res.docs[1].get('nested.embedding').isEqual(FieldValue.vector([10, 0])) + ).to.be.true; + expect( + res.docs[2].get('nested.embedding').isEqual(FieldValue.vector([1, 1])) + ).to.be.true; + }); + + it('supports findNearest with select to exclude vector data in response', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.createTestDocs([ + {foo: 1}, + {foo: 2, embedding: FieldValue.vector([10, 10])}, + {foo: 3, embedding: FieldValue.vector([1, 1])}, + {foo: 4, embedding: FieldValue.vector([10, 0])}, + {foo: 5, embedding: FieldValue.vector([20, 0])}, + {foo: 6, embedding: FieldValue.vector([100, 100])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', 'in', [1, 2, 3, 4, 5, 6]) + .select('foo') + .findNearest({ + vectorField: 'embedding', + queryVector: [10, 10], + limit: 10, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(5); + expect(res.docs[0].get('foo')).to.equal(2); + expect(res.docs[1].get('foo')).to.equal(4); + expect(res.docs[2].get('foo')).to.equal(3); + expect(res.docs[3].get('foo')).to.equal(5); + expect(res.docs[4].get('foo')).to.equal(6); + + res.docs.forEach(ds => expect(ds.get('embedding')).to.be.undefined); + }); + + it('supports findNearest limits', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const embeddingVector = []; + const queryVector = []; + for (let i = 0; i < 2048; i++) { + embeddingVector.push(i + 1); + queryVector.push(i - 1); + } + + const collectionReference = await indexTestHelper.createTestDocs([ + {embedding: FieldValue.vector(embeddingVector)}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: queryVector, + limit: 1000, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(1); + expect( + (res.docs[0].get('embedding') as VectorValue).toArray() + ).to.deep.equal(embeddingVector); + }); + + describe('preview API (deprecated)', () => { + it('supports findNearest with EUCLIDEAN', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + {foo: 'bar', embedding: FieldValue.vector([10, 10])}, + {foo: 'bar', embedding: FieldValue.vector([1, 1.1])}, + {foo: 'x', embedding: FieldValue.vector([1, 1])}, + {foo: 'bar', embedding: FieldValue.vector([10, 0])}, + {foo: 'bar', embedding: FieldValue.vector([-100, -100])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest('embedding', [1, 1], { + limit: 3, + distanceMeasure: 'EUCLIDEAN', + }); + + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); + expect( + res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 1.1])) + ).to.be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([10, 0]))) + .to.be.true; + expect( + res.docs[2].get('embedding').isEqual(FieldValue.vector([10, 10])) + ).to.be.true; }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(3); - expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([100, 100]))) - .to.be.true; - expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([20, 0]))).to - .be.true; - expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([10, 0]))).to - .be.true; - }); + it('supports findNearest with COSINE', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.createTestDocs([ + {foo: 'bar'}, + {foo: 'bar', embedding: FieldValue.vector([10, 10])}, + {foo: 'bar', embedding: FieldValue.vector([1, 1.1])}, + {foo: 'x', embedding: FieldValue.vector([1, 1])}, + {foo: 'bar', embedding: FieldValue.vector([10, 0])}, + {foo: 'bar', embedding: FieldValue.vector([-100, -100])}, + ]); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest('embedding', [1, 1], { + limit: 3, + distanceMeasure: 'COSINE', + }); - it('findNearest works with converters', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const res = await vectorQuery.get(); + expect(res.size).to.equal(3); + expect( + res.docs[0].get('embedding').isEqual(FieldValue.vector([10, 10])) + ).to.be.true; + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1.1])) + ).to.be.true; + expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([10, 0]))) + .to.be.true; + }); + }); - class FooDistance { - constructor( - readonly foo: string, - readonly embedding: Array - ) {} - } + describe('requesting computed distance', () => { + it('supports COSINE distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - const fooConverter = { - toFirestore(d: FooDistance): DocumentData { - return {title: d.foo, embedding: FieldValue.vector(d.embedding)}; - }, - fromFirestore(snapshot: QueryDocumentSnapshot): FooDistance { - const data = snapshot.data(); - return new FooDistance(data.foo, data.embedding.toArray()); - }, - }; + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([1, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([0, 1])}, + '4': {foo: 'bar', embedding: FieldValue.vector([0, -0.1])}, + '5': {foo: 'bar', embedding: FieldValue.vector([-1, 0])}, + }); - const collectionRef = await indexTestHelper.createTestDocs([ - {foo: 'bar', embedding: FieldValue.vector([5, 5])}, - ]); + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'COSINE', + distanceResultField: 'distance', + }); + + const res = await vectorQuery.get(); - const vectorQuery = indexTestHelper - .query(collectionRef) - .withConverter(fooConverter) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'EUCLIDEAN', + expect(res.size).to.equal(4); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 0]))) + .to.be.true; + expect(res.docs[0].get('distance')).to.equal(0); + + expect(res.docs[1].get('distance')).to.equal(1); + expect(res.docs[2].get('distance')).to.equal(1); + + expect(res.docs[3].get('embedding').isEqual(FieldValue.vector([-1, 0]))) + .to.be.true; + expect(res.docs[3].get('distance')).to.equal(2); }); - const res = await vectorQuery.get(); + it('supports EUCLIDEAN distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - expect(res.size).to.equal(1); - expect(res.docs[0].data().foo).to.equal('bar'); - expect(res.docs[0].data().embedding).to.deep.equal([5, 5]); - }); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([1, -0.1])}, + '5': {foo: 'bar', embedding: FieldValue.vector([4, 4])}, + }); - it('supports findNearest skipping fields of wrong types', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'EUCLIDEAN', + distanceResultField: 'distance', + }); - const collectionRef = await indexTestHelper.createTestDocs([ - {foo: 'bar'}, + const res = await vectorQuery.get(); - // These documents are skipped because it is not really a vector value - {foo: 'bar', embedding: [10, 10]}, - {foo: 'bar', embedding: 'not actually a vector'}, - {foo: 'bar', embedding: null}, + expect(res.size).to.equal(4); - // Actual vector values - {foo: 'bar', embedding: FieldValue.vector([9, 9])}, - {foo: 'bar', embedding: FieldValue.vector([50, 50])}, - {foo: 'bar', embedding: FieldValue.vector([100, 100])}, - ]); + expect( + res.docs[0].get('embedding').isEqual(FieldValue.vector([1, -0.1])) + ).to.be.true; + expect(res.docs[0].get('distance')).to.equal(0.1); + + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + expect(res.docs[1].get('distance')).to.equal(1); - const vectorQuery = indexTestHelper - .query(collectionRef) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 100, // Intentionally large to get all matches. - distanceMeasure: 'EUCLIDEAN', + expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([4, 4]))) + .to.be.true; + expect(res.docs[2].get('distance')).to.equal(5); + + expect( + res.docs[3].get('embedding').isEqual(FieldValue.vector([1, 100])) + ).to.be.true; + expect(res.docs[3].get('distance')).to.equal(100); }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(3); - expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([9, 9]))).to - .be.true; - expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([50, 50]))).to - .be.true; - expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([100, 100]))) - .to.be.true; - }); + it('supports DOT_PRODUCT distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - it('findNearest ignores mismatching dimensions', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([-20, 0])}, + '5': {foo: 'bar', embedding: FieldValue.vector([0.1, 4])}, + }); - const collectionRef = await indexTestHelper.createTestDocs([ - {foo: 'bar'}, + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'DOT_PRODUCT', + distanceResultField: 'distance', + }); - // Vectors with dimension mismatch - {foo: 'bar', embedding: FieldValue.vector([10])}, + const res = await vectorQuery.get(); - // Vectors with dimension match - {foo: 'bar', embedding: FieldValue.vector([9, 9])}, - {foo: 'bar', embedding: FieldValue.vector([50, 50])}, - ]); + expect(res.size).to.equal(4); + + expect(res.docs[0].get('distance')).to.equal(2); + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + + expect(res.docs[1].get('distance')).to.equal(1); + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 100])) + ).to.be.true; + + expect(res.docs[2].get('distance')).to.equal(0.1); + expect( + res.docs[2].get('embedding').isEqual(FieldValue.vector([0.1, 4])) + ).to.be.true; - const vectorQuery = indexTestHelper - .query(collectionRef) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'EUCLIDEAN', + expect(res.docs[3].get('distance')).to.equal(-20); + expect( + res.docs[3].get('embedding').isEqual(FieldValue.vector([-20, 0])) + ).to.be.true; }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(2); - expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([9, 9]))).to - .be.true; - expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([50, 50]))).to - .be.true; - }); + it('overwrites distance result field on conflict', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - it('supports findNearest on non-existent field', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': { + foo: 'bar', + embedding: FieldValue.vector([0, 1]), + distance: '100 miles', + }, + }); - const collectionRef = await indexTestHelper.createTestDocs([ - {foo: 'bar'}, - {foo: 'bar', otherField: [10, 10]}, - {foo: 'bar', otherField: 'not actually a vector'}, - {foo: 'bar', otherField: null}, - ]); + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'COSINE', + distanceResultField: 'distance', + }); - const vectorQuery = indexTestHelper - .query(collectionRef) - .where('foo', '==', 'bar') - .findNearest('embedding', [10, 10], { - limit: 3, - distanceMeasure: 'EUCLIDEAN', + const res = await vectorQuery.get(); + + expect(res.size).to.equal(1); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([0, 1]))) + .to.be.true; + expect(res.docs[0].get('distance')).to.equal(1); }); - const res = await vectorQuery.get(); + it('supports select queries', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - expect(res.size).to.equal(0); - }); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([1, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([0, 1])}, + '4': {foo: 'bar', embedding: FieldValue.vector([0, -0.1])}, + '5': {foo: 'bar', embedding: FieldValue.vector([-1, 0])}, + }); - it('supports findNearest on vector nested in a map', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const vectorQuery = indexTestHelper + .query(collectionReference) + // value of `distanceResultField` must also be in select statement + .select('embedding', 'distance') + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'COSINE', + distanceResultField: 'distance', + }); - const collectionReference = await indexTestHelper.createTestDocs([ - {nested: {foo: 'bar'}}, - {nested: {foo: 'xxx', embedding: FieldValue.vector([10, 10])}}, - {nested: {foo: 'bar', embedding: FieldValue.vector([1, 1])}}, - {nested: {foo: 'bar', embedding: FieldValue.vector([10, 0])}}, - {nested: {foo: 'bar', embedding: FieldValue.vector([20, 0])}}, - {nested: {foo: 'bar', embedding: FieldValue.vector([100, 100])}}, - ]); + const res = await vectorQuery.get(); - const vectorQuery = indexTestHelper - .query(collectionReference) - .findNearest('nested.embedding', [10, 10], { - limit: 3, - distanceMeasure: 'EUCLIDEAN', + expect(res.size).to.equal(4); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 0]))) + .to.be.true; + expect(res.docs[0].get('distance')).to.equal(0); + + expect(res.docs[1].get('distance')).to.equal(1); + expect(res.docs[2].get('distance')).to.equal(1); + + expect(res.docs[3].get('embedding').isEqual(FieldValue.vector([-1, 0]))) + .to.be.true; + expect(res.docs[3].get('distance')).to.equal(2); }); + }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(3); - expect( - res.docs[0].get('nested.embedding').isEqual(FieldValue.vector([10, 10])) - ).to.be.true; - expect( - res.docs[1].get('nested.embedding').isEqual(FieldValue.vector([10, 0])) - ).to.be.true; - expect( - res.docs[2].get('nested.embedding').isEqual(FieldValue.vector([1, 1])) - ).to.be.true; - }); + describe('querying with distance threshold', () => { + it('supports COSINE distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - it('supports findNearest with select to exclude vector data in response', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([1, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 1])}, + '4': {foo: 'bar', embedding: FieldValue.vector([0, -0.1])}, + '5': {foo: 'bar', embedding: FieldValue.vector([-1, 0])}, + }); - const collectionReference = await indexTestHelper.createTestDocs([ - {foo: 1}, - {foo: 2, embedding: FieldValue.vector([10, 10])}, - {foo: 3, embedding: FieldValue.vector([1, 1])}, - {foo: 4, embedding: FieldValue.vector([10, 0])}, - {foo: 5, embedding: FieldValue.vector([20, 0])}, - {foo: 6, embedding: FieldValue.vector([100, 100])}, - ]); + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'COSINE', + distanceThreshold: 1, + }); - const vectorQuery = indexTestHelper - .query(collectionReference) - .where('foo', 'in', [1, 2, 3, 4, 5, 6]) - .select('foo') - .findNearest('embedding', [10, 10], { - limit: 10, - distanceMeasure: 'EUCLIDEAN', + const res = await vectorQuery.get(); + + expect(res.size).to.equal(3); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([1, 0]))) + .to.be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 1]))) + .to.be.true; + expect( + res.docs[2].get('embedding').isEqual(FieldValue.vector([0, -0.1])) + ).to.be.true; }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(5); - expect(res.docs[0].get('foo')).to.equal(2); - expect(res.docs[1].get('foo')).to.equal(4); - expect(res.docs[2].get('foo')).to.equal(3); - expect(res.docs[3].get('foo')).to.equal(5); - expect(res.docs[4].get('foo')).to.equal(6); + it('supports EUCLIDEAN distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); - res.docs.forEach(ds => expect(ds.get('embedding')).to.be.undefined); - }); + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([1, -0.1])}, + '5': {foo: 'bar', embedding: FieldValue.vector([4, 4])}, + }); - it('supports findNearest limits', async () => { - const indexTestHelper = new IndexTestHelper(firestore); + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'EUCLIDEAN', + distanceThreshold: 5, + }); - const embeddingVector = []; - const queryVector = []; - for (let i = 0; i < 2048; i++) { - embeddingVector.push(i + 1); - queryVector.push(i - 1); - } + const res = await vectorQuery.get(); - const collectionReference = await indexTestHelper.createTestDocs([ - {embedding: FieldValue.vector(embeddingVector)}, - ]); + expect(res.size).to.equal(3); - const vectorQuery = indexTestHelper - .query(collectionReference) - .findNearest('embedding', queryVector, { - limit: 1000, - distanceMeasure: 'EUCLIDEAN', + expect( + res.docs[0].get('embedding').isEqual(FieldValue.vector([1, -0.1])) + ).to.be.true; + expect(res.docs[1].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + expect(res.docs[2].get('embedding').isEqual(FieldValue.vector([4, 4]))) + .to.be.true; }); - const res = await vectorQuery.get(); - expect(res.size).to.equal(1); - expect( - (res.docs[0].get('embedding') as VectorValue).toArray() - ).to.deep.equal(embeddingVector); + it('supports DOT_PRODUCT distance', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([-20, 0])}, + '5': {foo: 'bar', embedding: FieldValue.vector([0.1, 4])}, + }); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'DOT_PRODUCT', + distanceThreshold: 1, + }); + + const res = await vectorQuery.get(); + + expect(res.size).to.equal(2); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 100])) + ).to.be.true; + }); + + it('works with distance threshold', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([-20, 0])}, + '5': {foo: 'bar', embedding: FieldValue.vector([0.1, 4])}, + }); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 5, + distanceMeasure: 'DOT_PRODUCT', + distanceThreshold: 0.11, + distanceResultField: 'foo', + }); + + const res = await vectorQuery.get(); + + expect(res.size).to.equal(2); + + expect(res.docs[0].get('foo')).to.equal(2); + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + + expect(res.docs[1].get('foo')).to.equal(1); + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 100])) + ).to.be.true; + }); + + it('will not exceed limit even if there are more results more similar than distanceThreshold', async () => { + const indexTestHelper = new IndexTestHelper(firestore); + + const collectionReference = await indexTestHelper.setTestDocs({ + '1': {foo: 'bar'}, + '2': {foo: 'bar', embedding: FieldValue.vector([2, 0])}, + '3': {foo: 'bar', embedding: FieldValue.vector([1, 100])}, + '4': {foo: 'bar', embedding: FieldValue.vector([-20, 0])}, + '5': {foo: 'bar', embedding: FieldValue.vector([0.1, 4])}, + }); + + const vectorQuery = indexTestHelper + .query(collectionReference) + .findNearest({ + vectorField: 'embedding', + queryVector: [1, 0], + limit: 2, + distanceMeasure: 'DOT_PRODUCT', + distanceThreshold: 0.0, + }); + + const res = await vectorQuery.get(); + + expect(res.size).to.equal(2); + + expect(res.docs[0].get('embedding').isEqual(FieldValue.vector([2, 0]))) + .to.be.true; + + expect( + res.docs[1].get('embedding').isEqual(FieldValue.vector([1, 100])) + ).to.be.true; + }); + }); }); it('supports !=', async () => { diff --git a/dev/test/vector-query.ts b/dev/test/vector-query.ts index 4da0d6e66..9897e0afc 100644 --- a/dev/test/vector-query.ts +++ b/dev/test/vector-query.ts @@ -25,6 +25,7 @@ import { import { DocumentSnapshot, FieldValue, + FieldPath, Firestore, Query, Timestamp, @@ -86,12 +87,16 @@ describe('Vector(findNearest) query interface', () => { expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], limit: 10, distanceMeasure: 'COSINE', }) .isEqual( - queryA.findNearest('embedding', [40, 41, 42], { + queryA.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], distanceMeasure: 'COSINE', limit: 10, }) @@ -99,41 +104,110 @@ describe('Vector(findNearest) query interface', () => { ).to.be.true; expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', limit: 10, - distanceMeasure: 'COSINE', }) .isEqual( - queryB.findNearest('embedding', [40, 41, 42], { - distanceMeasure: 'COSINE', + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', limit: 10, }) ) ).to.be.true; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 0.125, + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 0.125, + }) + ) + ).to.be.true; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 0.125, + distanceResultField: new FieldPath('foo'), + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 0.125, + distanceResultField: new FieldPath('foo'), + }) + ) + ).to.be.true; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: 'distance', + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: new FieldPath('distance'), + }) + ) + ).to.be.true; expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], limit: 10, distanceMeasure: 'COSINE', }) .isEqual( - firestore - .collection('collectionId') - .findNearest('embedding', [40, 41, 42], { - distanceMeasure: 'COSINE', - limit: 10, - }) + firestore.collection('collectionId').findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'COSINE', + limit: 10, + }) ) ).to.be.false; expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], limit: 10, distanceMeasure: 'COSINE', }) .isEqual( - queryB.findNearest('embedding', [40, 42], { + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 42], distanceMeasure: 'COSINE', limit: 10, }) @@ -141,12 +215,16 @@ describe('Vector(findNearest) query interface', () => { ).to.be.false; expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], limit: 10, distanceMeasure: 'COSINE', }) .isEqual( - queryB.findNearest('embedding', [40, 42], { + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 42], distanceMeasure: 'COSINE', limit: 1000, }) @@ -154,19 +232,188 @@ describe('Vector(findNearest) query interface', () => { ).to.be.false; expect( queryA - .findNearest('embedding', [40, 41, 42], { + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], limit: 10, distanceMeasure: 'COSINE', }) .isEqual( - queryB.findNearest('embedding', [40, 42], { + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 42], distanceMeasure: 'EUCLIDEAN', - limit: 1000, + limit: 10, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 1.125, + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 0.125, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 1, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceThreshold: 1, + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: 'distance', + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: 'result', + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: new FieldPath('bar'), + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: new FieldPath('foo'), + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: new FieldPath('foo'), + }) + ) + ).to.be.false; + expect( + queryA + .findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, + distanceResultField: 'result', + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'EUCLIDEAN', + limit: 10, }) ) ).to.be.false; }); + it('generates equal vector queries with deprecated API', () => { + const queryA = firestore.collection('collectionId').where('foo', '==', 42); + const queryB = firestore.collection('collectionId').where('foo', '==', 42); + + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + queryB.findNearest({ + vectorField: 'embedding', + queryVector: [40, 41, 42], + distanceMeasure: 'COSINE', + limit: 10, + }) + ) + ).to.be.true; + expect( + queryA + .findNearest('foo', [40, 41, 42, 43], { + limit: 1, + distanceMeasure: 'DOT_PRODUCT', + }) + .isEqual( + queryB.findNearest({ + vectorField: 'foo', + queryVector: [40, 41, 42, 43], + distanceMeasure: 'DOT_PRODUCT', + limit: 1, + }) + ) + ).to.be.true; + }); + it('generates proto', async () => { const overrides: ApiOverride = { runQuery: request => { @@ -193,6 +440,26 @@ describe('Vector(findNearest) query interface', () => { }); it('validates inputs', async () => { + const query: Query = firestore.collection('collectionId'); + expect(() => { + query.findNearest({ + vectorField: 'embedding', + queryVector: [], + limit: 10, + distanceMeasure: 'EUCLIDEAN', + }); + }).to.throw('not a valid vector size'); + expect(() => { + query.findNearest({ + vectorField: 'embedding', + queryVector: [10, 1000], + limit: 0, + distanceMeasure: 'EUCLIDEAN', + }); + }).to.throw('not a valid positive limit number'); + }); + + it('validates inputs - preview (deprecated) API', async () => { const query: Query = firestore.collection('collectionId'); expect(() => { query.findNearest('embedding', [], { @@ -227,12 +494,12 @@ describe('Vector(findNearest) query interface', () => { return createInstance(overrides).then(firestoreInstance => { firestore = firestoreInstance; - const query = firestore - .collection('collectionId') - .findNearest('embedding', [1], { - limit: 2, - distanceMeasure: distanceMeasure, - }); + const query = firestore.collection('collectionId').findNearest({ + vectorField: 'embedding', + queryVector: [1], + limit: 2, + distanceMeasure: distanceMeasure, + }); return query.get().then(results => { expect(results.size).to.equal(2); expect(results.empty).to.be.false; @@ -268,9 +535,12 @@ describe('Vector(findNearest) query interface', () => { let counter = 0; return createInstance(overrides).then(firestoreInstance => { firestore = firestoreInstance; - const query = firestore - .collection('collectionId') - .findNearest('vector', [1], {limit: 10, distanceMeasure: 'COSINE'}); + const query = firestore.collection('collectionId').findNearest({ + vectorField: 'vector', + queryVector: [1], + limit: 10, + distanceMeasure: 'COSINE', + }); return query.get().then(results => { expect(++counter).to.equal(1); expect(results.size).to.equal(2); @@ -285,12 +555,12 @@ describe('Vector(findNearest) query interface', () => { it('handles stream exception at initialization', async () => { let attempts = 0; - const query = firestore - .collection('collectionId') - .findNearest('embedding', [1], { - limit: 100, - distanceMeasure: 'EUCLIDEAN', - }); + const query = firestore.collection('collectionId').findNearest({ + vectorField: 'embedding', + queryVector: [1], + limit: 100, + distanceMeasure: 'EUCLIDEAN', + }); query._queryUtil._stream = () => { ++attempts; @@ -322,7 +592,12 @@ describe('Vector(findNearest) query interface', () => { firestore = firestoreInstance; return firestore .collection('collectionId') - .findNearest('embedding', [1], {limit: 10, distanceMeasure: 'COSINE'}) + .findNearest({ + vectorField: 'embedding', + queryVector: [1], + limit: 10, + distanceMeasure: 'COSINE', + }) .get() .then(() => { throw new Error('Unexpected success in Promise'); diff --git a/package.json b/package.json index c82e0e890..655543ee3 100644 --- a/package.json +++ b/package.json @@ -73,9 +73,9 @@ "@google-cloud/trace-agent": "^8.0.0", "@googleapis/cloudtrace": "^1.1.2", "@google-cloud/cloud-rad": "^0.4.0", - "@opentelemetry/sdk-trace-node": "^1.24.1", - "@opentelemetry/context-async-hooks": "^1.24.1", "@google-cloud/opentelemetry-cloud-trace-exporter": "^2.0.0", + "@opentelemetry/context-async-hooks": "^1.24.1", + "@opentelemetry/sdk-trace-node": "^1.24.1", "@types/assert": "^1.4.0", "@types/chai": "^4.2.7", "@types/chai-as-promised": "^7.1.2", diff --git a/types/firestore.d.ts b/types/firestore.d.ts index 702e5ee84..51113bb1d 100644 --- a/types/firestore.d.ts +++ b/types/firestore.d.ts @@ -20,6 +20,7 @@ // Declare a global (ambient) namespace // (used when not using import statement, but just script include). + declare namespace FirebaseFirestore { /** Alias for `any` but used where a Firestore field value would be provided. */ export type DocumentFieldValue = any; @@ -2035,11 +2036,11 @@ declare namespace FirebaseFirestore { * `vectorField` against the given `queryVector` and returns the top documents that are closest * to the `queryVector`. * - * Only documents whose `vectorField` field is a `VectorValue` of the same dimension as `queryVector` + * Only documents whose `vectorField` field is a {@link VectorValue} of the same dimension as `queryVector` * participate in the query, all other documents are ignored. * * @example - * ```typescript + * ``` * // Returns the closest 10 documents whose Euclidean distance from their 'embedding' fields are closed to [41, 42]. * const vectorQuery = col.findNearest('embedding', [41, 42], {limit: 10, distanceMeasure: 'EUCLIDEAN'}); * @@ -2047,11 +2048,14 @@ declare namespace FirebaseFirestore { * querySnapshot.forEach(...); * ``` * - * @param vectorField The field path this vector query executes on. - * @param queryVector The vector value used to measure the distance from `vectorField` values in the documents. - * @param options Options control the vector query. `limit` specifies the upper bound of documents to return, must - * be a positive integer with a maximum value of 1000. `distanceMeasure` specifies what type of distance is - * calculated when performing the query. + * @param vectorField - A string or {@link FieldPath} specifying the vector field to search on. + * @param queryVector - The {@link VectorValue} used to measure the distance from `vectorField` values in the documents. + * @param options - Options control the vector query. `limit` specifies the upper bound of documents to return, must + * be a positive integer with a maximum value of 1000. `distanceMeasure` specifies what type of distance is calculated + * when performing the query. + * + * @deprecated Use the new {@link findNearest} implementation + * accepting a single `options` param. */ findNearest( vectorField: string | FieldPath, @@ -2062,6 +2066,38 @@ declare namespace FirebaseFirestore { } ): VectorQuery; + /** + * Returns a query that can perform vector distance (similarity) search with given parameters. + * + * The returned query, when executed, performs a distance (similarity) search on the specified + * `vectorField` against the given `queryVector` and returns the top documents that are closest + * to the `queryVector`. + * + * Only documents whose `vectorField` field is a {@link VectorValue} of the same dimension as `queryVector` + * participate in the query, all other documents are ignored. + * + * @example + * ``` + * // Returns the closest 10 documents whose Euclidean distance from their 'embedding' fields are closed to [41, 42]. + * const vectorQuery = col.findNearest({ + * vectorField: 'embedding', + * queryVector: [41, 42], + * limit: 10, + * distanceMeasure: 'EUCLIDEAN', + * distanceResultField: 'distance', + * distanceThreshold: 0.125 + * }); + * + * const querySnapshot = await aggregateQuery.get(); + * querySnapshot.forEach(...); + * ``` + * @param options - An argument specifying the behavior of the {@link VectorQuery} returned by this function. + * See {@link VectorQueryOptions}. + */ + findNearest( + options: VectorQueryOptions + ): VectorQuery; + /** * Returns true if this `Query` is equal to the provided one. * @@ -3192,6 +3228,50 @@ declare namespace FirebaseFirestore { */ readonly snapshot: T | null; } + + /** + * Specifies the behavior of the {@link VectorQuery} generated by a call to {@link Query.findNearest}. + */ + export interface VectorQueryOptions { + /** + * A string or {@link FieldPath} specifying the vector field to search on. + */ + vectorField: string | FieldPath; + + /** + * The {@link VectorValue} used to measure the distance from `vectorField` values in the documents. + */ + queryVector: VectorValue | Array; + + /** + * Specifies the upper bound of documents to return, must be a positive integer with a maximum value of 1000. + */ + limit: number; + + /** + * Specifies what type of distance is calculated when performing the query. + */ + distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + + /** + * Optionally specifies the name of a field that will be set on each returned DocumentSnapshot, + * which will contain the computed distance for the document. + */ + distanceResultField?: string | FieldPath; + + /** + * Specifies a threshold for which no less similar documents will be returned. The behavior + * of the specified `distanceMeasure` will affect the meaning of the distance threshold. + * + * - For `distanceMeasure: "EUCLIDEAN"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE euclidean_distance <= distanceThreshold + * - For `distanceMeasure: "COSINE"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE cosine_distance <= distanceThreshold + * - For `distanceMeasure: "DOT_PRODUCT"`, the meaning of `distanceThreshold` is: + * SELECT docs WHERE dot_product_distance >= distanceThreshold + */ + distanceThreshold?: number; + } } declare module '@google-cloud/firestore' { From 8118d37bf27b49d5e695bf18a9550a37ecc34ca2 Mon Sep 17 00:00:00 2001 From: "release-please[bot]" <55107282+release-please[bot]@users.noreply.github.com> Date: Thu, 5 Sep 2024 16:17:48 -0600 Subject: [PATCH 2/2] chore(main): release 7.10.0 (#2168) Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com> --- CHANGELOG.md | 9 +++++++++ package.json | 2 +- samples/package.json | 2 +- 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b5f254a8e..57e34cc6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,15 @@ [1]: https://www.npmjs.com/package/@google-cloud/firestore?activeTab=versions +## [7.10.0](https://github.com/googleapis/nodejs-firestore/compare/v7.9.0...v7.10.0) (2024-09-05) + + +### Features + +* Expose proto changes for the bulk delete api ([23ce891](https://github.com/googleapis/nodejs-firestore/commit/23ce89175ce315648ce9af3994cba0decc48ba47)) +* Expose the proto changes to support FindNearest.distance_result_field parameter and the FindNearest.distance_threshold parameter ([23ce891](https://github.com/googleapis/nodejs-firestore/commit/23ce89175ce315648ce9af3994cba0decc48ba47)) +* Return computed distance and set distance thresholds on VectorQueries ([#2090](https://github.com/googleapis/nodejs-firestore/issues/2090)) ([b5ca84f](https://github.com/googleapis/nodejs-firestore/commit/b5ca84f076ca0668e90ca3fc7dd878f732ccd956)) + ## [7.9.0](https://github.com/googleapis/nodejs-firestore/compare/v7.8.0...v7.9.0) (2024-06-25) diff --git a/package.json b/package.json index 655543ee3..bca4d84c5 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "@google-cloud/firestore", "description": "Firestore Client Library for Node.js", - "version": "7.9.0", + "version": "7.10.0", "license": "Apache-2.0", "author": "Google Inc.", "engines": { diff --git a/samples/package.json b/samples/package.json index c2328f7ce..d3a07204d 100644 --- a/samples/package.json +++ b/samples/package.json @@ -11,7 +11,7 @@ "test": "mocha --timeout 600000" }, "dependencies": { - "@google-cloud/firestore": "^7.9.0" + "@google-cloud/firestore": "^7.10.0" }, "devDependencies": { "chai": "^4.2.0",