Skip to content

Commit 90b6506

Browse files
NickChittleNickChittle
authored
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
1 parent dc13cb9 commit 90b6506

File tree

6 files changed

+260
-7
lines changed

6 files changed

+260
-7
lines changed

‎src/firestore/api-sort.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,7 @@ export function compareFieldOverride(a: Spec.FieldOverride, b: Spec.FieldOverrid
180180
* 1) Field path.
181181
* 2) Sort order (if it exists).
182182
* 3) Array config (if it exists).
183+
* 4) Vector config (if it exists).
183184
*/
184185
function compareIndexField(a: API.IndexField, b: API.IndexField): number {
185186
if (a.fieldPath !== b.fieldPath) {
@@ -194,6 +195,10 @@ function compareIndexField(a: API.IndexField, b: API.IndexField): number {
194195
return compareArrayConfig(a.arrayConfig, b.arrayConfig);
195196
}
196197

198+
if (a.vectorConfig !== b.vectorConfig) {
199+
return compareVectorConfig(a.vectorConfig, b.vectorConfig);
200+
}
201+
197202
return 0;
198203
}
199204

@@ -225,6 +230,19 @@ function compareArrayConfig(a?: API.ArrayConfig, b?: API.ArrayConfig): number {
225230
return ARRAY_CONFIG_SEQUENCE.indexOf(a) - ARRAY_CONFIG_SEQUENCE.indexOf(b);
226231
}
227232

233+
function compareVectorConfig(a?: API.VectorConfig, b?: API.VectorConfig): number {
234+
if (!a) {
235+
if (!b) {
236+
return 0;
237+
} else {
238+
return 1;
239+
}
240+
} else if (!b) {
241+
return -1;
242+
}
243+
return a.dimension - b.dimension;
244+
}
245+
228246
/**
229247
* Compare two arrays of objects by looking for the first
230248
* non-equal element and comparing them.

‎src/firestore/api-types.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,11 @@ export enum ArrayConfig {
2424
CONTAINS = "CONTAINS",
2525
}
2626

27+
export interface VectorConfig {
28+
dimension: number;
29+
flat?: {};
30+
}
31+
2732
export enum State {
2833
CREATING = "CREATING",
2934
READY = "READY",
@@ -53,6 +58,7 @@ export interface IndexField {
5358
fieldPath: string;
5459
order?: Order;
5560
arrayConfig?: ArrayConfig;
61+
vectorConfig?: VectorConfig;
5662
}
5763

5864
/**

‎src/firestore/api.ts

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,7 @@ export class FirestoreApi {
297297

298298
index.fields.forEach((field: any) => {
299299
validator.assertHas(field, "fieldPath");
300-
validator.assertHasOneOf(field, ["order", "arrayConfig"]);
300+
validator.assertHasOneOf(field, ["order", "arrayConfig", "vectorConfig"]);
301301

302302
if (field.order) {
303303
validator.assertEnum(field, "order", Object.keys(types.Order));
@@ -306,6 +306,11 @@ export class FirestoreApi {
306306
if (field.arrayConfig) {
307307
validator.assertEnum(field, "arrayConfig", Object.keys(types.ArrayConfig));
308308
}
309+
310+
if (field.vectorConfig) {
311+
validator.assertType("vectorConfig.dimension", field.vectorConfig.dimension, "number");
312+
validator.assertHas(field.vectorConfig, "flat");
313+
}
309314
});
310315
}
311316

@@ -548,6 +553,8 @@ export class FirestoreApi {
548553
f.order = field.order;
549554
} else if (field.arrayConfig) {
550555
f.arrayConfig = field.arrayConfig;
556+
} else if (field.vectorConfig) {
557+
f.vectorConfig = field.vectorConfig;
551558
} else if (field.mode === types.Mode.ARRAY_CONTAINS) {
552559
f.arrayConfig = types.ArrayConfig.CONTAINS;
553560
} else {

‎src/firestore/pretty-print.ts

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -236,10 +236,18 @@ export class PrettyPrint {
236236
return;
237237
}
238238

239-
// Normal field indexes have an "order" while array indexes have an "arrayConfig",
240-
// we want to display whichever one is present.
241-
const orderOrArrayConfig = field.order ? field.order : field.arrayConfig;
242-
result += `(${field.fieldPath},${orderOrArrayConfig}) `;
239+
// Normal field indexes have an "order", array indexes have an
240+
// "arrayConfig", and vector indexes have a "vectorConfig" we want to
241+
// display whichever one is present.
242+
let configString;
243+
if (field.order) {
244+
configString = field.order;
245+
} else if (field.arrayConfig) {
246+
configString = field.arrayConfig;
247+
} else if (field.vectorConfig) {
248+
configString = `VECTOR<${field.vectorConfig.dimension}>`;
249+
}
250+
result += `(${field.fieldPath},${configString}) `;
243251
});
244252

245253
return result;

‎src/test/firestore/indexes.spec.ts

Lines changed: 148 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,112 @@ describe("IndexValidation", () => {
6666
);
6767
});
6868

69+
it("should accept a valid vectorConfig index", () => {
70+
idx.validateSpec(
71+
idx.upgradeOldSpec({
72+
indexes: [
73+
{
74+
collectionGroup: "collection",
75+
queryScope: "COLLECTION",
76+
fields: [
77+
{
78+
fieldPath: "embedding",
79+
vectorConfig: {
80+
dimension: 100,
81+
flat: {},
82+
},
83+
},
84+
],
85+
},
86+
],
87+
}),
88+
);
89+
});
90+
91+
it("should accept a valid vectorConfig index after upgrade", () => {
92+
idx.validateSpec({
93+
indexes: [
94+
{
95+
collectionGroup: "collection",
96+
queryScope: "COLLECTION",
97+
fields: [
98+
{
99+
fieldPath: "embedding",
100+
vectorConfig: {
101+
dimension: 100,
102+
flat: {},
103+
},
104+
},
105+
],
106+
},
107+
],
108+
});
109+
});
110+
111+
it("should accept a valid vectorConfig index with another field", () => {
112+
idx.validateSpec({
113+
indexes: [
114+
{
115+
collectionGroup: "collection",
116+
queryScope: "COLLECTION",
117+
fields: [
118+
{ fieldPath: "foo", order: "ASCENDING" },
119+
{
120+
fieldPath: "embedding",
121+
vectorConfig: {
122+
dimension: 100,
123+
flat: {},
124+
},
125+
},
126+
],
127+
},
128+
],
129+
});
130+
});
131+
132+
it("should reject invalid vectorConfig dimension", () => {
133+
expect(() => {
134+
idx.validateSpec({
135+
indexes: [
136+
{
137+
collectionGroup: "collection",
138+
queryScope: "COLLECTION",
139+
fields: [
140+
{
141+
fieldPath: "embedding",
142+
vectorConfig: {
143+
dimension: "wrongType",
144+
flat: {},
145+
},
146+
},
147+
],
148+
},
149+
],
150+
});
151+
}).to.throw(FirebaseError, /Property "vectorConfig.dimension" must be of type number/);
152+
});
153+
154+
it("should reject invalid vectorConfig missing flat type", () => {
155+
expect(() => {
156+
idx.validateSpec({
157+
indexes: [
158+
{
159+
collectionGroup: "collection",
160+
queryScope: "COLLECTION",
161+
fields: [
162+
{
163+
fieldPath: "embedding",
164+
vectorConfig: {
165+
dimension: 100,
166+
},
167+
},
168+
],
169+
},
170+
],
171+
});
172+
}).to.throw(FirebaseError, /Must contain "flat"/);
173+
});
174+
69175
it("should reject an incomplete index spec", () => {
70176
expect(() => {
71177
idx.validateSpec({
@@ -96,7 +202,7 @@ describe("IndexValidation", () => {
96202
},
97203
],
98204
});
99-
}).to.throw(FirebaseError, /Must contain exactly one of "order,arrayConfig"/);
205+
}).to.throw(FirebaseError, /Must contain exactly one of "order,arrayConfig,vectorConfig"/);
100206
});
101207
});
102208
describe("IndexSpecMatching", () => {
@@ -532,7 +638,47 @@ describe("IndexSorting", () => {
532638
],
533639
};
534640

535-
expect([b, a, d, c].sort(sort.compareApiIndex)).to.eql([a, b, c, d]);
641+
const e: API.Index = {
642+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/e",
643+
queryScope: API.QueryScope.COLLECTION,
644+
fields: [
645+
{
646+
fieldPath: "fieldA",
647+
vectorConfig: {
648+
dimension: 100,
649+
flat: {},
650+
},
651+
},
652+
],
653+
};
654+
655+
const f: API.Index = {
656+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/f",
657+
queryScope: API.QueryScope.COLLECTION,
658+
fields: [
659+
{
660+
fieldPath: "fieldA",
661+
vectorConfig: {
662+
dimension: 200,
663+
flat: {},
664+
},
665+
},
666+
],
667+
};
668+
669+
// This Index is invalid, but is used to verify sort ordering on undefined
670+
// fields.
671+
const g: API.Index = {
672+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/g",
673+
queryScope: API.QueryScope.COLLECTION,
674+
fields: [
675+
{
676+
fieldPath: "fieldA",
677+
},
678+
],
679+
};
680+
681+
expect([b, a, d, g, f, e, c].sort(sort.compareApiIndex)).to.eql([a, b, c, d, e, f, g]);
536682
});
537683

538684
it("should correctly sort an array of API field overrides", () => {

‎src/test/firestore/pretty-print.test.ts

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import { expect } from "chai";
2+
import * as API from "../../firestore/api-types";
3+
import { PrettyPrint } from "../../firestore/pretty-print";
4+
5+
const printer = new PrettyPrint();
6+
7+
describe("prettyIndexString", () => {
8+
it("should correctly print an order type Index", () => {
9+
expect(
10+
printer.prettyIndexString(
11+
{
12+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a",
13+
queryScope: API.QueryScope.COLLECTION,
14+
fields: [
15+
{ fieldPath: "foo", order: API.Order.ASCENDING },
16+
{ fieldPath: "bar", order: API.Order.DESCENDING },
17+
],
18+
},
19+
false,
20+
),
21+
).to.contain("(foo,ASCENDING) (bar,DESCENDING) ");
22+
});
23+
24+
it("should correctly print a contains type Index", () => {
25+
expect(
26+
printer.prettyIndexString(
27+
{
28+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a",
29+
queryScope: API.QueryScope.COLLECTION,
30+
fields: [
31+
{ fieldPath: "foo", order: API.Order.ASCENDING },
32+
{ fieldPath: "baz", arrayConfig: API.ArrayConfig.CONTAINS },
33+
],
34+
},
35+
false,
36+
),
37+
).to.contain("(foo,ASCENDING) (baz,CONTAINS) ");
38+
});
39+
40+
it("should correctly print a vector type Index", () => {
41+
expect(
42+
printer.prettyIndexString(
43+
{
44+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a",
45+
queryScope: API.QueryScope.COLLECTION,
46+
fields: [{ fieldPath: "foo", vectorConfig: { dimension: 100, flat: {} } }],
47+
},
48+
false,
49+
),
50+
).to.contain("(foo,VECTOR<100>) ");
51+
});
52+
53+
it("should correctly print a vector type Index with other fields", () => {
54+
expect(
55+
printer.prettyIndexString(
56+
{
57+
name: "/projects/project/databases/(default)/collectionGroups/collectionB/indexes/a",
58+
queryScope: API.QueryScope.COLLECTION,
59+
fields: [
60+
{ fieldPath: "foo", order: API.Order.ASCENDING },
61+
{ fieldPath: "bar", vectorConfig: { dimension: 200, flat: {} } },
62+
],
63+
},
64+
false,
65+
),
66+
).to.contain("(foo,ASCENDING) (bar,VECTOR<200>) ");
67+
});
68+
});

0 commit comments

Comments
 (0)