diff --git a/api-report/firestore.api.md b/api-report/firestore.api.md index 483c28c6a..921ba16cf 100644 --- a/api-report/firestore.api.md +++ b/api-report/firestore.api.md @@ -717,6 +717,7 @@ export class FieldValue implements firestore.FieldValue { // Warning: (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" // Warning: (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration static serverTimestamp(): FieldValue; + static vector(values?: number[]): VectorValue; } // @public @@ -1003,6 +1004,10 @@ export class Query; }, AppModelType, DbModelType>; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + _createSnapshot(readTime: Timestamp, size: number, docs: () => Array>, changes: () => Array>): QuerySnapshot; // Warning: (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen // Warning: (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' // Warning: (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag @@ -1028,6 +1033,10 @@ export class Query, options: { + limit: number; + distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + }): 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 @@ -1047,7 +1056,7 @@ export class Query>; // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // - // @internal (undocumented) + // @internal _hasRetryTimedOut(methodName: string, startTime: number): boolean; // Warning: (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen // Warning: (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' @@ -1106,12 +1115,21 @@ export class Query; + readonly _queryOptions: QueryOptions; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // Warning: (ae-forgotten-export) The symbol "QueryUtil" needs to be exported by the entry point index.d.ts + // + // @internal (undocumented) + readonly _queryUtil: QueryUtil>; // Warning: (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen // Warning: (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' // 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 "{@" select(...fieldPaths: Array): Query; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal (undocumented) + readonly _serializer: Serializer; // Warning: (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen // Warning: (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' // Warning: (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag @@ -1424,6 +1442,74 @@ export class Transaction implements firestore.Transaction { update(documentRef: firestore.DocumentReference, dataOrField: firestore.UpdateData | string | firestore.FieldPath, ...preconditionOrValues: Array): 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); + // 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; + get(): Promise>; + isEqual(other: firestore.VectorQuery): boolean; + get query(): Query; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal (undocumented) + readonly _queryUtil: QueryUtil>; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + startAfter(...fieldValuesOrDocumentSnapshot: Array): VectorQuery; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + _stream(transactionId?: Uint8Array): NodeJS.ReadableStream; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + toProto(transactionIdOrReadTime?: Uint8Array | Timestamp): api.IRunQueryRequest; +} + +// @public +export class VectorQuerySnapshot implements firestore.VectorQuerySnapshot { + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + constructor(_query: VectorQuery, _readTime: Timestamp, _size: number, docs: () => Array>, changes: () => Array>); + docChanges(): Array>; + get docs(): Array>; + get empty(): boolean; + forEach(callback: (result: firestore.QueryDocumentSnapshot) => void, thisArg?: unknown): void; + isEqual(other: firestore.VectorQuerySnapshot): boolean; + get query(): VectorQuery; + get readTime(): Timestamp; + get size(): number; +} + +// Warning: (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration +// +// @public +export class VectorValue implements firestore.VectorValue { + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal + constructor(values: number[] | undefined); + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal (undocumented) + static _fromProto(valueArray: api.IValue): VectorValue; + isEqual(other: VectorValue): boolean; + toArray(): number[]; + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration + // + // @internal (undocumented) + _toProto(serializer: Serializer): api.IValue; +} + // Warning: (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration // // @public @@ -1570,24 +1656,26 @@ export class WriteResult implements firestore.WriteResult { // build/src/reference.d.ts:367:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration // build/src/reference.d.ts:398:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration // build/src/reference.d.ts:400:4 - (tsdoc-undefined-tag) The TSDoc tag "@class" is not defined in this configuration -// build/src/reference.d.ts:629:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference.d.ts:999:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration -// build/src/reference.d.ts:1005:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference.d.ts:1007:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1015:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference.d.ts:1017:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1017:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' -// build/src/reference.d.ts:1019:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1019:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' -// build/src/reference.d.ts:1021:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1023:24 - (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag -// build/src/reference.d.ts:1023:17 - (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" -// build/src/reference.d.ts:1032:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1034:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration -// build/src/reference.d.ts:1036:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference.d.ts:1220:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen -// build/src/reference.d.ts:1221:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration -// build/src/reference.d.ts:1478:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:800:4 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1228:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration +// build/src/reference.d.ts:1234:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1236:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1244:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1246:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1246:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' +// build/src/reference.d.ts:1248:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1248:15 - (tsdoc-param-tag-with-invalid-type) The @param block should not include a JSDoc-style '{type}' +// build/src/reference.d.ts:1250:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1252:24 - (tsdoc-escape-right-brace) The "}" character should be escaped using a backslash to avoid confusion with a TSDoc inline tag +// build/src/reference.d.ts:1252:17 - (tsdoc-malformed-inline-tag) Expecting a TSDoc tag starting with "{@" +// build/src/reference.d.ts:1261:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1263:8 - (tsdoc-undefined-tag) The TSDoc tag "@return" is not defined in this configuration +// build/src/reference.d.ts:1265:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1449:8 - (tsdoc-param-tag-missing-hyphen) The @param block should be followed by a parameter name and then a hyphen +// build/src/reference.d.ts:1450:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1715:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1825:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration +// build/src/reference.d.ts:1830: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/transaction.d.ts:239:8 - (tsdoc-undefined-tag) The TSDoc tag "@private" is not defined in this configuration diff --git a/dev/src/aggregate.ts b/dev/src/aggregate.ts index 8c5f73401..988a33901 100644 --- a/dev/src/aggregate.ts +++ b/dev/src/aggregate.ts @@ -1,4 +1,4 @@ -/*! +/** * Copyright 2023 Google LLC. All Rights Reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/dev/src/convert.ts b/dev/src/convert.ts index 713675daf..6e141a6e4 100644 --- a/dev/src/convert.ts +++ b/dev/src/convert.ts @@ -20,6 +20,7 @@ import {ApiMapValue, ProtobufJsValue} from './types'; import {validateObject} from './validate'; import api = google.firestore.v1; +import {RESERVED_MAP_KEY, RESERVED_MAP_KEY_VECTOR_VALUE} from './map-type'; /*! * @module firestore/convert @@ -112,53 +113,72 @@ function bytesFromJson(bytesValue: string | Uint8Array): Uint8Array { * @return The string value for 'valueType'. */ export function detectValueType(proto: ProtobufJsValue): string { + let valueType: string | undefined; + if (proto.valueType) { - return proto.valueType; - } + valueType = proto.valueType; + } else { + const detectedValues: string[] = []; - const detectedValues: string[] = []; + if (proto.stringValue !== undefined) { + detectedValues.push('stringValue'); + } + if (proto.booleanValue !== undefined) { + detectedValues.push('booleanValue'); + } + if (proto.integerValue !== undefined) { + detectedValues.push('integerValue'); + } + if (proto.doubleValue !== undefined) { + detectedValues.push('doubleValue'); + } + if (proto.timestampValue !== undefined) { + detectedValues.push('timestampValue'); + } + if (proto.referenceValue !== undefined) { + detectedValues.push('referenceValue'); + } + if (proto.arrayValue !== undefined) { + detectedValues.push('arrayValue'); + } + if (proto.nullValue !== undefined) { + detectedValues.push('nullValue'); + } + if (proto.mapValue !== undefined) { + detectedValues.push('mapValue'); + } + if (proto.geoPointValue !== undefined) { + detectedValues.push('geoPointValue'); + } + if (proto.bytesValue !== undefined) { + detectedValues.push('bytesValue'); + } - if (proto.stringValue !== undefined) { - detectedValues.push('stringValue'); - } - if (proto.booleanValue !== undefined) { - detectedValues.push('booleanValue'); - } - if (proto.integerValue !== undefined) { - detectedValues.push('integerValue'); - } - if (proto.doubleValue !== undefined) { - detectedValues.push('doubleValue'); - } - if (proto.timestampValue !== undefined) { - detectedValues.push('timestampValue'); - } - if (proto.referenceValue !== undefined) { - detectedValues.push('referenceValue'); - } - if (proto.arrayValue !== undefined) { - detectedValues.push('arrayValue'); - } - if (proto.nullValue !== undefined) { - detectedValues.push('nullValue'); - } - if (proto.mapValue !== undefined) { - detectedValues.push('mapValue'); - } - if (proto.geoPointValue !== undefined) { - detectedValues.push('geoPointValue'); - } - if (proto.bytesValue !== undefined) { - detectedValues.push('bytesValue'); + if (detectedValues.length !== 1) { + throw new Error( + `Unable to infer type value from '${JSON.stringify(proto)}'.` + ); + } + + valueType = detectedValues[0]; } - if (detectedValues.length !== 1) { - throw new Error( - `Unable to infer type value from '${JSON.stringify(proto)}'.` - ); + // Special handling of mapValues used to represent other data types + if (valueType === 'mapValue') { + const fields = proto.mapValue?.fields; + if (fields) { + const props = Object.keys(fields); + if ( + props.indexOf(RESERVED_MAP_KEY) !== -1 && + detectValueType(fields[RESERVED_MAP_KEY]) === 'stringValue' && + fields[RESERVED_MAP_KEY].stringValue === RESERVED_MAP_KEY_VECTOR_VALUE + ) { + valueType = 'vectorValue'; + } + } } - return detectedValues[0]; + return valueType; } /** @@ -240,7 +260,8 @@ export function valueFromJson(fieldValue: api.IValue): api.IValue { }, }; } - case 'mapValue': { + case 'mapValue': + case 'vectorValue': { const mapValue: ApiMapValue = {}; const fields = fieldValue.mapValue!.fields; if (fields) { diff --git a/dev/src/field-value.ts b/dev/src/field-value.ts index eb054067c..4e98b7201 100644 --- a/dev/src/field-value.ts +++ b/dev/src/field-value.ts @@ -22,6 +22,7 @@ import * as proto from '../protos/firestore_v1_proto_api'; import {FieldPath} from './path'; import {Serializer, validateUserInput} from './serializer'; +import {isPrimitiveArrayEqual} from './util'; import { invalidArgumentMessage, validateMinNumberOfArguments, @@ -30,6 +31,58 @@ import { import api = proto.google.firestore.v1; +/** + * Represent a vector type in Firestore documents. + * Create an instance with {@link FieldValue.vector}. + * + * @class VectorValue + */ +export class VectorValue implements firestore.VectorValue { + private readonly _values: number[]; + + /** + * @private + * @internal + */ + constructor(values: number[] | undefined) { + // Making a copy of the parameter. + this._values = (values || []).map(n => n); + } + + /** + * Returns a copy of the raw number array form of the vector. + */ + public toArray(): number[] { + return this._values.map(n => n); + } + + /** + * @private + * @internal + */ + _toProto(serializer: Serializer): api.IValue { + return serializer.encodeVector(this._values); + } + + /** + * @private + * @internal + */ + static _fromProto(valueArray: api.IValue): VectorValue { + const values = valueArray.arrayValue?.values?.map(v => { + return v.doubleValue!; + }); + return new VectorValue(values); + } + + /** + * Returns `true` if the two VectorValue has the same raw number arrays, returns `false` otherwise. + */ + isEqual(other: VectorValue): boolean { + return isPrimitiveArrayEqual(this._values, other._values); + } +} + /** * Sentinel values that can be used when writing documents with set(), create() * or update(). @@ -40,6 +93,17 @@ export class FieldValue implements firestore.FieldValue { /** @private */ constructor() {} + /** + * Creates a new `VectorValue` constructed with a copy of the given array of numbers. + * + * @param values - Create a `VectorValue` instance with a copy of this array of numbers. + * + * @returns A new `VectorValue` constructed with a copy of the given array of numbers. + */ + static vector(values?: number[]): VectorValue { + return new VectorValue(values); + } + /** * Returns a sentinel for use with update() or set() with {merge:true} to mark * a field for deletion. diff --git a/dev/src/index.ts b/dev/src/index.ts index 7998f145a..e7480d48d 100644 --- a/dev/src/index.ts +++ b/dev/src/index.ts @@ -91,12 +91,17 @@ export { QuerySnapshot, Query, } from './reference'; -export type {AggregateQuery, AggregateQuerySnapshot} from './reference'; +export type { + AggregateQuery, + AggregateQuerySnapshot, + VectorQuery, + VectorQuerySnapshot, +} from './reference'; export {BulkWriter} from './bulk-writer'; export type {BulkWriterError} from './bulk-writer'; export type {BundleBuilder} from './bundle'; export {DocumentSnapshot, QueryDocumentSnapshot} from './document'; -export {FieldValue} from './field-value'; +export {FieldValue, VectorValue} from './field-value'; export {Filter} from './filter'; export {WriteBatch, WriteResult} from './write-batch'; export {Transaction} from './transaction'; diff --git a/dev/src/map-type.ts b/dev/src/map-type.ts new file mode 100644 index 000000000..f87622071 --- /dev/null +++ b/dev/src/map-type.ts @@ -0,0 +1,19 @@ +/*! + * Copyright 2024 Google LLC. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +export const RESERVED_MAP_KEY = '__type__'; +export const RESERVED_MAP_KEY_VECTOR_VALUE = '__vector__'; +export const VECTOR_MAP_VECTORS_KEY = 'value'; diff --git a/dev/src/order.ts b/dev/src/order.ts index bcb854366..faa003107 100644 --- a/dev/src/order.ts +++ b/dev/src/order.ts @@ -34,7 +34,8 @@ enum TypeOrder { REF = 6, GEO_POINT = 7, ARRAY = 8, - OBJECT = 9, + VECTOR = 9, + OBJECT = 10, } /*! @@ -67,6 +68,8 @@ function typeOrder(val: api.IValue): TypeOrder { return TypeOrder.REF; case 'mapValue': return TypeOrder.OBJECT; + case 'vectorValue': + return TypeOrder.VECTOR; default: throw new Error('Unexpected value type: ' + valueType); } @@ -225,6 +228,26 @@ function compareObjects(left: ApiMapValue, right: ApiMapValue): number { return primitiveComparator(leftKeys.length, rightKeys.length); } +/*! + * @private + * @internal + */ +function compareVectors(left: ApiMapValue, right: ApiMapValue): number { + // The vector is a map, but only vector value is compared. + const leftArray = left?.['value']?.arrayValue?.values ?? []; + const rightArray = right?.['value']?.arrayValue?.values ?? []; + + const lengthCompare = primitiveComparator( + leftArray.length, + rightArray.length + ); + if (lengthCompare !== 0) { + return lengthCompare; + } + + return compareArrays(leftArray, rightArray); +} + /*! * @private * @internal @@ -267,6 +290,11 @@ export function compare(left: api.IValue, right: api.IValue): number { left.mapValue!.fields || {}, right.mapValue!.fields || {} ); + case TypeOrder.VECTOR: + return compareVectors( + left.mapValue!.fields || {}, + right.mapValue!.fields || {} + ); default: throw new Error(`Encountered unknown type order: ${leftType}`); } diff --git a/dev/src/reference.ts b/dev/src/reference.ts index 23f73ee6c..6afb5ff73 100644 --- a/dev/src/reference.ts +++ b/dev/src/reference.ts @@ -22,12 +22,15 @@ import {GoogleError} from 'google-gax'; import * as protos from '../protos/firestore_v1_proto_api'; +import {Aggregate, AggregateField, AggregateSpec} from './aggregate'; import { DocumentSnapshot, DocumentSnapshotBuilder, QueryDocumentSnapshot, } from './document'; import {DocumentChange} from './document-change'; +import {VectorValue} from './field-value'; +import {CompositeFilter, Filter, UnaryFilter} from './filter'; import {Firestore} from './index'; import {logger} from './logger'; import {compare} from './order'; @@ -45,7 +48,9 @@ import { autoId, Deferred, getTotalTimeout, + isArrayEqual, isPermanentRpcError, + isPrimitiveArrayEqual, mapToArray, requestTag, wrapError, @@ -60,8 +65,6 @@ import { import {DocumentWatch, QueryWatch} from './watch'; import {validateDocumentData, WriteBatch, WriteResult} from './write-batch'; import api = protos.google.firestore.v1; -import {CompositeFilter, Filter, UnaryFilter} from './filter'; -import {AggregateField, Aggregate, AggregateSpec} from './aggregate'; import {ExplainMetrics, ExplainResults} from './query-profile'; /** @@ -1188,6 +1191,273 @@ export class QuerySnapshot< } } +/** + * A `VectorQuerySnapshot` contains zero or more `QueryDocumentSnapshot` objects + * representing the results of a query. The documents can be accessed as an + * array via the `docs` property or enumerated using the `forEach` method. The + * number of documents can be determined via the `empty` and `size` + * properties. + */ +export class VectorQuerySnapshot< + AppModelType = firestore.DocumentData, + DbModelType extends firestore.DocumentData = firestore.DocumentData, +> implements firestore.VectorQuerySnapshot +{ + private _materializedDocs: Array< + QueryDocumentSnapshot + > | null = null; + private _materializedChanges: Array< + DocumentChange + > | null = null; + private _docs: + | (() => Array>) + | null = null; + private _changes: + | (() => Array>) + | null = null; + + /** + * @private + * @internal + * + * @param _query - The originating query. + * @param _readTime - The time when this query snapshot was obtained. + * @param _size - The number of documents in the result set. + * @param docs - A callback returning a sorted array of documents matching + * this query + * @param changes - A callback returning a sorted array of document change + * events for this snapshot. + */ + constructor( + private readonly _query: VectorQuery, + private readonly _readTime: Timestamp, + private readonly _size: number, + docs: () => Array>, + changes: () => Array> + ) { + this._docs = docs; + this._changes = changes; + } + + /** + * The `VectorQuery` on which you called get() in order to get this + * `VectorQuerySnapshot`. + * + * @readonly + * + * @example + * ``` + * let query = firestore.collection('col').where('foo', '==', 'bar'); + * + * query.findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}) + * .get().then(querySnapshot => { + * console.log(`Returned first batch of results`); + * let query = querySnapshot.query; + * return query.offset(10).get(); + * }).then(() => { + * console.log(`Returned second batch of results`); + * }); + * ``` + */ + get query(): VectorQuery { + return this._query; + } + + /** + * An array of all the documents in this `VectorQuerySnapshot`. + * + * @readonly + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then(querySnapshot => { + * let docs = querySnapshot.docs; + * for (let doc of docs) { + * console.log(`Document found at path: ${doc.ref.path}`); + * } + * }); + * ``` + */ + get docs(): Array> { + if (this._materializedDocs) { + return this._materializedDocs!; + } + this._materializedDocs = this._docs!(); + this._docs = null; + return this._materializedDocs!; + } + + /** + * `true` if there are no documents in the `VectorQuerySnapshot`. + * + * @readonly + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then(querySnapshot => { + * if (querySnapshot.empty) { + * console.log('No documents found.'); + * } + * }); + * ``` + */ + get empty(): boolean { + return this._size === 0; + } + + /** + * The number of documents in the `VectorQuerySnapshot`. + * + * @readonly + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then(querySnapshot => { + * console.log(`Found ${querySnapshot.size} documents.`); + * }); + * ``` + */ + get size(): number { + return this._size; + } + + /** + * The time this `VectorQuerySnapshot` was obtained. + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then((querySnapshot) => { + * let readTime = querySnapshot.readTime; + * console.log(`Query results returned at '${readTime.toDate()}'`); + * }); + * ``` + */ + get readTime(): Timestamp { + return this._readTime; + } + + /** + * Returns an array of the documents changes since the last snapshot. If + * this is the first snapshot, all documents will be in the list as added + * changes. + * + * @returns An array of the documents changes since the last snapshot. + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then(querySnapshot => { + * let changes = querySnapshot.docChanges(); + * for (let change of changes) { + * console.log(`A document was ${change.type}.`); + * } + * }); + * ``` + */ + docChanges(): Array> { + if (this._materializedChanges) { + return this._materializedChanges!; + } + this._materializedChanges = this._changes!(); + this._changes = null; + return this._materializedChanges!; + } + + /** + * Enumerates all of the documents in the `VectorQuerySnapshot`. This is a convenience + * method for running the same callback on each {@link QueryDocumentSnapshot} + * that is returned. + * + * @param callback - A callback to be called with a + * {@link QueryDocumentSnapshot} for each document in + * the snapshot. + * @param thisArg - The `this` binding for the callback.. + * + * @example + * ``` + * let query = firestore.collection('col') + * .findNearest("embedding", [0, 0], {limit: 10, distanceMeasure: "EUCLIDEAN"}); + * + * query.get().then(querySnapshot => { + * querySnapshot.forEach(documentSnapshot => { + * console.log(`Document found at path: ${documentSnapshot.ref.path}`); + * }); + * }); + * ``` + */ + forEach( + callback: ( + result: firestore.QueryDocumentSnapshot + ) => void, + thisArg?: unknown + ): void { + validateFunction('callback', callback); + + for (const doc of this.docs) { + callback.call(thisArg, doc); + } + } + + /** + * Returns true if the document data in this `VectorQuerySnapshot` is equal to the + * provided value. + * + * @param other - The value to compare against. + * @returns true if this `VectorQuerySnapshot` is equal to the provided + * value. + */ + isEqual( + other: firestore.VectorQuerySnapshot + ): boolean { + // Since the read time is different on every query read, we explicitly + // ignore all metadata in this comparison. + + if (this === other) { + return true; + } + + if (!(other instanceof VectorQuerySnapshot)) { + return false; + } + + if (this._size !== other._size) { + return false; + } + + if (!this._query.isEqual(other._query)) { + return false; + } + + if (this._materializedDocs && !this._materializedChanges) { + // If we have only materialized the documents, we compare them first. + return ( + isArrayEqual(this.docs, other.docs) && + isArrayEqual(this.docChanges(), other.docChanges()) + ); + } + + // Otherwise, we compare the changes first as we expect there to be fewer. + return ( + isArrayEqual(this.docChanges(), other.docChanges()) && + isArrayEqual(this.docs, other.docs) + ); + } +} + /** Internal representation of a query cursor before serialization. */ interface QueryCursor { before: boolean; @@ -1405,46 +1675,353 @@ export class QueryOptions< } } -/** - * A Query refers to a query which you can read or stream from. You can also - * construct refined Query objects by adding filters and ordering. - * - * @class Query - */ -export class Query< - AppModelType = firestore.DocumentData, - DbModelType extends firestore.DocumentData = firestore.DocumentData, -> implements firestore.Query -{ - private readonly _serializer: Serializer; - /** - * @internal - * @private - **/ - protected readonly _allowUndefined: boolean; - - /** - * @internal - * @private - * - * @param _firestore The Firestore Database client. - * @param _queryOptions Options that define the query. - */ +class QueryUtil< + AppModelType, + DbModelType extends firestore.DocumentData, + Template extends + | Query + | VectorQuery, +> { constructor( - /** - * @internal - * @private - **/ + /** @private */ readonly _firestore: Firestore, - /** - * @internal - * @private + /** @private */ + readonly _queryOptions: QueryOptions, + /** @private */ + readonly _serializer: Serializer + ) {} + + _get( + query: Template, + transactionIdOrReadTime?: Uint8Array | Timestamp, + retryWithCursor = true + ): Promise< + | QuerySnapshot + | VectorQuerySnapshot + > { + const docs: Array> = []; + + // Capture the error stack to preserve stack tracing across async calls. + const stack = Error().stack!; + + return new Promise((resolve, reject) => { + let readTime: Timestamp; + + this._stream(query, transactionIdOrReadTime, retryWithCursor) + .on('error', err => { + reject(wrapError(err, stack)); + }) + .on('data', result => { + readTime = result.readTime; + if (result.document) { + docs.push(result.document); + } + }) + .on('end', () => { + if (this._queryOptions.limitType === LimitType.Last) { + // The results for limitToLast queries need to be flipped since + // we reversed the ordering constraints before sending the query + // to the backend. + docs.reverse(); + } + + resolve( + query._createSnapshot( + readTime, + docs.length, + () => docs, + () => { + const changes: Array< + DocumentChange + > = []; + for (let i = 0; i < docs.length; ++i) { + changes.push(new DocumentChange('added', docs[i], -1, i)); + } + return changes; + } + ) + ); + }); + }); + } + + // This method exists solely to enable unit tests to mock it. + _isPermanentRpcError(err: GoogleError, methodName: string): boolean { + return isPermanentRpcError(err, methodName); + } + + _hasRetryTimedOut(methodName: string, startTime: number): boolean { + const totalTimeout = getTotalTimeout(methodName); + if (totalTimeout === 0) { + return false; + } + + return Date.now() - startTime >= totalTimeout; + } + + stream(query: Template): NodeJS.ReadableStream { + if (this._queryOptions.limitType === LimitType.Last) { + throw new Error( + 'Query results for queries that include limitToLast() ' + + 'constraints cannot be streamed. Use Query.get() instead.' + ); + } + + const responseStream = this._stream(query); + const transform = new Transform({ + objectMode: true, + transform(chunk, encoding, callback) { + callback(undefined, chunk.document); + }, + }); + + responseStream.pipe(transform); + responseStream.on('error', e => transform.destroy(e)); + return transform; + } + + _stream( + query: Template, + transactionIdOrReadTime?: Uint8Array | Timestamp, + retryWithCursor = true, + explainOptions?: firestore.ExplainOptions + ): NodeJS.ReadableStream { + const tag = requestTag(); + const startTime = Date.now(); + const isExplain = explainOptions !== undefined; + + let lastReceivedDocument: QueryDocumentSnapshot< + AppModelType, + DbModelType + > | null = null; + + let backendStream: Duplex; + const stream = new Transform({ + objectMode: true, + transform: (proto, enc, callback) => { + if (proto === NOOP_MESSAGE) { + callback(undefined); + return; + } + + const output: { + readTime?: Timestamp; + document?: QueryDocumentSnapshot; + explainMetrics?: ExplainMetrics; + } = {}; + + if (proto.readTime) { + output.readTime = Timestamp.fromProto(proto.readTime); + } + + if (proto.document) { + const document = this._firestore.snapshot_( + proto.document, + proto.readTime + ); + const finalDoc = new DocumentSnapshotBuilder< + AppModelType, + DbModelType + >(document.ref.withConverter(this._queryOptions.converter)); + // Recreate the QueryDocumentSnapshot with the DocumentReference + // containing the original converter. + finalDoc.fieldsProto = document._fieldsProto; + finalDoc.readTime = document.readTime; + finalDoc.createTime = document.createTime; + finalDoc.updateTime = document.updateTime; + lastReceivedDocument = finalDoc.build() as QueryDocumentSnapshot< + AppModelType, + DbModelType + >; + output.document = lastReceivedDocument; + } + + if (proto.explainMetrics) { + output.explainMetrics = ExplainMetrics._fromProto( + proto.explainMetrics, + this._serializer + ); + } + + callback(undefined, output); + + if (proto.done) { + logger('QueryUtil._stream', tag, 'Trigger Logical Termination.'); + backendStream.unpipe(stream); + backendStream.resume(); + backendStream.end(); + stream.end(); + } + }, + }); + + this._firestore + .initializeIfNeeded(tag) + .then(async () => { + // `toProto()` might throw an exception. We rely on the behavior of an + // async function to convert this exception into the rejected Promise we + // catch below. + let request = query.toProto(transactionIdOrReadTime, explainOptions); + + let streamActive: Deferred; + do { + streamActive = new Deferred(); + const methodName = 'runQuery'; + backendStream = await this._firestore.requestStream( + methodName, + /* bidirectional= */ false, + request, + tag + ); + backendStream.on('error', err => { + backendStream.unpipe(stream); + + // If a non-transactional query failed, attempt to restart. + // Transactional queries are retried via the transaction runner. + // Explain queries are not retried with a cursor. That would produce + // incorrect/partial profiling results. + if ( + !isExplain && + !transactionIdOrReadTime && + !this._isPermanentRpcError(err, 'runQuery') + ) { + logger( + 'QueryUtil._stream', + tag, + 'Query failed with retryable stream error:', + err + ); + + // Enqueue a "no-op" write into the stream and wait for it to be + // read by the downstream consumer. This ensures that all enqueued + // results in the stream are consumed, which will give us an accurate + // value for `lastReceivedDocument`. + stream.write(NOOP_MESSAGE, () => { + if (this._hasRetryTimedOut(methodName, startTime)) { + logger( + 'QueryUtil._stream', + tag, + 'Query failed with retryable stream error but the total retry timeout has exceeded.' + ); + stream.destroy(err); + streamActive.resolve(/* active= */ false); + } else if (lastReceivedDocument && retryWithCursor) { + logger( + 'Query._stream', + tag, + 'Query failed with retryable stream error and progress was made receiving ' + + 'documents, so the stream is being retried.' + ); + + // Restart the query but use the last document we received as + // the query cursor. Note that we do not use backoff here. The + // call to `requestStream()` will backoff should the restart + // fail before delivering any results. + if (this._queryOptions.requireConsistency) { + request = query + .startAfter(lastReceivedDocument) + .toProto(lastReceivedDocument.readTime); + } else { + request = query.startAfter(lastReceivedDocument).toProto(); + } + + // Set lastReceivedDocument to null before each retry attempt to ensure the retry makes progress + lastReceivedDocument = null; + + streamActive.resolve(/* active= */ true); + } else { + logger( + 'QueryUtil._stream', + tag, + `Query failed with retryable stream error however either retryWithCursor="${retryWithCursor}", or ` + + 'no progress was made receiving documents, so the stream is being closed.' + ); + stream.destroy(err); + streamActive.resolve(/* active= */ false); + } + }); + } else { + logger( + 'QueryUtil._stream', + tag, + 'Query failed with stream error:', + err + ); + stream.destroy(err); + streamActive.resolve(/* active= */ false); + } + }); + backendStream.on('end', () => { + streamActive.resolve(/* active= */ false); + }); + backendStream.resume(); + backendStream.pipe(stream); + } while (await streamActive.promise); + }) + .catch(e => stream.destroy(e)); + + return stream; + } +} + +/** + * A Query refers to a query which you can read or stream from. You can also + * construct refined Query objects by adding filters and ordering. + * + * @class Query + */ +export class Query< + AppModelType = firestore.DocumentData, + DbModelType extends firestore.DocumentData = firestore.DocumentData, +> implements firestore.Query +{ + /** + * @internal + * @private + **/ + readonly _serializer: Serializer; + /** + * @internal + * @private + **/ + protected readonly _allowUndefined: boolean; + /** + * @internal + * @private + **/ + readonly _queryUtil: QueryUtil< + AppModelType, + DbModelType, + Query + >; + + /** + * @internal + * @private + * + * @param _firestore The Firestore Database client. + * @param _queryOptions Options that define the query. + */ + constructor( + /** + * @internal + * @private + **/ + readonly _firestore: Firestore, + /** + * @internal + * @private **/ - protected readonly _queryOptions: QueryOptions + readonly _queryOptions: QueryOptions ) { this._serializer = new Serializer(_firestore); this._allowUndefined = !!this._firestore._settings.ignoreUndefinedProperties; + this._queryUtil = new QueryUtil< + AppModelType, + DbModelType, + Query + >(_firestore, _queryOptions, this._serializer); } /** @@ -1922,6 +2499,64 @@ export class Query< ); } + /** + * 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('embedding', [41, 42], {limit: 10, distanceMeasure: 'EUCLIDEAN'}); + * + * const querySnapshot = await aggregateQuery.get(); + * querySnapshot.forEach(...); + * ``` + * + * @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. + */ + findNearest( + vectorField: string | firestore.FieldPath, + queryVector: firestore.VectorValue | Array, + options: { + limit: number; + distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + } + ): VectorQuery { + validateFieldPath('vectorField', vectorField); + + if (options.limit <= 0) { + throw invalidArgumentMessage('options.limit', 'positive limit number'); + } + + if ( + (Array.isArray(queryVector) + ? queryVector.length + : queryVector.toArray().length) === 0 + ) { + throw invalidArgumentMessage( + 'queryVector', + 'vector size must be larger than 0' + ); + } + + return new VectorQuery( + this, + vectorField, + queryVector, + new VectorQueryOptions(options.limit, options.distanceMeasure) + ); + } + /** * Returns true if this `Query` is equal to the provided value. * @@ -2435,51 +3070,9 @@ export class Query< _get( transactionIdOrReadTime?: Uint8Array | Timestamp ): Promise> { - const docs: Array> = []; - - // Capture the error stack to preserve stack tracing across async calls. - const stack = Error().stack!; - - return new Promise((resolve, reject) => { - let readTime: Timestamp; - - this._stream(transactionIdOrReadTime) - .on('error', err => { - reject(wrapError(err, stack)); - }) - .on('data', result => { - readTime = result.readTime; - if (result.document) { - docs.push(result.document); - } - }) - .on('end', () => { - if (this._queryOptions.limitType === LimitType.Last) { - // The results for limitToLast queries need to be flipped since - // we reversed the ordering constraints before sending the query - // to the backend. - docs.reverse(); - } - - resolve( - new QuerySnapshot( - this, - readTime, - docs.length, - () => docs, - () => { - const changes: Array< - DocumentChange - > = []; - for (let i = 0; i < docs.length; ++i) { - changes.push(new DocumentChange('added', docs[i], -1, i)); - } - return changes; - } - ) - ); - }); - }); + return this._queryUtil._get(this, transactionIdOrReadTime) as Promise< + QuerySnapshot + >; } /** @@ -2504,24 +3097,7 @@ export class Query< * ``` */ stream(): NodeJS.ReadableStream { - if (this._queryOptions.limitType === LimitType.Last) { - throw new Error( - 'Query results for queries that include limitToLast() ' + - 'constraints cannot be streamed. Use Query.get() instead.' - ); - } - - const responseStream = this._stream(); - const transform = new Transform({ - objectMode: true, - transform(chunk, encoding, callback) { - callback(undefined, chunk.document); - }, - }); - - responseStream.pipe(transform); - responseStream.on('error', e => transform.destroy(e)); - return transform; + return this._queryUtil.stream(this); } /** @@ -2722,235 +3298,62 @@ export class Query< } if (this._queryOptions.hasFieldOrders()) { - structuredQuery.orderBy = this._queryOptions.fieldOrders.map(o => - o.toProto() - ); - } - - structuredQuery.startAt = this.toCursor(this._queryOptions.startAt); - structuredQuery.endAt = this.toCursor(this._queryOptions.endAt); - - if (this._queryOptions.limit) { - structuredQuery.limit = {value: this._queryOptions.limit}; - } - - structuredQuery.offset = this._queryOptions.offset; - structuredQuery.select = this._queryOptions.projection; - - return structuredQuery; - } - - /** - * @internal - * @private - * This method exists solely to enable unit tests to mock it. - */ - _isPermanentRpcError(err: GoogleError, methodName: string): boolean { - return isPermanentRpcError(err, methodName); - } - - /** - * @internal - * @private - */ - _hasRetryTimedOut(methodName: string, startTime: number): boolean { - const totalTimeout = getTotalTimeout(methodName); - if (totalTimeout === 0) { - return false; - } - - return Date.now() - startTime >= totalTimeout; - } - - /** - * Internal streaming method that accepts an optional transaction ID. - * - * @param transactionIdOrReadTime A transaction ID or the read time at which - * to execute the query. - * @param explainOptions Options to use for explaining the query (if any). - * @private - * @internal - * @returns A stream of document results. - */ - _stream( - transactionIdOrReadTime?: Uint8Array | Timestamp, - explainOptions?: firestore.ExplainOptions - ): NodeJS.ReadableStream { - const tag = requestTag(); - const startTime = Date.now(); - const isExplain = explainOptions !== undefined; - - let lastReceivedDocument: QueryDocumentSnapshot< - AppModelType, - DbModelType - > | null = null; - - let backendStream: Duplex; - const stream = new Transform({ - objectMode: true, - transform: (proto, enc, callback) => { - if (proto === NOOP_MESSAGE) { - callback(undefined); - return; - } - - const output: { - readTime?: Timestamp; - document?: QueryDocumentSnapshot; - explainMetrics?: ExplainMetrics; - } = {}; - - if (proto.readTime) { - output.readTime = Timestamp.fromProto(proto.readTime); - } - - if (proto.document) { - const document = this.firestore.snapshot_( - proto.document, - proto.readTime - ); - const finalDoc = new DocumentSnapshotBuilder< - AppModelType, - DbModelType - >(document.ref.withConverter(this._queryOptions.converter)); - // Recreate the QueryDocumentSnapshot with the DocumentReference - // containing the original converter. - finalDoc.fieldsProto = document._fieldsProto; - finalDoc.readTime = document.readTime; - finalDoc.createTime = document.createTime; - finalDoc.updateTime = document.updateTime; - lastReceivedDocument = finalDoc.build() as QueryDocumentSnapshot< - AppModelType, - DbModelType - >; - output.document = lastReceivedDocument; - } - - if (proto.explainMetrics) { - output.explainMetrics = ExplainMetrics._fromProto( - proto.explainMetrics, - this._serializer - ); - } - - callback(undefined, output); - - if (proto.done) { - logger('Query._stream', tag, 'Trigger Logical Termination.'); - backendStream.unpipe(stream); - backendStream.resume(); - backendStream.end(); - stream.end(); - } - }, - }); - - this.firestore - .initializeIfNeeded(tag) - .then(async () => { - // `toProto()` might throw an exception. We rely on the behavior of an - // async function to convert this exception into the rejected Promise we - // catch below. - let request = this.toProto(transactionIdOrReadTime, explainOptions); - - let streamActive: Deferred; - do { - streamActive = new Deferred(); - const methodName = 'runQuery'; - backendStream = await this._firestore.requestStream( - methodName, - /* bidirectional= */ false, - request, - tag - ); - backendStream.on('error', err => { - backendStream.unpipe(stream); - - // If a non-transactional query failed, attempt to restart. - // Transactional queries are retried via the transaction runner. - // Explain queries are not retried with a cursor. That would produce - // incorrect/partial profiling results. - if ( - !isExplain && - !transactionIdOrReadTime && - !this._isPermanentRpcError(err, 'runQuery') - ) { - logger( - 'Query._stream', - tag, - 'Query failed with retryable stream error:', - err - ); + structuredQuery.orderBy = this._queryOptions.fieldOrders.map(o => + o.toProto() + ); + } - // Enqueue a "no-op" write into the stream and wait for it to be - // read by the downstream consumer. This ensures that all enqueued - // results in the stream are consumed, which will give us an accurate - // value for `lastReceivedDocument`. - stream.write(NOOP_MESSAGE, () => { - if (this._hasRetryTimedOut(methodName, startTime)) { - logger( - 'Query._stream', - tag, - 'Query failed with retryable stream error but the total retry timeout has exceeded.' - ); - stream.destroy(err); - streamActive.resolve(/* active= */ false); - } else if (lastReceivedDocument) { - logger( - 'Query._stream', - tag, - 'Query failed with retryable stream error and progress was made receiving ' + - 'documents, so the stream is being retried.' - ); + structuredQuery.startAt = this.toCursor(this._queryOptions.startAt); + structuredQuery.endAt = this.toCursor(this._queryOptions.endAt); - // Restart the query but use the last document we received as - // the query cursor. Note that we do not use backoff here. The - // call to `requestStream()` will backoff should the restart - // fail before delivering any results. - if (this._queryOptions.requireConsistency) { - request = this.startAfter(lastReceivedDocument).toProto( - lastReceivedDocument.readTime - ); - } else { - request = this.startAfter(lastReceivedDocument).toProto(); - } + if (this._queryOptions.limit) { + structuredQuery.limit = {value: this._queryOptions.limit}; + } - // Set lastReceivedDocument to null before each retry attempt to ensure the retry makes progress - lastReceivedDocument = null; + structuredQuery.offset = this._queryOptions.offset; + structuredQuery.select = this._queryOptions.projection; - streamActive.resolve(/* active= */ true); - } else { - logger( - 'Query._stream', - tag, - 'Query failed with retryable stream error however no progress was made receiving ' + - 'documents, so the stream is being closed.' - ); - stream.destroy(err); - streamActive.resolve(/* active= */ false); - } - }); - } else { - logger( - 'Query._stream', - tag, - 'Query failed with stream error:', - err - ); - stream.destroy(err); - streamActive.resolve(/* active= */ false); - } - }); - backendStream.on('end', () => { - streamActive.resolve(/* active= */ false); - }); - backendStream.resume(); - backendStream.pipe(stream); - } while (await streamActive.promise); - }) - .catch(e => stream.destroy(e)); + return structuredQuery; + } - return stream; + /** + * @internal + * @private + * This method exists solely to maintain backward compatability. + */ + _isPermanentRpcError(err: GoogleError, methodName: string): boolean { + return this._queryUtil._isPermanentRpcError(err, methodName); + } + + /** + * @internal + * @private + * This method exists solely to maintain backward compatability. + */ + _hasRetryTimedOut(methodName: string, startTime: number): boolean { + return this._queryUtil._hasRetryTimedOut(methodName, startTime); + } + + /** + * Internal streaming method that accepts an optional transaction ID. + * + * @param transactionIdOrReadTime A transaction ID or the read time at which + * to execute the query. + * @param explainOptions Options to use for explaining the query (if any). + * @private + * @internal + * @returns A stream of document results. + */ + _stream( + transactionIdOrReadTime?: Uint8Array | Timestamp, + explainOptions?: firestore.ExplainOptions + ): NodeJS.ReadableStream { + return this._queryUtil._stream( + this, + transactionIdOrReadTime, + true, + explainOptions + ); } /** @@ -3117,6 +3520,27 @@ export class Query< this._queryOptions.withConverter(converter ?? defaultConverter()) ); } + + /** + * Construct the resulting snapshot for this query with given documents. + * + * @private + * @internal + */ + _createSnapshot( + readTime: Timestamp, + size: number, + docs: () => Array>, + changes: () => Array> + ): QuerySnapshot { + return new QuerySnapshot( + this, + readTime, + size, + docs, + changes + ); + } } /** @@ -3896,6 +4320,216 @@ export class AggregateQuerySnapshot< } } +class VectorQueryOptions { + constructor( + readonly limit: number, + readonly distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT' + ) {} + + isEqual(other: VectorQueryOptions): boolean { + if (this === other) { + return true; + } + if (!(other instanceof VectorQueryOptions)) { + return false; + } + + return ( + this.limit === other.limit && + this.distanceMeasure === other.distanceMeasure + ); + } +} + +/** + * A query that finds the documents whose vector fields are closest to a certain query vector. + * Create an instance of `VectorQuery` with {@link Query.findNearest}. + */ +export class VectorQuery< + AppModelType = firestore.DocumentData, + DbModelType extends firestore.DocumentData = firestore.DocumentData, +> implements firestore.VectorQuery +{ + /** + * @internal + * @private + **/ + readonly _queryUtil: QueryUtil< + AppModelType, + DbModelType, + VectorQuery + >; + + /** + * @private + * @internal + */ + constructor( + private readonly _query: Query, + private readonly vectorField: string | firestore.FieldPath, + private readonly queryVector: firestore.VectorValue | Array, + private readonly options: VectorQueryOptions + ) { + this._queryUtil = new QueryUtil< + AppModelType, + DbModelType, + VectorQuery + >(_query._firestore, _query._queryOptions, _query._serializer); + } + + /** The query whose results participants in the vector search. Filtering + * performed by the query will apply before the vector search. + **/ + get query(): Query { + return this._query; + } + + /** + * @private + * @internal + */ + private get _rawVectorField(): string { + return typeof this.vectorField === 'string' + ? this.vectorField + : this.vectorField.toString(); + } + + /** + * @private + * @internal + */ + private get _rawQueryVector(): Array { + return Array.isArray(this.queryVector) + ? this.queryVector + : this.queryVector.toArray(); + } + + /** + * Executes this vector search query. + * + * @returns A promise that will be resolved with the results of the query. + */ + get(): Promise> { + return this._queryUtil._get( + this, + /*transactionId*/ undefined, + // VectorQuery cannot be retried with cursors as they do not support cursors yet. + /*retryWithCursor*/ false + ) as Promise>; + } + + /** + * Internal streaming method that accepts an optional transaction ID. + * + * @param transactionId - A transaction ID. + * @private + * @internal + * @returns A stream of document results. + */ + _stream(transactionId?: Uint8Array): NodeJS.ReadableStream { + return this._queryUtil._stream( + this, + transactionId, + /*retryWithCursor*/ false + ); + } + + /** + * Internal method for serializing a query to its RunAggregationQuery proto + * representation with an optional transaction id. + * + * @private + * @internal + * @returns Serialized JSON for the query. + */ + toProto( + transactionIdOrReadTime?: Uint8Array | Timestamp + ): api.IRunQueryRequest { + const queryProto = this._query.toProto(transactionIdOrReadTime); + + const queryVector = Array.isArray(this.queryVector) + ? new VectorValue(this.queryVector) + : (this.queryVector as VectorValue); + + queryProto.structuredQuery!.findNearest = { + limit: {value: this.options.limit}, + distanceMeasure: this.options.distanceMeasure, + vectorField: { + fieldPath: FieldPath.fromArgument(this.vectorField).formattedName, + }, + queryVector: queryVector._toProto(this._query._serializer), + }; + return queryProto; + } + + /** + * Construct the resulting vector snapshot for this query with given documents. + * + * @private + * @internal + */ + _createSnapshot( + readTime: Timestamp, + size: number, + docs: () => Array>, + changes: () => Array> + ): VectorQuerySnapshot { + return new VectorQuerySnapshot( + this, + readTime, + size, + docs, + changes + ); + } + + /** + * Construct a new vector query whose result will start after To support stream(). + * This now throws an exception because cursors are not supported from the backend for vector queries yet. + * + * @private + * @internal + * @returns Serialized JSON for the query. + */ + startAfter( + // eslint-disable-next-line @typescript-eslint/no-unused-vars + ...fieldValuesOrDocumentSnapshot: Array + ): VectorQuery { + throw new Error( + 'Unimplemented: Vector query does not support cursors yet.' + ); + } + + /** + * Compares this object with the given object for equality. + * + * This object is considered "equal" to the other object if and only if + * `other` performs the same vector distance search as this `VectorQuery` and + * the underlying Query of `other` compares equal to that of this object + * using `Query.isEqual()`. + * + * @param other - The object to compare to this object for equality. + * @returns `true` if this object is "equal" to the given object, as + * defined above, or `false` otherwise. + */ + isEqual(other: firestore.VectorQuery): boolean { + if (this === other) { + return true; + } + if (!(other instanceof VectorQuery)) { + return false; + } + if (!this.query.isEqual(other.query)) { + return false; + } + return ( + this._rawVectorField === other._rawVectorField && + isPrimitiveArrayEqual(this._rawQueryVector, other._rawQueryVector) && + this.options.isEqual(other.options) + ); + } +} + /** * Validates the input string as a field order direction. * @@ -4004,32 +4638,6 @@ function validateQueryValue( }); } -/** - * Verifies equality for an array of objects using the `isEqual` interface. - * - * @private - * @internal - * @param left Array of objects supporting `isEqual`. - * @param right Array of objects supporting `isEqual`. - * @return True if arrays are equal. - */ -function isArrayEqual boolean}>( - left: T[], - right: T[] -): boolean { - if (left.length !== right.length) { - return false; - } - - for (let i = 0; i < left.length; ++i) { - if (!left[i].isEqual(right[i])) { - return false; - } - } - - return true; -} - /** * Returns the first non-undefined value or `undefined` if no such value exists. * @private diff --git a/dev/src/serializer.ts b/dev/src/serializer.ts index 415c2b82f..ace53ce4c 100644 --- a/dev/src/serializer.ts +++ b/dev/src/serializer.ts @@ -18,8 +18,8 @@ import {DocumentData} from '@google-cloud/firestore'; import * as proto from '../protos/firestore_v1_proto_api'; +import {DeleteTransform, FieldTransform, VectorValue} from './field-value'; import {detectGoogleProtobufValueType, detectValueType} from './convert'; -import {DeleteTransform, FieldTransform} from './field-value'; import {GeoPoint} from './geo-point'; import {DocumentReference, Firestore} from './index'; import {FieldPath, QualifiedResourcePath} from './path'; @@ -29,6 +29,11 @@ import {isEmpty, isObject, isPlainObject} from './util'; import {customObjectMessage, invalidArgumentMessage} from './validate'; import api = proto.google.firestore.v1; +import { + RESERVED_MAP_KEY, + RESERVED_MAP_KEY_VECTOR_VALUE, + VECTOR_MAP_VECTORS_KEY, +} from './map-type'; /** * The maximum depth of a Firestore object. @@ -168,6 +173,10 @@ export class Serializer { }; } + if (val instanceof VectorValue) { + return val._toProto(this); + } + if (isObject(val)) { const toProto = val['toProto']; if (typeof toProto === 'function') { @@ -217,6 +226,31 @@ export class Serializer { throw new Error(`Cannot encode value: ${val}`); } + /** + * @private + */ + encodeVector(rawVector: number[]): api.IValue { + // A Firestore Vector is a map with reserved key/value pairs. + return { + mapValue: { + fields: { + [RESERVED_MAP_KEY]: { + stringValue: RESERVED_MAP_KEY_VECTOR_VALUE, + }, + [VECTOR_MAP_VECTORS_KEY]: { + arrayValue: { + values: rawVector.map(value => { + return { + doubleValue: value, + }; + }), + }, + }, + }, + }, + }; + } + /** * Decodes a single Firestore 'Value' Protobuf. * @@ -263,15 +297,20 @@ export class Serializer { return null; } case 'mapValue': { - const obj: DocumentData = {}; const fields = proto.mapValue!.fields; if (fields) { + const obj: DocumentData = {}; for (const prop of Object.keys(fields)) { obj[prop] = this.decodeValue(fields[prop]); } + return obj; + } else { + return {}; } - - return obj; + } + case 'vectorValue': { + const fields = proto.mapValue!.fields!; + return VectorValue._fromProto(fields[VECTOR_MAP_VECTORS_KEY]); } case 'geoPointValue': { return GeoPoint.fromProto(proto.geoPointValue!); @@ -444,6 +483,8 @@ export function validateUserInput( 'If you want to ignore undefined values, enable `ignoreUndefinedProperties`.' ); } + } else if (value instanceof VectorValue) { + // OK } else if (value instanceof DeleteTransform) { if (inArray) { throw new Error( diff --git a/dev/src/util.ts b/dev/src/util.ts index 6781032b4..c4b4754c1 100644 --- a/dev/src/util.ts +++ b/dev/src/util.ts @@ -289,3 +289,55 @@ export function mapToArray( } return result; } + +/** + * Verifies equality for an array of objects using the `isEqual` interface. + * + * @private + * @internal + * @param left Array of objects supporting `isEqual`. + * @param right Array of objects supporting `isEqual`. + * @return True if arrays are equal. + */ +export function isArrayEqual boolean}>( + left: T[], + right: T[] +): boolean { + if (left.length !== right.length) { + return false; + } + + for (let i = 0; i < left.length; ++i) { + if (!left[i].isEqual(right[i])) { + return false; + } + } + + return true; +} + +/** + * Verifies equality for an array of primitives. + * + * @private + * @internal + * @param left Array of primitives. + * @param right Array of primitives. + * @return True if arrays are equal. + */ +export function isPrimitiveArrayEqual( + left: T[], + right: T[] +): boolean { + if (left.length !== right.length) { + return false; + } + + for (let i = 0; i < left.length; ++i) { + if (left[i] !== right[i]) { + return false; + } + } + + return true; +} diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts index 70934b62f..c21cfdbc6 100644 --- a/dev/system-test/firestore.ts +++ b/dev/system-test/firestore.ts @@ -19,6 +19,7 @@ import { QuerySnapshot, SetOptions, Settings, + VectorValue, WithFieldValue, } from '@google-cloud/firestore'; @@ -58,6 +59,7 @@ import {QueryPartition} from '../src/query-partition'; import {CollectionGroup} from '../src/collection-group'; import IBundleElement = firestore.IBundleElement; import {Filter} from '../src/filter'; +import {IndexTestHelper} from './index_test_helper'; use(chaiAsPromised); @@ -1301,6 +1303,31 @@ describe('DocumentReference class', () => { return promise; }); + it('can write and read vector embeddings', async () => { + const ref = randomCol.doc(); + await ref.create({ + vector0: FieldValue.vector([0.0]), + vector1: FieldValue.vector([1, 2, 3.99]), + }); + await ref.set({ + vector0: FieldValue.vector([0.0]), + vector1: FieldValue.vector([1, 2, 3.99]), + vector2: FieldValue.vector([0, 0, 0]), + }); + await ref.update({ + vector3: FieldValue.vector([-1, -200, -999]), + }); + + const snap1 = await ref.get(); + expect(snap1.get('vector0').isEqual(FieldValue.vector([0.0]))).to.be.true; + expect(snap1.get('vector1').isEqual(FieldValue.vector([1, 2, 3.99]))).to.be + .true; + expect(snap1.get('vector2').isEqual(FieldValue.vector([0, 0, 0]))).to.be + .true; + expect(snap1.get('vector3').isEqual(FieldValue.vector([-1, -200, -999]))).to + .be.true; + }); + describe('watch', () => { const currentDeferred = new DeferredPromise(); @@ -1594,6 +1621,89 @@ describe('DocumentReference class', () => { const result2 = await ref2.get(); expect(result2.data()).to.deep.equal([1, 2, 3]); }); + + it('can listen to documents with vectors', async () => { + const ref = randomCol.doc(); + const initialDeferred = new Deferred(); + const createDeferred = new Deferred(); + const setDeferred = new Deferred(); + const updateDeferred = new Deferred(); + const deleteDeferred = new Deferred(); + + const expected = [ + initialDeferred, + createDeferred, + setDeferred, + updateDeferred, + deleteDeferred, + ]; + let idx = 0; + let document: DocumentSnapshot | null = null; + + const unlisten = randomCol + .where('purpose', '==', 'vector tests') + .onSnapshot(snap => { + expected[idx].resolve(); + idx += 1; + if (snap.docs.length > 0) { + document = snap.docs[0]; + } else { + document = null; + } + }); + + await initialDeferred.promise; + expect(document).to.be.null; + + await ref.create({ + purpose: 'vector tests', + vector0: FieldValue.vector([0.0]), + vector1: FieldValue.vector([1, 2, 3.99]), + }); + + await createDeferred.promise; + expect(document).to.be.not.null; + expect(document!.get('vector0').isEqual(FieldValue.vector([0.0]))).to.be + .true; + expect(document!.get('vector1').isEqual(FieldValue.vector([1, 2, 3.99]))).to + .be.true; + + await ref.set({ + purpose: 'vector tests', + vector0: FieldValue.vector([0.0]), + vector1: FieldValue.vector([1, 2, 3.99]), + vector2: FieldValue.vector([0, 0, 0]), + }); + await setDeferred.promise; + expect(document).to.be.not.null; + expect(document!.get('vector0').isEqual(FieldValue.vector([0.0]))).to.be + .true; + expect(document!.get('vector1').isEqual(FieldValue.vector([1, 2, 3.99]))).to + .be.true; + expect(document!.get('vector2').isEqual(FieldValue.vector([0, 0, 0]))).to.be + .true; + + await ref.update({ + vector3: FieldValue.vector([-1, -200, -999]), + }); + await updateDeferred.promise; + expect(document).to.be.not.null; + expect(document!.get('vector0').isEqual(FieldValue.vector([0.0]))).to.be + .true; + expect(document!.get('vector1').isEqual(FieldValue.vector([1, 2, 3.99]))).to + .be.true; + expect(document!.get('vector2').isEqual(FieldValue.vector([0, 0, 0]))).to.be + .true; + expect( + document!.get('vector3').isEqual(FieldValue.vector([-1, -200, -999])) + ).to.be.true; + + await ref.delete(); + await deleteDeferred.promise; + expect(document).to.be.null; + + unlisten(); + }); }); describe('runs query on a large collection', () => { @@ -1785,6 +1895,328 @@ describe('Query class', () => { }); }); + 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 vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest('embedding', [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 vectorQuery = indexTestHelper + .query(collectionReference) + .where('foo', '==', 'bar') + .findNearest('embedding', [10, 10], { + limit: 3, + distanceMeasure: 'COSINE', + }); + + const res = await vectorQuery.get(); + + expect(res.size).to.equal(3); + + 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; + }); + + it('supports findNearest by DOT_PRODUCT 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 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(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('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('embedding', [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('embedding', [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[2].get('embedding').isEqual(FieldValue.vector([100, 100]))) + .to.be.true; + }); + + 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('embedding', [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; + }); + + it('supports findNearest on non-existent field', 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 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(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('nested.embedding', [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('embedding', [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('embedding', 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); + }); + it('supports !=', async () => { await addDocs( {zip: NaN}, @@ -2946,6 +3378,72 @@ describe('Query class', () => { unsubscribe(); }); + + it('SDK orders vector field same way as backend', async () => { + // We validate that the SDK orders the vector field the same way as the backend + // by comparing the sort order of vector fields from a Query.get() and + // Query.onSnapshot(). Query.onSnapshot() will return sort order of the SDK, + // and Query.get() will return sort order of the backend. + + // Test data in the order that we expect the backend to sort it. + const docsInOrder = [ + {embedding: [1, 2, 3, 4, 5, 6]}, + {embedding: [100]}, + {embedding: FieldValue.vector([Number.NEGATIVE_INFINITY])}, + {embedding: FieldValue.vector([-100])}, + {embedding: FieldValue.vector([100])}, + {embedding: FieldValue.vector([Number.POSITIVE_INFINITY])}, + {embedding: FieldValue.vector([1, 2])}, + {embedding: FieldValue.vector([2, 2])}, + {embedding: FieldValue.vector([1, 2, 3])}, + {embedding: FieldValue.vector([1, 2, 3, 4])}, + {embedding: FieldValue.vector([1, 2, 3, 4, 5])}, + {embedding: FieldValue.vector([1, 2, 100, 4, 4])}, + {embedding: FieldValue.vector([100, 2, 3, 4, 5])}, + {embedding: {HELLO: 'WORLD'}}, + {embedding: {hello: 'world'}}, + ]; + + const expectedSnapshots = []; + const expectedChanges = []; + + for (let i = 0; i < docsInOrder.length; i++) { + const dr = await randomCol.add(docsInOrder[i]); + expectedSnapshots.push(snapshot(dr.id, docsInOrder[i])); + expectedChanges.push(added(dr.id, docsInOrder[i])); + } + + const orderedQuery = randomCol.orderBy('embedding'); + + const unsubscribe = orderedQuery.onSnapshot( + snapshot => { + currentDeferred.resolve(snapshot); + }, + err => { + currentDeferred.reject!(err); + } + ); + + const watchSnapshot = await waitForSnapshot(); + unsubscribe(); + + const getSnapshot = await orderedQuery.get(); + + // Compare the snapshot (including sort order) of a snapshot + // from Query.onSnapshot() to an actual snapshot from Query.get() + snapshotsEqual(watchSnapshot, { + docs: getSnapshot.docs, + docChanges: getSnapshot.docChanges(), + }); + + // Compare the snapshot (including sort order) of a snapshot + // from Query.onSnapshot() to the expected sort order from + // the backend. + snapshotsEqual(watchSnapshot, { + docs: expectedSnapshots, + docChanges: expectedChanges, + }); + }); }); (process.env.FIRESTORE_EMULATOR_HOST === undefined diff --git a/dev/system-test/index_test_helper.ts b/dev/system-test/index_test_helper.ts new file mode 100644 index 000000000..0e3ef58c5 --- /dev/null +++ b/dev/system-test/index_test_helper.ts @@ -0,0 +1,205 @@ +/*! + * Copyright 2024 Google LLC. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + CollectionReference, + DocumentReference, + DocumentSnapshot, + FieldPath, + Filter, + Firestore, + Query, + Timestamp, +} from '../src'; +import {autoId} from '../src/util'; + +import { + DocumentData, + QuerySnapshot, + WithFieldValue, + UpdateData, +} from '@google-cloud/firestore'; +export const INDEX_TEST_COLLECTION = 'index-test-collection'; + +/** + * This helper class is designed to facilitate integration testing of Firestore queries that + * require manually created indexes within a controlled testing environment. + * + *

Key Features: + * + *

    + *
  • Runs tests against the dedicated test collection with predefined indexes. + *
  • Automatically associates a test ID with documents for data isolation. + *
  • Utilizes TTL policy for automatic test data cleanup. + *
  • Constructs Firestore queries with test ID filters. + *
+ */ +export class IndexTestHelper { + private readonly testId: string; + private readonly TEST_ID_FIELD: string = 'testId'; + private readonly TTL_FIELD: string = 'expireAt'; + + // Creates a new instance of the CompositeIndexTestHelper class, with a unique test + // identifier for data isolation. + constructor(public readonly db: Firestore) { + this.testId = 'test-id-' + autoId(); + } + + // Runs a test with specified documents in the INDEX_TEST_COLLECTION. + async setTestDocs(docs: { + [key: string]: DocumentData; + }): Promise { + const testDocs = this.prepareTestDocuments(docs); + const collectionRef = this.db.collection(INDEX_TEST_COLLECTION); + for (const id in testDocs) { + const ref = collectionRef.doc(id); + await ref.set(testDocs[id]); + } + return collectionRef; + } + + // Runs a test with specified documents in the INDEX_TEST_COLLECTION. + async createTestDocs(docs: DocumentData[]): Promise { + // convert docsArray without IDs to a map with IDs + const docsMap = docs.reduce<{[key: string]: DocumentData}>( + (result, doc) => { + result[autoId()] = doc; + return result; + }, + {} + ); + return this.setTestDocs(docsMap); + } + + // Runs a test on INDEX_TEST_COLLECTION. + async withTestCollection(): Promise { + const collectionRef = this.db.collection(INDEX_TEST_COLLECTION); + return collectionRef; + } + + // Hash the document key with testId. + private toHashedId(docId: string): string { + return docId + '-' + this.testId; + } + + private toHashedIds(docs: string[]): string[] { + return docs.map(docId => this.toHashedId(docId)); + } + + // Adds test-specific fields to a document, including the testId and expiration date. + addTestSpecificFieldsToDoc(doc: DocumentData): DocumentData { + return { + ...doc, + [this.TEST_ID_FIELD]: this.testId, + [this.TTL_FIELD]: new Timestamp( // Expire test data after 24 hours + Timestamp.now().seconds + 24 * 60 * 60, + Timestamp.now().nanoseconds + ), + }; + } + + // Remove test-specific fields from a document, including the testId and expiration date. + private removeTestSpecificFieldsFromDoc(doc: DocumentData): void { + doc._document?.data?.delete(new FieldPath(this.TEST_ID_FIELD)); + doc._document?.data?.delete(new FieldPath(this.TTL_FIELD)); + } + + // Helper method to hash document keys and add test-specific fields for the provided documents. + private prepareTestDocuments(docs: {[key: string]: DocumentData}): { + [key: string]: DocumentData; + } { + const result: {[key: string]: DocumentData} = {}; + for (const key in docs) { + // eslint-disable-next-line no-prototype-builtins + if (docs.hasOwnProperty(key)) { + result[this.toHashedId(key)] = this.addTestSpecificFieldsToDoc( + docs[key] + ); + } + } + return result; + } + + // Adds a filter on test id for a query. + query(query_: Query, ...filters: Filter[]): Query { + return filters.reduce>( + (query, filter) => { + return query.where(filter); + }, + query_.where(this.TEST_ID_FIELD, '==', this.testId) + ); + } + + // Get document reference from a document key. + getDocRef( + coll: CollectionReference, + docId: string + ): DocumentReference { + if (!docId.includes('test-id-')) { + docId = this.toHashedId(docId); + } + return coll.doc(docId); + } + + // Adds a document to a Firestore collection with test-specific fields. + addDoc( + reference: CollectionReference, + data: object + ): Promise> { + const processedData = this.addTestSpecificFieldsToDoc( + data + ) as WithFieldValue; + return reference.add(processedData); + } + + // Sets a document in Firestore with test-specific fields. + async setDoc( + reference: DocumentReference, + data: object + ): Promise { + const processedData = this.addTestSpecificFieldsToDoc( + data + ) as WithFieldValue; + await reference.set(processedData); + } + + async updateDoc( + reference: DocumentReference, + data: UpdateData + ): Promise { + await reference.update(data); + } + + async deleteDoc(reference: DocumentReference): Promise { + await reference.delete(); + } + + // Retrieves a single document from Firestore with test-specific fields removed. + async getDoc(docRef: DocumentReference): Promise> { + const docSnapshot = await docRef.get(); + this.removeTestSpecificFieldsFromDoc(docSnapshot); + return docSnapshot; + } + + // Retrieves multiple documents from Firestore with test-specific fields removed. + async getDocs(query_: Query): Promise> { + const querySnapshot = await this.query(query_).get(); + querySnapshot.forEach(doc => { + this.removeTestSpecificFieldsFromDoc(doc); + }); + return querySnapshot; + } +} diff --git a/dev/test/document.ts b/dev/test/document.ts index bc140aeef..bb33d8fd9 100644 --- a/dev/test/document.ts +++ b/dev/test/document.ts @@ -471,6 +471,43 @@ describe('serialize document', () => { return ref.set({ref}); }); }); + + it('is able to translate FirestoreVector to internal representation with set', () => { + const overrides: ApiOverride = { + commit: request => { + requestEquals( + request, + set({ + document: document('documentId', 'embedding1', { + mapValue: { + fields: { + __type__: { + stringValue: '__vector__', + }, + value: { + arrayValue: { + values: [ + {doubleValue: 0}, + {doubleValue: 1}, + {doubleValue: 2}, + ], + }, + }, + }, + }, + }), + }) + ); + return response(writeResult(1)); + }, + }; + + return createInstance(overrides).then(firestore => { + return firestore.doc('collectionId/documentId').set({ + embedding1: FieldValue.vector([0, 1, 2]), + }); + }); + }); }); describe('deserialize document', () => { @@ -599,6 +636,46 @@ describe('deserialize document', () => { }); }); + it('deserializes FirestoreVector', () => { + const overrides: ApiOverride = { + batchGetDocuments: () => { + return stream( + found( + document('documentId', 'embedding', { + mapValue: { + fields: { + __type__: { + stringValue: '__vector__', + }, + value: { + arrayValue: { + values: [ + {doubleValue: -41.0}, + {doubleValue: 0}, + {doubleValue: 42}, + ], + }, + }, + }, + }, + }) + ) + ); + }, + }; + + return createInstance(overrides).then(firestore => { + return firestore + .doc('collectionId/documentId') + .get() + .then(res => { + expect(res.get('embedding')).to.deep.equal( + FieldValue.vector([-41.0, 0, 42]) + ); + }); + }); + }); + it("doesn't deserialize unsupported types", () => { const overrides: ApiOverride = { batchGetDocuments: () => { diff --git a/dev/test/index.ts b/dev/test/index.ts index 5e1713376..bd9b2cae6 100644 --- a/dev/test/index.ts +++ b/dev/test/index.ts @@ -21,7 +21,7 @@ import {GoogleError, GrpcClient, Status} from 'google-gax'; import {google} from '../protos/firestore_v1_proto_api'; import * as Firestore from '../src'; -import {DocumentSnapshot, FieldPath} from '../src'; +import {DocumentSnapshot, FieldPath, FieldValue} from '../src'; import {setTimeoutHandler} from '../src/backoff'; import {QualifiedResourcePath} from '../src/path'; import { @@ -125,6 +125,31 @@ const allSupportedTypesProtobufJs = document( }, }, }, + 'vectorValue', + { + mapValue: { + fields: { + __type__: { + stringValue: '__vector__', + }, + value: { + arrayValue: { + values: [ + { + doubleValue: 0.1, + }, + { + doubleValue: 0.2, + }, + { + doubleValue: 0.3, + }, + ], + }, + }, + }, + }, + }, 'emptyObject', { mapValue: {}, @@ -211,6 +236,30 @@ const allSupportedTypesJson = { emptyObject: { mapValue: {}, }, + vectorValue: { + mapValue: { + fields: { + __type__: { + stringValue: '__vector__', + }, + value: { + arrayValue: { + values: [ + { + doubleValue: 0.1, + }, + { + doubleValue: 0.2, + }, + { + doubleValue: 0.3, + }, + ], + }, + }, + }, + }, + }, pathValue: { referenceValue: `${DATABASE_ROOT}/documents/collection/document`, }, @@ -244,6 +293,7 @@ const allSupportedTypesInput = { negativeInfinityValue: -Infinity, objectValue: {foo: 'bar'}, emptyObject: {}, + vectorValue: FieldValue.vector([0.1, 0.2, 0.3]), dateValue: new Date('Mar 18, 1985 08:20:00.123 GMT+0100 (CET)'), timestampValue: Firestore.Timestamp.fromDate( new Date('Mar 18, 1985 08:20:00.123 GMT+0100 (CET)') @@ -272,6 +322,7 @@ const allSupportedTypesOutput: {[field: string]: unknown} = { negativeInfinityValue: -Infinity, objectValue: {foo: 'bar'}, emptyObject: {}, + vectorValue: FieldValue.vector([0.1, 0.2, 0.3]), dateValue: Firestore.Timestamp.fromDate( new Date('Mar 18, 1985 08:20:00.123 GMT+0100 (CET)') ), diff --git a/dev/test/query.ts b/dev/test/query.ts index 97dce3453..aaf7f5829 100644 --- a/dev/test/query.ts +++ b/dev/test/query.ts @@ -725,7 +725,7 @@ describe('query interface', () => { let attempts = 0; const query = firestore.collection('collectionId'); - query._stream = () => { + query._queryUtil._stream = () => { ++attempts; throw new Error('Expected error'); }; @@ -896,7 +896,7 @@ describe('query interface', () => { return createInstance(overrides).then(firestoreInstance => { firestore = firestoreInstance; const query = firestore.collection('collectionId'); - query._hasRetryTimedOut = () => false; + query._queryUtil._hasRetryTimedOut = () => false; return query .get() .then(() => { @@ -956,7 +956,7 @@ describe('query interface', () => { firestore = firestoreInstance; const query = firestore.collection('collectionId'); // Fake our timeout check to fail after 10 retry attempts - query._hasRetryTimedOut = (methodName, startTime) => { + query._queryUtil._hasRetryTimedOut = (methodName, startTime) => { expect(methodName).to.equal('runQuery'); expect(startTime).to.be.lessThanOrEqual(Date.now()); return attempts >= 10; @@ -3664,7 +3664,7 @@ describe('query resumption', () => { // the async iterator into an array. firestore = await createInstance(overrides); const query = firestore.collection('collectionId'); - query._isPermanentRpcError = mockIsPermanentRpcError; + query._queryUtil._isPermanentRpcError = mockIsPermanentRpcError; const iterator = query .stream() [Symbol.asyncIterator]() as AsyncIterator; diff --git a/dev/test/vector-query.ts b/dev/test/vector-query.ts new file mode 100644 index 000000000..a4d08e027 --- /dev/null +++ b/dev/test/vector-query.ts @@ -0,0 +1,335 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {afterEach, beforeEach, it} from 'mocha'; +import {fieldFiltersQuery, queryEquals, result} from './query'; +import { + ApiOverride, + createInstance, + stream, + streamWithoutEnd, + verifyInstance, +} from './util/helpers'; +import { + DocumentSnapshot, + FieldValue, + Firestore, + Query, + Timestamp, +} from '../src'; +import {expect, use} from 'chai'; +import {google} from '../protos/firestore_v1_proto_api'; +import api = google.firestore.v1; +import * as chaiAsPromised from 'chai-as-promised'; +import {setTimeoutHandler} from '../src/backoff'; +use(chaiAsPromised); + +export function findNearestQuery( + fieldPath: string, + queryVector: Array, + limit: number, + measure: api.StructuredQuery.FindNearest.DistanceMeasure +): api.IStructuredQuery { + return { + findNearest: { + vectorField: {fieldPath}, + queryVector: { + mapValue: { + fields: { + __type__: {stringValue: '__vector__'}, + value: { + arrayValue: { + values: queryVector.map(n => { + return {doubleValue: n}; + }), + }, + }, + }, + }, + }, + limit: {value: limit}, + distanceMeasure: measure, + }, + }; +} + +describe('Vector(findNearest) query interface', () => { + let firestore: Firestore; + + beforeEach(() => { + setTimeoutHandler(setImmediate); + return createInstance().then(firestoreInstance => { + firestore = firestoreInstance; + }); + }); + + afterEach(async () => { + await verifyInstance(firestore); + setTimeoutHandler(setTimeout); + }); + + it('has isEqual() method', () => { + 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( + queryA.findNearest('embedding', [40, 41, 42], { + distanceMeasure: 'COSINE', + limit: 10, + }) + ) + ).to.be.true; + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + queryB.findNearest('embedding', [40, 41, 42], { + distanceMeasure: 'COSINE', + limit: 10, + }) + ) + ).to.be.true; + + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + firestore + .collection('collectionId') + .findNearest('embedding', [40, 41, 42], { + distanceMeasure: 'COSINE', + limit: 10, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + queryB.findNearest('embedding', [40, 42], { + distanceMeasure: 'COSINE', + limit: 10, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + queryB.findNearest('embedding', [40, 42], { + distanceMeasure: 'COSINE', + limit: 1000, + }) + ) + ).to.be.false; + expect( + queryA + .findNearest('embedding', [40, 41, 42], { + limit: 10, + distanceMeasure: 'COSINE', + }) + .isEqual( + queryB.findNearest('embedding', [40, 42], { + distanceMeasure: 'EUCLIDEAN', + limit: 1000, + }) + ) + ).to.be.false; + }); + + it('generates proto', async () => { + const overrides: ApiOverride = { + runQuery: request => { + queryEquals( + request, + fieldFiltersQuery('foo', 'EQUAL', 'bar'), + findNearestQuery('embedding', [3, 4, 5], 100, 'COSINE') + ); + return stream(); + }, + }; + + return createInstance(overrides).then(firestoreInstance => { + firestore = firestoreInstance; + const query: Query = firestore.collection('collectionId'); + const vectorQuery = query + .where('foo', '==', 'bar') + .findNearest('embedding', FieldValue.vector([3, 4, 5]), { + limit: 100, + distanceMeasure: 'COSINE', + }); + return vectorQuery.get(); + }); + }); + + it('validates inputs', async () => { + const query: Query = firestore.collection('collectionId'); + expect(() => { + query.findNearest('embedding', [], { + limit: 10, + distanceMeasure: 'EUCLIDEAN', + }); + }).to.throw('not a valid vector size'); + expect(() => { + query.findNearest('embedding', [10, 1000], { + limit: 0, + distanceMeasure: 'EUCLIDEAN', + }); + }).to.throw('not a valid positive limit number'); + }); + + const distanceMeasure: ('EUCLIDEAN' | 'DOT_PRODUCT' | 'COSINE')[] = [ + 'EUCLIDEAN', + 'DOT_PRODUCT', + 'COSINE', + ]; + distanceMeasure.forEach(distanceMeasure => { + it(`returns results when distanceMeasure is ${distanceMeasure}`, async () => { + const overrides: ApiOverride = { + runQuery: request => { + queryEquals( + request, + findNearestQuery('embedding', [1], 2, distanceMeasure) + ); + return stream(result('first'), result('second')); + }, + }; + + return createInstance(overrides).then(firestoreInstance => { + firestore = firestoreInstance; + const query = firestore + .collection('collectionId') + .findNearest('embedding', [1], { + limit: 2, + distanceMeasure: distanceMeasure, + }); + return query.get().then(results => { + expect(results.size).to.equal(2); + expect(results.empty).to.be.false; + expect(results.readTime.isEqual(new Timestamp(5, 6))).to.be.true; + expect(results.docs[0].id).to.equal('first'); + expect(results.docs[1].id).to.equal('second'); + expect(results.docChanges()).to.have.length(2); + + let count = 0; + + results.forEach(doc => { + expect(doc instanceof DocumentSnapshot).to.be.true; + expect(doc.createTime.isEqual(new Timestamp(1, 2))).to.be.true; + expect(doc.updateTime.isEqual(new Timestamp(3, 4))).to.be.true; + expect(doc.readTime.isEqual(new Timestamp(5, 6))).to.be.true; + ++count; + }); + + expect(2).to.equal(count); + }); + }); + }); + }); + + it('successful return without ending the stream on get()', async () => { + const overrides: ApiOverride = { + runQuery: request => { + queryEquals(request, findNearestQuery('vector', [1], 10, 'COSINE')); + return streamWithoutEnd(result('first'), result('second', true)); + }, + }; + + let counter = 0; + return createInstance(overrides).then(firestoreInstance => { + firestore = firestoreInstance; + const query = firestore + .collection('collectionId') + .findNearest('vector', [1], {limit: 10, distanceMeasure: 'COSINE'}); + return query.get().then(results => { + expect(++counter).to.equal(1); + expect(results.size).to.equal(2); + expect(results.empty).to.be.false; + expect(results.readTime.isEqual(new Timestamp(5, 6))).to.be.true; + expect(results.docs[0].id).to.equal('first'); + expect(results.docs[1].id).to.equal('second'); + expect(results.docChanges()).to.have.length(2); + }); + }); + }); + + it('handles stream exception at initialization', async () => { + let attempts = 0; + const query = firestore + .collection('collectionId') + .findNearest('embedding', [1], { + limit: 100, + distanceMeasure: 'EUCLIDEAN', + }); + + query._queryUtil._stream = () => { + ++attempts; + throw new Error('Expected error'); + }; + + return query + .get() + .then(() => { + throw new Error('Unexpected success in Promise'); + }) + .catch(err => { + expect(err.message).to.equal('Expected error'); + expect(attempts).to.equal(1); + }); + }); + + it('handles stream exception during initialization', async () => { + let attempts = 0; + + const overrides: ApiOverride = { + runQuery: () => { + ++attempts; + return stream(new Error('Expected error')); + }, + }; + + return createInstance(overrides).then(firestoreInstance => { + firestore = firestoreInstance; + return firestore + .collection('collectionId') + .findNearest('embedding', [1], {limit: 10, distanceMeasure: 'COSINE'}) + .get() + .then(() => { + throw new Error('Unexpected success in Promise'); + }) + .catch(err => { + expect(err.message).to.equal('Expected error'); + expect(attempts).to.equal(5); + }); + }); + }); +}); diff --git a/docfx.json b/docfx.json new file mode 100644 index 000000000..e69de29bb diff --git a/package.json b/package.json index e867290e2..ad366aa0e 100644 --- a/package.json +++ b/package.json @@ -65,8 +65,8 @@ "dependencies": { "fast-deep-equal": "^3.1.1", "functional-red-black-tree": "^1.0.1", - "google-gax": "^4.0.4", - "protobufjs": "^7.2.5" + "google-gax": "^4.3.1", + "protobufjs": "^7.2.6" }, "devDependencies": { "@google-cloud/cloud-rad": "^0.4.0", diff --git a/types/firestore.d.ts b/types/firestore.d.ts index e293f030d..702e5ee84 100644 --- a/types/firestore.d.ts +++ b/types/firestore.d.ts @@ -2028,6 +2028,40 @@ declare namespace FirebaseFirestore { aggregateSpec: T ): AggregateQuery; + /** + * 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 `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'}); + * + * const querySnapshot = await aggregateQuery.get(); + * 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. + */ + findNearest( + vectorField: string | FieldPath, + queryVector: VectorValue | Array, + options: { + limit: number; + distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; + } + ): VectorQuery; + /** * Returns true if this `Query` is equal to the provided one. * @@ -2117,6 +2151,68 @@ declare namespace FirebaseFirestore { isEqual(other: QuerySnapshot): boolean; } + /** + * A `VectorQuerySnapshot` contains zero or more `QueryDocumentSnapshot` objects + * representing the results of a query. The documents can be accessed as an + * array via the `docs` property or enumerated using the `forEach` method. The + * number of documents can be determined via the `empty` and `size` + * properties. + */ + export class VectorQuerySnapshot< + AppModelType = DocumentData, + DbModelType extends DocumentData = DocumentData, + > { + private constructor(); + + /** + * The query on which you called `get` in order to get this + * `VectorQuerySnapshot`. + */ + readonly query: VectorQuery; + + /** An array of all the documents in the QuerySnapshot. */ + readonly docs: Array>; + + /** The number of documents in the QuerySnapshot. */ + readonly size: number; + + /** True if there are no documents in the QuerySnapshot. */ + readonly empty: boolean; + + /** The time this query snapshot was obtained. */ + readonly readTime: Timestamp; + + /** + * Returns an array of the documents changes since the last snapshot. If + * this is the first snapshot, all documents will be in the list as added + * changes. + */ + docChanges(): DocumentChange[]; + + /** + * Enumerates all of the documents in the QuerySnapshot. + * + * @param callback A callback to be called with a `DocumentSnapshot` for + * each document in the snapshot. + * @param thisArg The `this` binding for the callback. + */ + forEach( + callback: ( + result: QueryDocumentSnapshot + ) => void, + thisArg?: any + ): void; + + /** + * Returns true if the document data in this `VectorQuerySnapshot` is equal to the + * provided one. + * + * @param other The `VectorQuerySnapshot` to compare against. + * @return true if this `VectorQuerySnapshot` is equal to the provided one. + */ + isEqual(other: VectorQuerySnapshot): boolean; + } + /** * The type of `DocumentChange` may be 'added', 'removed', or 'modified'. */ @@ -2567,6 +2663,57 @@ declare namespace FirebaseFirestore { ): boolean; } + /** + * A query that finds the document whose vector fields are closest to a certain vector. + */ + export class VectorQuery< + AppModelType = DocumentData, + DbModelType extends DocumentData = DocumentData, + > { + private constructor(); + + /** The query whose results participants in the distance search. */ + readonly query: Query; + + /** + * Executes this query. + * + * @return A promise that will be resolved with the results of the query. + */ + get(): Promise>; + + /** + * Compares this object with the given object for equality. + * + * This object is considered "equal" to the other object if and only if + * `other` performs the same vector distance search as this `VectorQuery` and + * the underlying Query of `other` compares equal to that of this object + * using `Query.isEqual()`. + * + * @param other The object to compare to this object for equality. + * @return `true` if this object is "equal" to the given object, as + * defined above, or `false` otherwise. + */ + isEqual(other: VectorQuery): boolean; + } + + /** + * Represent a vector type in Firestore documents. + */ + export class VectorValue { + private constructor(values: number[] | undefined); + + /** + * Returns a copy of the raw number array form of the vector. + */ + toArray(): number[]; + + /** + * Returns true if the two `VectorValue` has the same raw number arrays, returns false otherwise. + */ + isEqual(other: VectorValue): boolean; + } + /** * Sentinel values that can be used when writing document fields with set(), * create() or update(). @@ -2637,6 +2784,11 @@ declare namespace FirebaseFirestore { */ static arrayRemove(...elements: any[]): FieldValue; + /** + * @return A new `VectorValue` constructed with a copy of the given array of number. + */ + static vector(values?: number[]): VectorValue; + /** * Returns true if this `FieldValue` is equal to the provided one. * diff --git a/types/protos/firestore_admin_v1_proto_api.d.ts b/types/protos/firestore_admin_v1_proto_api.d.ts index 216fa60d6..d236a5d6b 100644 --- a/types/protos/firestore_admin_v1_proto_api.d.ts +++ b/types/protos/firestore_admin_v1_proto_api.d.ts @@ -1291,6 +1291,9 @@ export namespace google { /** IndexField arrayConfig */ arrayConfig?: (google.firestore.admin.v1.Index.IndexField.ArrayConfig|null); + + /** IndexField vectorConfig */ + vectorConfig?: (google.firestore.admin.v1.Index.IndexField.IVectorConfig|null); } /** Represents an IndexField. */ @@ -1311,8 +1314,11 @@ export namespace google { /** IndexField arrayConfig. */ public arrayConfig?: (google.firestore.admin.v1.Index.IndexField.ArrayConfig|null); + /** IndexField vectorConfig. */ + public vectorConfig?: (google.firestore.admin.v1.Index.IndexField.IVectorConfig|null); + /** IndexField valueMode. */ - public valueMode?: ("order"|"arrayConfig"); + public valueMode?: ("order"|"arrayConfig"|"vectorConfig"); /** * Creates an IndexField message from a plain object. Also converts values to their respective internal types. @@ -1345,6 +1351,108 @@ export namespace google { /** ArrayConfig enum. */ type ArrayConfig = "ARRAY_CONFIG_UNSPECIFIED"| "CONTAINS"; + + /** Properties of a VectorConfig. */ + interface IVectorConfig { + + /** VectorConfig dimension */ + dimension?: (number|null); + + /** VectorConfig flat */ + flat?: (google.firestore.admin.v1.Index.IndexField.VectorConfig.IFlatIndex|null); + } + + /** Represents a VectorConfig. */ + class VectorConfig implements IVectorConfig { + + /** + * Constructs a new VectorConfig. + * @param [properties] Properties to set + */ + constructor(properties?: google.firestore.admin.v1.Index.IndexField.IVectorConfig); + + /** VectorConfig dimension. */ + public dimension: number; + + /** VectorConfig flat. */ + public flat?: (google.firestore.admin.v1.Index.IndexField.VectorConfig.IFlatIndex|null); + + /** VectorConfig type. */ + public type?: "flat"; + + /** + * Creates a VectorConfig message from a plain object. Also converts values to their respective internal types. + * @param object Plain object + * @returns VectorConfig + */ + public static fromObject(object: { [k: string]: any }): google.firestore.admin.v1.Index.IndexField.VectorConfig; + + /** + * Creates a plain object from a VectorConfig message. Also converts values to other types if specified. + * @param message VectorConfig + * @param [options] Conversion options + * @returns Plain object + */ + public static toObject(message: google.firestore.admin.v1.Index.IndexField.VectorConfig, options?: $protobuf.IConversionOptions): { [k: string]: any }; + + /** + * Converts this VectorConfig to JSON. + * @returns JSON object + */ + public toJSON(): { [k: string]: any }; + + /** + * Gets the default type url for VectorConfig + * @param [typeUrlPrefix] your custom typeUrlPrefix(default "type.googleapis.com") + * @returns The default type url + */ + public static getTypeUrl(typeUrlPrefix?: string): string; + } + + namespace VectorConfig { + + /** Properties of a FlatIndex. */ + interface IFlatIndex { + } + + /** Represents a FlatIndex. */ + class FlatIndex implements IFlatIndex { + + /** + * Constructs a new FlatIndex. + * @param [properties] Properties to set + */ + constructor(properties?: google.firestore.admin.v1.Index.IndexField.VectorConfig.IFlatIndex); + + /** + * Creates a FlatIndex message from a plain object. Also converts values to their respective internal types. + * @param object Plain object + * @returns FlatIndex + */ + public static fromObject(object: { [k: string]: any }): google.firestore.admin.v1.Index.IndexField.VectorConfig.FlatIndex; + + /** + * Creates a plain object from a FlatIndex message. Also converts values to other types if specified. + * @param message FlatIndex + * @param [options] Conversion options + * @returns Plain object + */ + public static toObject(message: google.firestore.admin.v1.Index.IndexField.VectorConfig.FlatIndex, options?: $protobuf.IConversionOptions): { [k: string]: any }; + + /** + * Converts this FlatIndex to JSON. + * @returns JSON object + */ + public toJSON(): { [k: string]: any }; + + /** + * Gets the default type url for FlatIndex + * @param [typeUrlPrefix] your custom typeUrlPrefix(default "type.googleapis.com") + * @returns The default type url + */ + public static getTypeUrl(typeUrlPrefix?: string): string; + } + } } /** QueryScope enum. */