Skip to content

Commit efdf171

Browse files
authored
Add support for float16 data type in vector embedding (#36471)
### Packages impacted by this PR @azure/cosmos ### Issues associated with this PR ### Describe the problem that is addressed by this PR This PR adds support for `float16` data type in Vector Embeddings. ### What are the possible designs available to address the problem? If there are more than one possible design, why was the one in this PR chosen? ### Are there test cases added in this PR? _(If not, why?)_ Yes ### Provide a list of related PRs _(if any)_ ### Command used to generate this PR:**_(Applicable only to SDK release request PRs)_ ### Checklists - [ ] Added impacted package name to the issue description - [ ] Does this PR needs any fixes in the SDK Generator?** _(If so, create an Issue in the [Autorest/typescript](https://github.com/Azure/autorest.typescript) repository and link it here)_ - [ ] Added a changelog (if necessary)
1 parent d971f49 commit efdf171

File tree

3 files changed

+43
-4
lines changed

3 files changed

+43
-4
lines changed

sdk/cosmosdb/cosmos/review/cosmos-node.api.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2775,6 +2775,7 @@ export interface VectorEmbedding {
27752775

27762776
// @public
27772777
export enum VectorEmbeddingDataType {
2778+
Float16 = "float16",
27782779
Float32 = "float32",
27792780
Int8 = "int8",
27802781
UInt8 = "uint8"

sdk/cosmosdb/cosmos/src/documents/VectorEmbeddingPolicy.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ export interface VectorEmbedding {
3636
* Represents the data type of the vector.
3737
*/
3838
export enum VectorEmbeddingDataType {
39+
/**
40+
* 16-bit floating point number.
41+
*/
42+
Float16 = "float16",
3943
/**
4044
* 32-bit floating point number.
4145
*/

sdk/cosmosdb/cosmos/test/public/functional/NonStreamingQueryPolicy.spec.ts

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import type { Database } from "../../../src/client/Database/Database.js";
1212
import type { Container } from "../../../src/client/index.js";
1313
import { describe, it, assert, beforeAll, afterAll } from "vitest";
1414

15-
// Skipping these tests as they are not supported by public emulator
1615
describe("Vector search feature", async () => {
1716
describe("VectorEmbeddingPolicy", async () => {
1817
let database: Database;
@@ -70,8 +69,7 @@ describe("Vector search feature", async () => {
7069
}
7170
});
7271

73-
// skipping the test case for now. Will enable it once the changes are live on backend
74-
it.skip("validate VectorEmbeddingPolicy", async () => {
72+
it("validate VectorEmbeddingPolicy", async () => {
7573
const indexingPolicy: IndexingPolicy = {
7674
vectorIndexes: [
7775
{ path: "/vector1", type: VectorIndexType.Flat },
@@ -126,7 +124,6 @@ describe("Vector search feature", async () => {
126124
assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[0].path === "/vector1");
127125
assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[1].path === "/vector2");
128126
assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[2].path === "/vector3");
129-
130127
assert(containerdef.indexingPolicy.vectorIndexes.length === 3);
131128
assert(containerdef.indexingPolicy.vectorIndexes[0].path === "/vector1");
132129
assert(containerdef.indexingPolicy.vectorIndexes[1].path === "/vector2");
@@ -142,7 +139,44 @@ describe("Vector search feature", async () => {
142139
assert(containerdef.indexingPolicy.vectorIndexes[1].vectorIndexShardKey[0] === "/Country");
143140
assert(containerdef.indexingPolicy.vectorIndexes[2].vectorIndexShardKey[0] === "/ZipCode");
144141
});
142+
it.skip("validate VectorEmbeddingPolicy with Float16 data type", async () => {
143+
const indexingPolicy: IndexingPolicy = {
144+
vectorIndexes: [{ path: "/vector1", type: VectorIndexType.QuantizedFlat }],
145+
};
146+
const vectorEmbeddingPolicy: VectorEmbeddingPolicy = {
147+
vectorEmbeddings: [
148+
{
149+
path: "/vector1",
150+
dataType: VectorEmbeddingDataType.Float16,
151+
dimensions: 500,
152+
distanceFunction: VectorEmbeddingDistanceFunction.Euclidean,
153+
},
154+
],
155+
};
156+
const containerName = "JSApp-vector embedding container float16";
157+
try {
158+
const { resource: containerdef } = await database.containers.createIfNotExists({
159+
id: containerName,
160+
vectorEmbeddingPolicy: vectorEmbeddingPolicy,
161+
indexingPolicy: indexingPolicy,
162+
});
145163

164+
assert(containerdef.vectorEmbeddingPolicy !== undefined);
165+
assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings.length === 1);
166+
assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[0].path === "/vector1");
167+
assert(
168+
containerdef.vectorEmbeddingPolicy.vectorEmbeddings[0].dataType ===
169+
VectorEmbeddingDataType.Float16,
170+
);
171+
} catch (e) {
172+
assert.fail(`Container creation failed for Float16 data type with error: ${e}`);
173+
}
174+
try {
175+
await database.container(containerName).delete();
176+
} catch {
177+
// Ignore if container is not found
178+
}
179+
});
146180
it("should fail to create vector indexing policy", async () => {
147181
const vectorEmbeddingPolicy: VectorEmbeddingPolicy = {
148182
vectorEmbeddings: [

0 commit comments

Comments
 (0)