Skip to content

Commit d64b354

Browse files
MarkDuckworthMarkDuckworth
committed
POC implementation of findNearest
1 parent ff0475c commit d64b354

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+4441
-1940
lines changed

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

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,9 @@ export function getCount<AppModelType, DbModelType extends DocumentData>(query:
226226
export function getDoc<AppModelType, DbModelType extends DocumentData>(reference: DocumentReference<AppModelType, DbModelType>): Promise<DocumentSnapshot<AppModelType, DbModelType>>;
227227

228228
// @public
229+
export function getDocs<AppModelType, DbModelType extends DocumentData>(query: VectorQuery<AppModelType, DbModelType>): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
230+
231+
// @public (undocumented)
229232
export function getDocs<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>): Promise<QuerySnapshot<AppModelType, DbModelType>>;
230233

231234
// @public
@@ -320,7 +323,7 @@ export class QueryEndAtConstraint extends QueryConstraint {
320323
}
321324

322325
// @public
323-
export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType>): boolean;
326+
export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>): boolean;
324327

325328
// @public
326329
export class QueryFieldFilterConstraint extends QueryConstraint {
@@ -390,7 +393,7 @@ export interface Settings {
390393
}
391394

392395
// @public
393-
export function snapshotEqual<AppModelType, DbModelType extends DocumentData>(left: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType>, right: DocumentSnapshot<AppModelType, DbModelType> | QuerySnapshot<AppModelType, DbModelType>): boolean;
396+
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;
394397

395398
// @public
396399
export function startAfter<AppModelType, DbModelType extends DocumentData>(snapshot: DocumentSnapshot<AppModelType, DbModelType>): QueryStartAtConstraint;
@@ -463,6 +466,37 @@ export function updateDoc<AppModelType, DbModelType extends DocumentData>(refere
463466
// @public
464467
export function vector(values?: number[]): VectorValue;
465468

469+
// @public (undocumented)
470+
export class VectorQuery<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
471+
protected constructor();
472+
get query(): Query<AppModelType, DbModelType>;
473+
}
474+
475+
// @public (undocumented)
476+
export interface VectorQueryOptions {
477+
// (undocumented)
478+
distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT';
479+
// (undocumented)
480+
distanceResultField?: string | FieldPath;
481+
// (undocumented)
482+
distanceThreshold?: number;
483+
// (undocumented)
484+
limit: number;
485+
// (undocumented)
486+
queryVector: VectorValue | number[];
487+
// (undocumented)
488+
vectorField: string | FieldPath;
489+
}
490+
491+
// @public (undocumented)
492+
export class VectorQuerySnapshot<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
493+
get docs(): Array<QueryDocumentSnapshot<AppModelType, DbModelType>>;
494+
get empty(): boolean;
495+
forEach(callback: (result: QueryDocumentSnapshot<AppModelType, DbModelType>) => void, thisArg?: unknown): void;
496+
readonly query: VectorQuery<AppModelType, DbModelType>;
497+
get size(): number;
498+
}
499+
466500
// @public
467501
export class VectorValue {
468502
/* Excluded from this release type: __constructor */

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

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,9 @@ export abstract class FieldValue {
222222
abstract isEqual(other: FieldValue): boolean;
223223
}
224224

225+
// @public (undocumented)
226+
export function findNearest<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>, options: VectorQueryOptions): VectorQuery<AppModelType, DbModelType>;
227+
225228
// @public
226229
export class Firestore {
227230
get app(): FirebaseApp;
@@ -299,6 +302,9 @@ export function getDocsFromCache<AppModelType, DbModelType extends DocumentData>
299302
// @public
300303
export function getDocsFromServer<AppModelType, DbModelType extends DocumentData>(query: Query<AppModelType, DbModelType>): Promise<QuerySnapshot<AppModelType, DbModelType>>;
301304

305+
// @public
306+
export function getDocsFromServer<AppModelType, DbModelType extends DocumentData>(vectorQuery: VectorQuery<AppModelType, DbModelType>): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
307+
302308
// @public
303309
export function getFirestore(): Firestore;
304310

@@ -578,7 +584,7 @@ export class QueryEndAtConstraint extends QueryConstraint {
578584
}
579585

580586
// @public
581-
export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType>): boolean;
587+
export function queryEqual<AppModelType, DbModelType extends DocumentData>(left: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>, right: Query<AppModelType, DbModelType> | VectorQuery<AppModelType, DbModelType>): boolean;
582588

583589
// @public
584590
export class QueryFieldFilterConstraint extends QueryConstraint {
@@ -748,6 +754,38 @@ export function updateDoc<AppModelType, DbModelType extends DocumentData>(refere
748754
// @public
749755
export function vector(values?: number[]): VectorValue;
750756

757+
// @public (undocumented)
758+
export class VectorQuery<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
759+
protected constructor();
760+
get query(): Query<AppModelType, DbModelType>;
761+
}
762+
763+
// @public (undocumented)
764+
export interface VectorQueryOptions {
765+
// (undocumented)
766+
distanceMeasure: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT';
767+
// (undocumented)
768+
distanceResultField?: string | FieldPath;
769+
// (undocumented)
770+
distanceThreshold?: number;
771+
// (undocumented)
772+
limit: number;
773+
// (undocumented)
774+
queryVector: VectorValue | number[];
775+
// (undocumented)
776+
vectorField: string | FieldPath;
777+
}
778+
779+
// @public
780+
export class VectorQuerySnapshot<AppModelType = DocumentData, DbModelType extends DocumentData = DocumentData> {
781+
get docs(): Array<QueryDocumentSnapshot<AppModelType, DbModelType>>;
782+
get empty(): boolean;
783+
forEach(callback: (result: QueryDocumentSnapshot<AppModelType, DbModelType>) => void, thisArg?: unknown): void;
784+
readonly metadata: SnapshotMetadata;
785+
readonly query: VectorQuery<AppModelType, DbModelType>;
786+
get size(): number;
787+
}
788+
751789
// @public
752790
export class VectorValue {
753791
/* Excluded from this release type: __constructor */

‎packages/firestore/lite/index.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,9 @@ export {
140140
} from '../src/lite-api/snapshot';
141141

142142
export { VectorValue } from '../src/lite-api/vector_value';
143+
export { VectorQuery } from '../src/lite-api/vector_query';
144+
export { VectorQuerySnapshot } from '../src/lite-api/vector_query_snapshot';
145+
export { VectorQueryOptions } from '../src/lite-api/vector_query_options';
143146

144147
export { WriteBatch, writeBatch } from '../src/lite-api/write_batch';
145148

‎packages/firestore/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@
120120
"chai-exclude": "2.1.0",
121121
"json-stable-stringify": "1.1.1",
122122
"protobufjs": "7.2.6",
123+
"protobufjs-cli": "^1.1.3",
123124
"rollup": "2.79.1",
124125
"rollup-plugin-copy": "3.5.0",
125126
"rollup-plugin-copy-assets": "2.0.3",

‎packages/firestore/src/api.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,12 @@ export {
176176
vector
177177
} from './api/field_value_impl';
178178

179-
export { VectorValue } from './lite-api/vector_value';
179+
export { VectorValue } from './api/vector_value';
180+
export { VectorQuery } from './api/vector_query';
181+
export { VectorQuerySnapshot } from './api/vector_query_snapshot';
182+
export { VectorQueryOptions } from './api/vector_query_options';
183+
184+
export { findNearest } from './api/find_nearest';
180185

181186
export { LogLevelString as LogLevel, setLogLevel } from './util/log';
182187

‎packages/firestore/src/api/find_nearest.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/**
2+
* @license
3+
* Copyright 2024 Google LLC
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
export { findNearest } from '../lite-api/find_nearest';

‎packages/firestore/src/api/reference_impl.ts

Lines changed: 93 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ import {
2424
NextFn,
2525
PartialObserver
2626
} from '../api/observer';
27+
import { QueryDocumentSnapshot } from '../api/snapshot';
28+
import { VectorQuerySnapshot } from '../api/vector_query_snapshot';
2729
import { ListenerDataSource } from '../core/event_manager';
2830
import {
2931
firestoreClientAddSnapshotsInSyncListener,
@@ -32,6 +34,7 @@ import {
3234
firestoreClientGetDocumentsViaSnapshotListener,
3335
firestoreClientGetDocumentViaSnapshotListener,
3436
firestoreClientListen,
37+
firestoreClientRunVectorQuery,
3538
firestoreClientWrite
3639
} from '../core/firestore_client';
3740
import { newQueryForPath, Query as InternalQuery } from '../core/query';
@@ -59,6 +62,7 @@ import {
5962
parseUpdateVarargs
6063
} from '../lite-api/user_data_reader';
6164
import { AbstractUserDataWriter } from '../lite-api/user_data_writer';
65+
import { VectorQuery } from '../lite-api/vector_query';
6266
import { DeleteMutation, Mutation, Precondition } from '../model/mutation';
6367
import { debugAssert } from '../util/assert';
6468
import { ByteString } from '../util/byte_string';
@@ -270,17 +274,96 @@ export function getDocsFromServer<
270274
DbModelType extends DocumentData
271275
>(
272276
query: Query<AppModelType, DbModelType>
273-
): Promise<QuerySnapshot<AppModelType, DbModelType>> {
274-
query = cast<Query<AppModelType, DbModelType>>(query, Query);
275-
const firestore = cast(query.firestore, Firestore);
276-
const client = ensureFirestoreConfigured(firestore);
277-
const userDataWriter = new ExpUserDataWriter(firestore);
277+
): Promise<QuerySnapshot<AppModelType, DbModelType>>;
278278

279-
return firestoreClientGetDocumentsViaSnapshotListener(client, query._query, {
280-
source: 'server'
281-
}).then(
282-
snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
283-
);
279+
/**
280+
* Executes the query and returns the results as a `QuerySnapshot` from the
281+
* server. Returns an error if the network is not available.
282+
*
283+
* @returns A `Promise` that will be resolved with the results of the query.
284+
*/
285+
export function getDocsFromServer<
286+
AppModelType,
287+
DbModelType extends DocumentData
288+
>(
289+
vectorQuery: VectorQuery<AppModelType, DbModelType>
290+
): Promise<VectorQuerySnapshot<AppModelType, DbModelType>>;
291+
292+
/**
293+
* Executes the query and returns the results as a `QuerySnapshot` from the
294+
* server. Returns an error if the network is not available.
295+
*
296+
* @returns A `Promise` that will be resolved with the results of the query.
297+
*/
298+
export function getDocsFromServer<
299+
AppModelType,
300+
DbModelType extends DocumentData
301+
>(
302+
queryOrVectorQuery:
303+
| Query<AppModelType, DbModelType>
304+
| VectorQuery<AppModelType, DbModelType>
305+
): Promise<
306+
| QuerySnapshot<AppModelType, DbModelType>
307+
| VectorQuerySnapshot<AppModelType, DbModelType>
308+
> {
309+
// If the query is a Query instance
310+
if (
311+
'type' in queryOrVectorQuery &&
312+
(queryOrVectorQuery.type! === 'query' ||
313+
queryOrVectorQuery.type! === 'collection')
314+
) {
315+
const query = cast<Query<AppModelType, DbModelType>>(
316+
queryOrVectorQuery,
317+
Query
318+
);
319+
320+
const firestore = cast(query.firestore, Firestore);
321+
const client = ensureFirestoreConfigured(firestore);
322+
const userDataWriter = new ExpUserDataWriter(firestore);
323+
324+
return firestoreClientGetDocumentsViaSnapshotListener(
325+
client,
326+
query._query,
327+
{
328+
source: 'server'
329+
}
330+
).then(
331+
snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
332+
);
333+
} else {
334+
// the query is a VectorQuery instance
335+
const vectorQuery: VectorQuery<AppModelType, DbModelType> = cast<
336+
VectorQuery<AppModelType, DbModelType>
337+
>(queryOrVectorQuery, VectorQuery);
338+
339+
const firestore = cast(vectorQuery.query.firestore, Firestore);
340+
const client = ensureFirestoreConfigured(firestore);
341+
const userDataWriter = new ExpUserDataWriter(firestore);
342+
343+
// Run the aggregation and convert the results
344+
return firestoreClientRunVectorQuery(client, vectorQuery._vectorQuery).then(
345+
result => {
346+
const docs = result.map(
347+
doc =>
348+
new QueryDocumentSnapshot<AppModelType, DbModelType>(
349+
firestore,
350+
userDataWriter,
351+
doc.key,
352+
doc,
353+
new SnapshotMetadata(false, false),
354+
vectorQuery.query.converter
355+
)
356+
);
357+
358+
return new VectorQuerySnapshot<AppModelType, DbModelType>(
359+
firestore,
360+
userDataWriter,
361+
vectorQuery,
362+
docs
363+
);
364+
}
365+
);
366+
}
284367
}
285368

286369
/**

‎packages/firestore/src/api/snapshot.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,12 @@ import {
2828
} from '../lite-api/reference';
2929
import {
3030
DocumentSnapshot as LiteDocumentSnapshot,
31-
fieldPathFromArgument,
3231
FirestoreDataConverter as LiteFirestoreDataConverter
3332
} from '../lite-api/snapshot';
34-
import { UntypedFirestoreDataConverter } from '../lite-api/user_data_reader';
33+
import {
34+
fieldPathFromArgument,
35+
UntypedFirestoreDataConverter
36+
} from '../lite-api/user_data_reader';
3537
import { AbstractUserDataWriter } from '../lite-api/user_data_writer';
3638
import { Document } from '../model/document';
3739
import { DocumentKey } from '../model/document_key';

‎packages/firestore/src/api/vector_query.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/**
2+
* @license
3+
* Copyright 2024 Google LLC
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
export { VectorQuery } from '../lite-api/vector_query';

‎packages/firestore/src/api/vector_query_options.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/**
2+
* @license
3+
* Copyright 2024 Google LLC
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
export { VectorQueryOptions } from '../lite-api/vector_query_options';

0 commit comments

Comments
 (0)