feat: Query profiling for VectorQuery (#2045)

feat: Query profiling for VectorQuery
googleapis · May 23, 2024 · d406f14 · d406f14
1 parent 1e949b8
commit d406f14
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 7 deletions.
diff --git 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<ExplainResults<VectorQuerySnapshot<AppModelType, DbModelType>>> {
+    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<VectorQuerySnapshot<AppModelType, DbModelType>> {
-    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<QueryResponse<VectorQuerySnapshot<AppModelType, DbModelType>>> {
+    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
@@ -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');