From d406f14612a4890e405913aadc75c7ee22993f2b Mon Sep 17 00:00:00 2001 From: Mark Duckworth <1124037+MarkDuckworth@users.noreply.github.com> Date: Thu, 23 May 2024 11:50:17 -0600 Subject: [PATCH] feat: Query profiling for VectorQuery (#2045) feat: Query profiling for VectorQuery --- dev/src/reference/vector-query.ts | 50 +++++++++++++++++++--- dev/system-test/firestore.ts | 71 +++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/dev/src/reference/vector-query.ts b/dev/src/reference/vector-query.ts index b6b2dd0cc..4df93e60f 100644 --- a/dev/src/reference/vector-query.ts +++ b/dev/src/reference/vector-query.ts @@ -28,6 +28,8 @@ import {QueryUtil} from './query-util'; import {Query} from './query'; import {VectorQueryOptions} from './vector-query-options'; import {VectorQuerySnapshot} from './vector-query-snapshot'; +import {ExplainResults} from '../query-profile'; +import {QueryResponse} from './types'; /** * A query that finds the documents whose vector fields are closest to a certain query vector. @@ -92,24 +94,52 @@ export class VectorQuery< : this.queryVector.toArray(); } + /** + * Plans and optionally executes this vector search query. Returns a Promise that will be + * resolved with the planner information, statistics from the query execution (if any), + * and the query results (if any). + * + * @return A Promise that will be resolved with the planner information, statistics + * from the query execution (if any), and the query results (if any). + */ + async explain( + options?: firestore.ExplainOptions + ): Promise>> { + if (options === undefined) { + options = {}; + } + const {result, explainMetrics} = await this._getResponse(options); + if (!explainMetrics) { + throw new Error('No explain results'); + } + return new ExplainResults(explainMetrics, result || null); + } + /** * Executes this vector search query. * * @returns A promise that will be resolved with the results of the query. */ async get(): Promise> { - const {result} = await this._queryUtil._getResponse( - this, - /*transactionId*/ undefined, - // VectorQuery cannot be retried with cursors as they do not support cursors yet. - /*retryWithCursor*/ false - ); + const {result} = await this._getResponse(); if (!result) { throw new Error('No VectorQuerySnapshot result'); } return result; } + _getResponse( + explainOptions?: firestore.ExplainOptions + ): Promise>> { + return this._queryUtil._getResponse( + this, + /*transactionOrReadTime*/ undefined, + // VectorQuery cannot be retried with cursors as they do not support cursors yet. + /*retryWithCursor*/ false, + explainOptions + ); + } + /** * Internal streaming method that accepts an optional transaction ID. * @@ -135,7 +165,8 @@ export class VectorQuery< * @returns Serialized JSON for the query. */ toProto( - transactionOrReadTime?: Uint8Array | Timestamp | api.ITransactionOptions + transactionOrReadTime?: Uint8Array | Timestamp | api.ITransactionOptions, + explainOptions?: firestore.ExplainOptions ): api.IRunQueryRequest { const queryProto = this._query.toProto(transactionOrReadTime); @@ -151,6 +182,11 @@ export class VectorQuery< }, queryVector: queryVector._toProto(this._query._serializer), }; + + if (explainOptions) { + queryProto.explainOptions = explainOptions; + } + return queryProto; } diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts index c21cfdbc6..5f7d683d7 100644 --- a/dev/system-test/firestore.ts +++ b/dev/system-test/firestore.ts @@ -429,6 +429,77 @@ describe('Firestore class', () => { expect(explainResults.snapshot!.data().count).to.equal(3); }); + it('can plan a vector query', 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 explainResults = await indexTestHelper + .query(collectionReference) + .findNearest('embedding', FieldValue.vector([1, 3]), { + limit: 10, + distanceMeasure: 'COSINE', + }) + .explain({analyze: false}); + + const metrics = explainResults.metrics; + + const plan = metrics.planSummary; + expect(plan).to.not.be.null; + expect(Object.keys(plan.indexesUsed).length).to.be.greaterThan(0); + + expect(metrics.executionStats).to.be.null; + expect(explainResults.snapshot).to.be.null; + }); + + it('can profile a vector query', 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 explainResults = await indexTestHelper + .query(collectionReference) + .findNearest('embedding', FieldValue.vector([1, 3]), { + limit: 10, + distanceMeasure: 'COSINE', + }) + .explain({analyze: true}); + + const metrics = explainResults.metrics; + expect(metrics.planSummary).to.not.be.null; + expect( + Object.keys(metrics.planSummary.indexesUsed).length + ).to.be.greaterThan(0); + + expect(metrics.executionStats).to.not.be.null; + const stats = metrics.executionStats!; + + expect(stats.readOperations).to.be.greaterThan(0); + expect(stats.resultsReturned).to.be.equal(5); + expect( + stats.executionDuration.nanoseconds > 0 || + stats.executionDuration.seconds > 0 + ).to.be.true; + expect(Object.keys(stats.debugStats).length).to.be.greaterThan(0); + + expect(explainResults.snapshot).to.not.be.null; + expect(explainResults.snapshot!.docs.length).to.equal(5); + }); + it('getAll() supports array destructuring', () => { const ref1 = randomCol.doc('doc1'); const ref2 = randomCol.doc('doc2');