WIP: Firestore vector query

common/api-review/firestore-lite.api.md

@@ -226,6 +226,9 @@ export function getCount<AppModelType, DbModelType extends DocumentData>(query:
 export function getDoc<AppModelType, DbModelType extends DocumentData>(reference: DocumentReference<AppModelType, DbModelType>): Promise<DocumentSnapshot<AppModelType, DbModelType>>;
 
 // @public
+export function getDocs<AppModelType, DbModelType extends DocumentData>(query: VectorQuery<AppModelType, DbModelType>): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
+
+// @public (undocumented)
 export function getDocs<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>): Promise<QuerySnapshot<AppModelType, DbModelType>>;
 
 // @public
@@ -320,7 +323,7 @@ export class QueryEndAtConstraint extends QueryConstraint {
 }
 
 // @public
-export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType>): boolean;
+export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>): boolean;
 
 // @public
 export class QueryFieldFilterConstraint extends QueryConstraint {
@@ -390,7 +393,7 @@ export interface Settings {
 }
 
 // @public
-export function snapshotEqual<AppModelType, DbModelType extends DocumentData>(left: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType>, right: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType>): boolean;
+export function snapshotEqual<AppModelType, DbModelType extends DocumentData>(left: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType> | VectorQuerySnapshot<AppModelType, DbModelType>, right: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType> | VectorQuerySnapshot<AppModelType, DbModelType>): boolean;
 
 // @public
 export function startAfter<AppModelType, DbModelType extends DocumentData>(snapshot: DocumentSnapshot<AppModelType, DbModelType>): QueryStartAtConstraint;
@@ -463,6 +466,37 @@ export function updateDoc<AppModelType, DbModelType extends DocumentData>(refere
 // @public
 export function vector(values?: number[]): VectorValue;
 
+// @public (undocumented)
+export class VectorQuery<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
+    protected constructor();
+    get query(): Query<AppModelType, DbModelType>;
+}
+
+// @public (undocumented)
+export interface VectorQueryOptions {
+    // (undocumented)
+    distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT';
+    // (undocumented)
+    distanceResultField?: string | FieldPath;
+    // (undocumented)
+    distanceThreshold?: number;
+    // (undocumented)
+    limit: number;
+    // (undocumented)
+    queryVector: VectorValue | number[];
+    // (undocumented)
+    vectorField: string | FieldPath;
+}
+
+// @public (undocumented)
+export class VectorQuerySnapshot<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
+    get docs(): Array<QueryDocumentSnapshot<AppModelType, DbModelType>>;
+    get empty(): boolean;
+    forEach(callback: (result: QueryDocumentSnapshot<AppModelType, DbModelType>) => void, thisArg?: unknown): void;
+    readonly query: VectorQuery<AppModelType, DbModelType>;
+    get size(): number;
+}
+
 // @public
 export class VectorValue {
     /* Excluded from this release type: __constructor */

common/api-review/firestore.api.md

@@ -222,6 +222,9 @@ export abstract class FieldValue {
     abstract isEqual(other: FieldValue): boolean;
 }
 
+// @public (undocumented)
+export function findNearest<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>, options: VectorQueryOptions): VectorQuery<AppModelType, DbModelType>;
+
 // @public
 export class Firestore {
     get app(): FirebaseApp;
@@ -299,6 +302,9 @@ export function getDocsFromCache<AppModelType, DbModelType extends DocumentData>
 // @public
 export function getDocsFromServer<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>): Promise<QuerySnapshot<AppModelType, DbModelType>>;
 
+// @public
+export function getDocsFromServer<AppModelType, DbModelType extends DocumentData>(vectorQuery: VectorQuery<AppModelType, DbModelType>): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
+
 // @public
 export function getFirestore(): Firestore;
 
@@ -578,7 +584,7 @@ export class QueryEndAtConstraint extends QueryConstraint {
 }
 
 // @public
-export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType>): boolean;
+export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>): boolean;
 
 // @public
 export class QueryFieldFilterConstraint extends QueryConstraint {
@@ -748,6 +754,38 @@ export function updateDoc<AppModelType, DbModelType extends DocumentData>(refere
 // @public
 export function vector(values?: number[]): VectorValue;
 
+// @public (undocumented)
+export class VectorQuery<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
+    protected constructor();
+    get query(): Query<AppModelType, DbModelType>;
+}
+
+// @public (undocumented)
+export interface VectorQueryOptions {
+    // (undocumented)
+    distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT';
+    // (undocumented)
+    distanceResultField?: string | FieldPath;
+    // (undocumented)
+    distanceThreshold?: number;
+    // (undocumented)
+    limit: number;
+    // (undocumented)
+    queryVector: VectorValue | number[];
+    // (undocumented)
+    vectorField: string | FieldPath;
+}
+
+// @public
+export class VectorQuerySnapshot<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
+    get docs(): Array<QueryDocumentSnapshot<AppModelType, DbModelType>>;
+    get empty(): boolean;
+    forEach(callback: (result: QueryDocumentSnapshot<AppModelType, DbModelType>) => void, thisArg?: unknown): void;
+    readonly metadata: SnapshotMetadata;
+    readonly query: VectorQuery<AppModelType, DbModelType>;
+    get size(): number;
+}
+
 // @public
 export class VectorValue {
     /* Excluded from this release type: __constructor */

packages/firestore/lite/index.ts

@@ -140,6 +140,9 @@ export {
 } from '../src/lite-api/snapshot';
 
 export { VectorValue } from '../src/lite-api/vector_value';
+export { VectorQuery } from '../src/lite-api/vector_query';
+export { VectorQuerySnapshot } from '../src/lite-api/vector_query_snapshot';
+export { VectorQueryOptions } from '../src/lite-api/vector_query_options';
 
 export { WriteBatch, writeBatch } from '../src/lite-api/write_batch';
 

packages/firestore/package.json

@@ -120,6 +120,7 @@
     "chai-exclude": "2.1.0",
     "json-stable-stringify": "1.1.1",
     "protobufjs": "7.2.6",
+    "protobufjs-cli": "^1.1.3",
     "rollup": "2.79.1",
     "rollup-plugin-copy": "3.5.0",
     "rollup-plugin-copy-assets": "2.0.3",

packages/firestore/src/api.ts

@@ -176,7 +176,12 @@ export {
   vector
 } from './api/field_value_impl';
 
-export { VectorValue } from './lite-api/vector_value';
+export { VectorValue } from './api/vector_value';
+export { VectorQuery } from './api/vector_query';
+export { VectorQuerySnapshot } from './api/vector_query_snapshot';
+export { VectorQueryOptions } from './api/vector_query_options';
+
+export { findNearest } from './api/find_nearest';
 
 export { LogLevelString as LogLevel, setLogLevel } from './util/log';
 

packages/firestore/src/api/find_nearest.ts

@@ -0,0 +1,18 @@
+/**
+ * @license
+ * 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.
+ */
+
+export { findNearest } from '../lite-api/find_nearest';

packages/firestore/src/api/reference_impl.ts

@@ -24,6 +24,8 @@ import {
   NextFn,
   PartialObserver
 } from '../api/observer';
+import { QueryDocumentSnapshot } from '../api/snapshot';
+import { VectorQuerySnapshot } from '../api/vector_query_snapshot';
 import { ListenerDataSource } from '../core/event_manager';
 import {
   firestoreClientAddSnapshotsInSyncListener,
@@ -32,6 +34,7 @@ import {
   firestoreClientGetDocumentsViaSnapshotListener,
   firestoreClientGetDocumentViaSnapshotListener,
   firestoreClientListen,
+  firestoreClientRunVectorQuery,
   firestoreClientWrite
 } from '../core/firestore_client';
 import { newQueryForPath, Query as InternalQuery } from '../core/query';
@@ -59,6 +62,7 @@ import {
   parseUpdateVarargs
 } from '../lite-api/user_data_reader';
 import { AbstractUserDataWriter } from '../lite-api/user_data_writer';
+import { VectorQuery } from '../lite-api/vector_query';
 import { DeleteMutation, Mutation, Precondition } from '../model/mutation';
 import { debugAssert } from '../util/assert';
 import { ByteString } from '../util/byte_string';
@@ -270,17 +274,96 @@ export function getDocsFromServer<
   DbModelType extends DocumentData
 >(
   query: Query<AppModelType, DbModelType>
-): Promise<QuerySnapshot<AppModelType, DbModelType>> {
-  query = cast<Query<AppModelType, DbModelType>>(query, Query);
-  const firestore = cast(query.firestore, Firestore);
-  const client = ensureFirestoreConfigured(firestore);
-  const userDataWriter = new ExpUserDataWriter(firestore);
+): Promise<QuerySnapshot<AppModelType, DbModelType>>;
 
-  return firestoreClientGetDocumentsViaSnapshotListener(client, query._query, {
-    source: 'server'
-  }).then(
-    snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
-  );
+/**
+ * Executes the query and returns the results as a `QuerySnapshot` from the
+ * server. Returns an error if the network is not available.
+ *
+ * @returns A `Promise` that will be resolved with the results of the query.
+ */
+export function getDocsFromServer<
+  AppModelType,
+  DbModelType extends DocumentData
+>(
+  vectorQuery: VectorQuery<AppModelType, DbModelType>
+): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
+
+/**
+ * Executes the query and returns the results as a `QuerySnapshot` from the
+ * server. Returns an error if the network is not available.
+ *
+ * @returns A `Promise` that will be resolved with the results of the query.
+ */
+export function getDocsFromServer<
+  AppModelType,
+  DbModelType extends DocumentData
+>(
+  queryOrVectorQuery:
+    | Query<AppModelType, DbModelType>
+    | VectorQuery<AppModelType, DbModelType>
+): Promise<
+  | QuerySnapshot<AppModelType, DbModelType>
+  | VectorQuerySnapshot<AppModelType, DbModelType>
+> {
+  // If the query is a Query instance
+  if (
+    'type' in queryOrVectorQuery &&
+    (queryOrVectorQuery.type! === 'query' ||
+      queryOrVectorQuery.type! === 'collection')
+  ) {
+    const query = cast<Query<AppModelType, DbModelType>>(
+      queryOrVectorQuery,
+      Query
+    );
+
+    const firestore = cast(query.firestore, Firestore);
+    const client = ensureFirestoreConfigured(firestore);
+    const userDataWriter = new ExpUserDataWriter(firestore);
+
+    return firestoreClientGetDocumentsViaSnapshotListener(
+      client,
+      query._query,
+      {
+        source: 'server'
+      }
+    ).then(
+      snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
+    );
+  } else {
+    // the query is a VectorQuery instance
+    const vectorQuery: VectorQuery<AppModelType, DbModelType> = cast<
+      VectorQuery<AppModelType, DbModelType>
+    >(queryOrVectorQuery, VectorQuery);
+
+    const firestore = cast(vectorQuery.query.firestore, Firestore);
+    const client = ensureFirestoreConfigured(firestore);
+    const userDataWriter = new ExpUserDataWriter(firestore);
+
+    // Run the aggregation and convert the results
+    return firestoreClientRunVectorQuery(client, vectorQuery._vectorQuery).then(
+      result => {
+        const docs = result.map(
+          doc =>
+            new QueryDocumentSnapshot<AppModelType, DbModelType>(
+              firestore,
+              userDataWriter,
+              doc.key,
+              doc,
+              new SnapshotMetadata(false, false),
+              vectorQuery.query.converter
+            )
+        );
+
+        return new VectorQuerySnapshot<AppModelType, DbModelType>(
+          firestore,
+          userDataWriter,
+          vectorQuery,
+          docs
+        );
+      }
+    );
+  }
 }
 
 /**

packages/firestore/src/api/snapshot.ts

@@ -28,10 +28,12 @@ import {
 } from '../lite-api/reference';
 import {
   DocumentSnapshot as LiteDocumentSnapshot,
-  fieldPathFromArgument,
   FirestoreDataConverter as LiteFirestoreDataConverter
 } from '../lite-api/snapshot';
-import { UntypedFirestoreDataConverter } from '../lite-api/user_data_reader';
+import {
+  fieldPathFromArgument,
+  UntypedFirestoreDataConverter
+} from '../lite-api/user_data_reader';
 import { AbstractUserDataWriter } from '../lite-api/user_data_writer';
 import { Document } from '../model/document';
 import { DocumentKey } from '../model/document_key';

packages/firestore/src/api/vector_query.ts

@@ -0,0 +1,18 @@
+/**
+ * @license
+ * 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.
+ */
+
+export { VectorQuery } from '../lite-api/vector_query';

packages/firestore/src/api/vector_query_options.ts

@@ -0,0 +1,18 @@
+/**
+ * @license
+ * 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.
+ */
+
+export { VectorQueryOptions } from '../lite-api/vector_query_options';