Skip to content

Commit

Permalink
Vector config support (#6900)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@NickChittle
NickChittle committed Mar 25, 2024
1 parent dc13cb9 commit 90b6506
Show file tree
Hide file tree
Showing 6 changed files with 260 additions and 7 deletions.
18 changes: 18 additions & 0 deletions src/firestore/api-sort.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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;
}

Expand Down Expand Up @@ -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.
Expand Down
6 changes: 6 additions & 0 deletions src/firestore/api-types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@ export enum ArrayConfig {
CONTAINS = "CONTAINS",
}

export interface VectorConfig {
dimension: number;
flat?: {};
}

export enum State {
CREATING = "CREATING",
READY = "READY",
Expand Down Expand Up @@ -53,6 +58,7 @@ export interface IndexField {
fieldPath: string;
order?: Order;
arrayConfig?: ArrayConfig;
vectorConfig?: VectorConfig;
}

/**
Expand Down
9 changes: 8 additions & 1 deletion src/firestore/api.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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");
}
});
}

Expand Down Expand Up @@ -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 {
Expand Down
16 changes: 12 additions & 4 deletions src/firestore/pretty-print.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
150 changes: 148 additions & 2 deletions src/test/firestore/indexes.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down Expand Up @@ -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", () => {
Expand Down Expand Up @@ -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", () => {
Expand Down
68 changes: 68 additions & 0 deletions src/test/firestore/pretty-print.test.ts
View file
Original file line number Diff line number Diff line change
@@ -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>) ");
});
});

0 comments on commit 90b6506

Please sign in to comment.