From 90b65061eadcdbeeadf5a4bd229c0b4d6bd614d8 Mon Sep 17 00:00:00 2001 From: NickChittle Date: Mon, 25 Mar 2024 14:43:10 -0400 Subject: [PATCH] Vector config support (#6900) * Support Vector Config in the indexes.json file * Fix up unnecessary dependencies in test file * Fix up linter warnings * Fixes for printer test * Some small adjustments to the sort method --- src/firestore/api-sort.ts | 18 +++ src/firestore/api-types.ts | 6 + src/firestore/api.ts | 9 +- src/firestore/pretty-print.ts | 16 ++- src/test/firestore/indexes.spec.ts | 150 +++++++++++++++++++++++- src/test/firestore/pretty-print.test.ts | 68 +++++++++++ 6 files changed, 260 insertions(+), 7 deletions(-) create mode 100644 src/test/firestore/pretty-print.test.ts diff --git a/src/firestore/api-sort.ts b/src/firestore/api-sort.ts index 387bfa31f2e..3991a831749 100644 --- a/src/firestore/api-sort.ts +++ b/src/firestore/api-sort.ts @@ -180,6 +180,7 @@ export function compareFieldOverride(a: Spec.FieldOverride, b: Spec.FieldOverrid * 1) Field path. * 2) Sort order (if it exists). * 3) Array config (if it exists). + * 4) Vector config (if it exists). */ function compareIndexField(a: API.IndexField, b: API.IndexField): number { if (a.fieldPath !== b.fieldPath) { @@ -194,6 +195,10 @@ function compareIndexField(a: API.IndexField, b: API.IndexField): number { return compareArrayConfig(a.arrayConfig, b.arrayConfig); } + if (a.vectorConfig !== b.vectorConfig) { + return compareVectorConfig(a.vectorConfig, b.vectorConfig); + } + return 0; } @@ -225,6 +230,19 @@ function compareArrayConfig(a?: API.ArrayConfig, b?: API.ArrayConfig): number { return ARRAY_CONFIG_SEQUENCE.indexOf(a) - ARRAY_CONFIG_SEQUENCE.indexOf(b); } +function compareVectorConfig(a?: API.VectorConfig, b?: API.VectorConfig): number { + if (!a) { + if (!b) { + return 0; + } else { + return 1; + } + } else if (!b) { + return -1; + } + return a.dimension - b.dimension; +} + /** * Compare two arrays of objects by looking for the first * non-equal element and comparing them. diff --git a/src/firestore/api-types.ts b/src/firestore/api-types.ts index f72843f32a4..9e564947923 100644 --- a/src/firestore/api-types.ts +++ b/src/firestore/api-types.ts @@ -24,6 +24,11 @@ export enum ArrayConfig { CONTAINS = "CONTAINS", } +export interface VectorConfig { + dimension: number; + flat?: {}; +} + export enum State { CREATING = "CREATING", READY = "READY", @@ -53,6 +58,7 @@ export interface IndexField { fieldPath: string; order?: Order; arrayConfig?: ArrayConfig; + vectorConfig?: VectorConfig; } /** diff --git a/src/firestore/api.ts b/src/firestore/api.ts index ddc204917aa..8360f5f82df 100644 --- a/src/firestore/api.ts +++ b/src/firestore/api.ts @@ -297,7 +297,7 @@ export class FirestoreApi { index.fields.forEach((field: any) => { validator.assertHas(field, "fieldPath"); - validator.assertHasOneOf(field, ["order", "arrayConfig"]); + validator.assertHasOneOf(field, ["order", "arrayConfig", "vectorConfig"]); if (field.order) { validator.assertEnum(field, "order", Object.keys(types.Order)); @@ -306,6 +306,11 @@ export class FirestoreApi { if (field.arrayConfig) { validator.assertEnum(field, "arrayConfig", Object.keys(types.ArrayConfig)); } + + if (field.vectorConfig) { + validator.assertType("vectorConfig.dimension", field.vectorConfig.dimension, "number"); + validator.assertHas(field.vectorConfig, "flat"); + } }); } @@ -548,6 +553,8 @@ export class FirestoreApi { f.order = field.order; } else if (field.arrayConfig) { f.arrayConfig = field.arrayConfig; + } else if (field.vectorConfig) { + f.vectorConfig = field.vectorConfig; } else if (field.mode === types.Mode.ARRAY_CONTAINS) { f.arrayConfig = types.ArrayConfig.CONTAINS; } else { diff --git a/src/firestore/pretty-print.ts b/src/firestore/pretty-print.ts index c4e69ba6f80..0fdd4e4d3a4 100644 --- a/src/firestore/pretty-print.ts +++ b/src/firestore/pretty-print.ts @@ -236,10 +236,18 @@ export class PrettyPrint { return; } - // Normal field indexes have an "order" while array indexes have an "arrayConfig", - // we want to display whichever one is present. - const orderOrArrayConfig = field.order ? field.order : field.arrayConfig; - result += `(${field.fieldPath},${orderOrArrayConfig}) `; + // Normal field indexes have an "order", array indexes have an + // "arrayConfig", and vector indexes have a "vectorConfig" we want to + // display whichever one is present. + let configString; + if (field.order) { + configString = field.order; + } else if (field.arrayConfig) { + configString = field.arrayConfig; + } else if (field.vectorConfig) { + configString = `VECTOR<${field.vectorConfig.dimension}>`; + } + result += `(${field.fieldPath},${configString}) `; }); return result; diff --git a/src/test/firestore/indexes.spec.ts b/src/test/firestore/indexes.spec.ts index da196fbd34e..b54dd00ad99 100644 --- a/src/test/firestore/indexes.spec.ts +++ b/src/test/firestore/indexes.spec.ts @@ -66,6 +66,112 @@ describe("IndexValidation", () => { ); }); + it("should accept a valid vectorConfig index", () => { + idx.validateSpec( + idx.upgradeOldSpec({ + indexes: [ + { + collectionGroup: "collection", + queryScope: "COLLECTION", + fields: [ + { + fieldPath: "embedding", + vectorConfig: { + dimension: 100, + flat: {}, + }, + }, + ], + }, + ], + }), + ); + }); + + it("should accept a valid vectorConfig index after upgrade", () => { + idx.validateSpec({ + indexes: [ + { + collectionGroup: "collection", + queryScope: "COLLECTION", + fields: [ + { + fieldPath: "embedding", + vectorConfig: { + dimension: 100, + flat: {}, + }, + }, + ], + }, + ], + }); + }); + + it("should accept a valid vectorConfig index with another field", () => { + idx.validateSpec({ + indexes: [ + { + collectionGroup: "collection", + queryScope: "COLLECTION", + fields: [ + { fieldPath: "foo", order: "ASCENDING" }, + { + fieldPath: "embedding", + vectorConfig: { + dimension: 100, + flat: {}, + }, + }, + ], + }, + ], + }); + }); + + it("should reject invalid vectorConfig dimension", () => { + expect(() => { + idx.validateSpec({ + indexes: [ + { + collectionGroup: "collection", + queryScope: "COLLECTION", + fields: [ + { + fieldPath: "embedding", + vectorConfig: { + dimension: "wrongType", + flat: {}, + }, + }, + ], + }, + ], + }); + }).to.throw(FirebaseError, /Property "vectorConfig.dimension" must be of type number/); + }); + + it("should reject invalid vectorConfig missing flat type", () => { + expect(() => { + idx.validateSpec({ + indexes: [ + { + collectionGroup: "collection", + queryScope: "COLLECTION", + fields: [ + { + fieldPath: "embedding", + vectorConfig: { + dimension: 100, + }, + }, + ], + }, + ], + }); + }).to.throw(FirebaseError, /Must contain "flat"/); + }); + it("should reject an incomplete index spec", () => { expect(() => { idx.validateSpec({ @@ -96,7 +202,7 @@ describe("IndexValidation", () => { }, ], }); - }).to.throw(FirebaseError, /Must contain exactly one of "order,arrayConfig"/); + }).to.throw(FirebaseError, /Must contain exactly one of "order,arrayConfig,vectorConfig"/); }); }); describe("IndexSpecMatching", () => { @@ -532,7 +638,47 @@ describe("IndexSorting", () => { ], }; - expect([b, a, d, c].sort(sort.compareApiIndex)).to.eql([a, b, c, d]); + const e: API.Index = { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/e", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { + fieldPath: "fieldA", + vectorConfig: { + dimension: 100, + flat: {}, + }, + }, + ], + }; + + const f: API.Index = { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/f", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { + fieldPath: "fieldA", + vectorConfig: { + dimension: 200, + flat: {}, + }, + }, + ], + }; + + // This Index is invalid, but is used to verify sort ordering on undefined + // fields. + const g: API.Index = { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/g", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { + fieldPath: "fieldA", + }, + ], + }; + + expect([b, a, d, g, f, e, c].sort(sort.compareApiIndex)).to.eql([a, b, c, d, e, f, g]); }); it("should correctly sort an array of API field overrides", () => { diff --git a/src/test/firestore/pretty-print.test.ts b/src/test/firestore/pretty-print.test.ts new file mode 100644 index 00000000000..1daf10c5d67 --- /dev/null +++ b/src/test/firestore/pretty-print.test.ts @@ -0,0 +1,68 @@ +import { expect } from "chai"; +import * as API from "../../firestore/api-types"; +import { PrettyPrint } from "../../firestore/pretty-print"; + +const printer = new PrettyPrint(); + +describe("prettyIndexString", () => { + it("should correctly print an order type Index", () => { + expect( + printer.prettyIndexString( + { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { fieldPath: "foo", order: API.Order.ASCENDING }, + { fieldPath: "bar", order: API.Order.DESCENDING }, + ], + }, + false, + ), + ).to.contain("(foo,ASCENDING) (bar,DESCENDING) "); + }); + + it("should correctly print a contains type Index", () => { + expect( + printer.prettyIndexString( + { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { fieldPath: "foo", order: API.Order.ASCENDING }, + { fieldPath: "baz", arrayConfig: API.ArrayConfig.CONTAINS }, + ], + }, + false, + ), + ).to.contain("(foo,ASCENDING) (baz,CONTAINS) "); + }); + + it("should correctly print a vector type Index", () => { + expect( + printer.prettyIndexString( + { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a", + queryScope: API.QueryScope.COLLECTION, + fields: [{ fieldPath: "foo", vectorConfig: { dimension: 100, flat: {} } }], + }, + false, + ), + ).to.contain("(foo,VECTOR<100>) "); + }); + + it("should correctly print a vector type Index with other fields", () => { + expect( + printer.prettyIndexString( + { + name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a", + queryScope: API.QueryScope.COLLECTION, + fields: [ + { fieldPath: "foo", order: API.Order.ASCENDING }, + { fieldPath: "bar", vectorConfig: { dimension: 200, flat: {} } }, + ], + }, + false, + ), + ).to.contain("(foo,ASCENDING) (bar,VECTOR<200>) "); + }); +});