From e2b9711305efb80824909e7295b138bf64fe58cd Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 13:41:19 -0800 Subject: [PATCH 1/8] Adds Support for QuantizerType in IndexingPolicy --- .../com/azure/cosmos/rx/VectorIndexTest.java | 51 +++++++++++++++++-- .../cosmos/implementation/Constants.java | 1 + .../cosmos/models/CosmosVectorIndexSpec.java | 30 +++++++++++ .../azure/cosmos/models/QuantizerType.java | 32 ++++++++++++ 4 files changed, 110 insertions(+), 4 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index 6390f8ba2df0..196d2d275745 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -26,6 +26,7 @@ import com.azure.cosmos.models.PartitionKeyDefinition; import com.azure.cosmos.models.CosmosVectorIndexSpec; import com.azure.cosmos.models.CosmosVectorIndexType; +import com.azure.cosmos.models.QuantizerType; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; @@ -279,6 +280,22 @@ public void shouldValidateVectorEmbeddingPolicySerializationAndDeserialization() validateVectorEmbeddingPolicy(cosmosVectorEmbeddingPolicy, expectedCosmosVectorEmbeddingPolicy); } + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void shouldValidateVectorIndexesSerializationAndDeserialization() throws JsonProcessingException { + IndexingPolicy indexingPolicy = new IndexingPolicy(); + indexingPolicy.setVectorIndexes(populateVectorIndexes()); + List expectedVectorIndexes = indexingPolicy.getVectorIndexes(); + + // Validate Vector Indexes Serialization + String actualVectorIndexesJSON = simpleObjectMapper.writeValueAsString(expectedVectorIndexes); + String expectedVectorIndexesJSON = getVectorIndexesAsString(); + assertThat(expectedVectorIndexesJSON).isEqualTo(actualVectorIndexesJSON); + + // Validate Vector Indexes Deserialization + List actualVectorIndexes = Arrays.asList(simpleObjectMapper.readValue(actualVectorIndexesJSON, CosmosVectorIndexSpec[].class)); + validateVectorIndexes(expectedVectorIndexes, actualVectorIndexes); + } + private void validateCollectionProperties(CosmosContainerProperties collectionDefinition, CosmosContainerProperties collectionProperties) { assertThat(collectionProperties.getVectorEmbeddingPolicy()).isNotNull(); assertThat(collectionProperties.getVectorEmbeddingPolicy().getVectorEmbeddings()).isNotNull(); @@ -301,13 +318,14 @@ private void validateVectorEmbeddingPolicy(CosmosVectorEmbeddingPolicy actual, C } } - private void validateVectorIndexes(List actual, List expected) { + private void validateVectorIndexes(List expected, List actual) { assertThat(expected).hasSameSizeAs(actual); for (int i = 0; i < expected.size(); i++) { assertThat(expected.get(i).getPath()).isEqualTo(actual.get(i).getPath()); assertThat(expected.get(i).getType()).isEqualTo(actual.get(i).getType()); if (Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.QUANTIZED_FLAT.toString()) || Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) { + assertThat(expected.get(i).getQuantizerType()).isEqualTo(actual.get(i).getQuantizerType()); assertThat(expected.get(i).getQuantizationSizeInBytes()).isEqualTo(actual.get(i).getQuantizationSizeInBytes()); assertThat(expected.get(i).getVectorIndexShardKeys()).isEqualTo(actual.get(i).getVectorIndexShardKeys()); } @@ -326,17 +344,26 @@ private List populateVectorIndexes() { CosmosVectorIndexSpec cosmosVectorIndexSpec2 = new CosmosVectorIndexSpec(); cosmosVectorIndexSpec2.setPath("/vector2"); cosmosVectorIndexSpec2.setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString()); + cosmosVectorIndexSpec2.setQuantizerType(QuantizerType.product); cosmosVectorIndexSpec2.setQuantizationSizeInBytes(2); cosmosVectorIndexSpec2.setVectorIndexShardKeys(Arrays.asList("/zipCode")); CosmosVectorIndexSpec cosmosVectorIndexSpec3 = new CosmosVectorIndexSpec(); cosmosVectorIndexSpec3.setPath("/vector3"); cosmosVectorIndexSpec3.setType(CosmosVectorIndexType.DISK_ANN.toString()); + cosmosVectorIndexSpec3.setQuantizerType(QuantizerType.product); cosmosVectorIndexSpec3.setQuantizationSizeInBytes(2); cosmosVectorIndexSpec3.setIndexingSearchListSize(30); cosmosVectorIndexSpec3.setVectorIndexShardKeys(Arrays.asList("/country/city")); - return Arrays.asList(cosmosVectorIndexSpec1, cosmosVectorIndexSpec2, cosmosVectorIndexSpec3); + CosmosVectorIndexSpec cosmosVectorIndexSpec4 = new CosmosVectorIndexSpec(); + cosmosVectorIndexSpec4.setPath("/vector4"); + cosmosVectorIndexSpec4.setType(CosmosVectorIndexType.DISK_ANN.toString()); + cosmosVectorIndexSpec4.setQuantizerType(QuantizerType.spherical); + cosmosVectorIndexSpec4.setIndexingSearchListSize(30); + cosmosVectorIndexSpec4.setVectorIndexShardKeys(Arrays.asList("/country/city")); + + return Arrays.asList(cosmosVectorIndexSpec1, cosmosVectorIndexSpec2, cosmosVectorIndexSpec3, cosmosVectorIndexSpec4); } private List populateEmbeddings() { @@ -357,14 +384,30 @@ private List populateEmbeddings() { embedding3.setDataType(CosmosVectorDataType.UINT8); embedding3.setEmbeddingDimensions(3); embedding3.setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); - return Arrays.asList(embedding1, embedding2, embedding3); + + CosmosVectorEmbedding embedding4 = new CosmosVectorEmbedding(); + embedding4.setPath("/vector4"); + embedding4.setDataType(CosmosVectorDataType.UINT8); + embedding4.setEmbeddingDimensions(3); + embedding4.setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); + return Arrays.asList(embedding1, embedding2, embedding3, embedding4); } private String getVectorEmbeddingPolicyAsString() { return "{\"vectorEmbeddings\":[" + "{\"path\":\"/vector1\",\"dataType\":\"int8\",\"dimensions\":3,\"distanceFunction\":\"cosine\"}," + "{\"path\":\"/vector2\",\"dataType\":\"float32\",\"dimensions\":3,\"distanceFunction\":\"dotproduct\"}," + - "{\"path\":\"/vector3\",\"dataType\":\"uint8\",\"dimensions\":3,\"distanceFunction\":\"euclidean\"}" + + "{\"path\":\"/vector3\",\"dataType\":\"uint8\",\"dimensions\":3,\"distanceFunction\":\"euclidean\"}," + + "{\"path\":\"/vector4\",\"dataType\":\"uint8\",\"dimensions\":3,\"distanceFunction\":\"euclidean\"}" + "]}"; } + + private String getVectorIndexesAsString() { + return "[" + + "{\"type\":\"flat\",\"path\":\"/vector1\"}," + + "{\"type\":\"quantizedFlat\",\"vectorIndexShardKeys\":[\"/zipCode\"],\"quantizerType\":\"product\",\"path\":\"/vector2\",\"quantizationByteSize\":2}," + + "{\"type\":\"diskANN\",\"indexingSearchListSize\":30,\"vectorIndexShardKeys\":[\"/country/city\"],\"quantizerType\":\"product\",\"path\":\"/vector3\",\"quantizationByteSize\":2}," + + "{\"type\":\"diskANN\",\"indexingSearchListSize\":30,\"vectorIndexShardKeys\":[\"/country/city\"],\"quantizerType\":\"spherical\",\"path\":\"/vector4\"}" + + "]"; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java index 3ea927bf9b72..340ef505f151 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java @@ -121,6 +121,7 @@ public static final class Properties { public static final String ORDER = "order"; public static final String SPATIAL_INDEXES = "spatialIndexes"; public static final String TYPES = "types"; + public static final String QUANTIZER_TYPE = "QuantizerType"; // Full text search public static final String FULL_TEXT_INDEXES = "fullTextIndexes"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java index 42f1ba663c16..c6f1adaa8a57 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java @@ -26,6 +26,8 @@ public final class CosmosVectorIndexSpec { private Integer indexingSearchListSize; @JsonInclude(JsonInclude.Include.NON_NULL) private List vectorIndexShardKeys; + @JsonInclude(JsonInclude.Include.NON_NULL) + private QuantizerType quantizerType; private final JsonSerializable jsonSerializable; /** @@ -84,6 +86,34 @@ public CosmosVectorIndexSpec setType(String type) { return this; } + /** + * Gets quantizer type. + * + * @return the quantizer type. + */ + public QuantizerType getQuantizerType() { + if (this.quantizerType == null) { + this.quantizerType = this.jsonSerializable.getObject(Constants.Properties.QUANTIZER_TYPE, QuantizerType.class); + } + return this.quantizerType; + } + + /** + * Set quantizer type. + * + * @param quantizerType The quantizer type + * @return the SpatialSpec. + */ + public CosmosVectorIndexSpec setQuantizerType(QuantizerType quantizerType) { + if (quantizerType != null) { + this.quantizerType = quantizerType; + this.jsonSerializable.set(Constants.Properties.QUANTIZER_TYPE, quantizerType); + } else { + this.quantizerType = null; + } + return this; + } + /** * Gets the quantization byte size * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java new file mode 100644 index 000000000000..eeebe27c4541 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.models; + +/** + * Defines the quantizer type of vector index path specification in the Azure Cosmos DB service. + */ +public enum QuantizerType { + /** + * Represent a product quantizer type. + */ + product("product"), + + /** + * Represent a spherical quantizer type. + */ + spherical("spherical"); + + + QuantizerType(String overWireValue) { + this.overWireValue = overWireValue; + } + + private final String overWireValue; + + @Override + public String toString() { + return this.overWireValue; + } +} + From 7d389e9be3c050ac70400555efe8ffdb096dd47b Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 14:28:52 -0800 Subject: [PATCH 2/8] Updated syntax of class initializations --- .../com/azure/cosmos/rx/VectorIndexTest.java | 127 +++++++++--------- 1 file changed, 64 insertions(+), 63 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index 196d2d275745..a29dbf1c66d0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -289,11 +289,11 @@ public void shouldValidateVectorIndexesSerializationAndDeserialization() throws // Validate Vector Indexes Serialization String actualVectorIndexesJSON = simpleObjectMapper.writeValueAsString(expectedVectorIndexes); String expectedVectorIndexesJSON = getVectorIndexesAsString(); - assertThat(expectedVectorIndexesJSON).isEqualTo(actualVectorIndexesJSON); + assertThat(actualVectorIndexesJSON).isEqualTo(expectedVectorIndexesJSON); // Validate Vector Indexes Deserialization List actualVectorIndexes = Arrays.asList(simpleObjectMapper.readValue(actualVectorIndexesJSON, CosmosVectorIndexSpec[].class)); - validateVectorIndexes(expectedVectorIndexes, actualVectorIndexes); + validateVectorIndexes(actualVectorIndexes, expectedVectorIndexes); } private void validateCollectionProperties(CosmosContainerProperties collectionDefinition, CosmosContainerProperties collectionProperties) { @@ -303,7 +303,7 @@ private void validateCollectionProperties(CosmosContainerProperties collectionDe collectionDefinition.getVectorEmbeddingPolicy()); assertThat(collectionProperties.getIndexingPolicy().getVectorIndexes()).isNotNull(); - validateVectorIndexes(collectionDefinition.getIndexingPolicy().getVectorIndexes(), collectionProperties.getIndexingPolicy().getVectorIndexes()); + validateVectorIndexes(collectionProperties.getIndexingPolicy().getVectorIndexes(), collectionDefinition.getIndexingPolicy().getVectorIndexes()); } private void validateVectorEmbeddingPolicy(CosmosVectorEmbeddingPolicy actual, CosmosVectorEmbeddingPolicy expected) { @@ -318,78 +318,79 @@ private void validateVectorEmbeddingPolicy(CosmosVectorEmbeddingPolicy actual, C } } - private void validateVectorIndexes(List expected, List actual) { - assertThat(expected).hasSameSizeAs(actual); - for (int i = 0; i < expected.size(); i++) { - assertThat(expected.get(i).getPath()).isEqualTo(actual.get(i).getPath()); - assertThat(expected.get(i).getType()).isEqualTo(actual.get(i).getType()); - if (Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.QUANTIZED_FLAT.toString()) || - Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) { - assertThat(expected.get(i).getQuantizerType()).isEqualTo(actual.get(i).getQuantizerType()); - assertThat(expected.get(i).getQuantizationSizeInBytes()).isEqualTo(actual.get(i).getQuantizationSizeInBytes()); - assertThat(expected.get(i).getVectorIndexShardKeys()).isEqualTo(actual.get(i).getVectorIndexShardKeys()); + private void validateVectorIndexes(List actual, List expected) { + assertThat(actual).hasSameSizeAs(expected); + for (int i = 0; i < actual.size(); i++) { + assertThat(actual.get(i).getPath()).isEqualTo(expected.get(i).getPath()); + assertThat(actual.get(i).getType()).isEqualTo(expected.get(i).getType()); + if (Objects.equals(actual.get(i).getType(), CosmosVectorIndexType.QUANTIZED_FLAT.toString()) || + Objects.equals(actual.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) { + assertThat(actual.get(i).getQuantizerType()).isEqualTo(expected.get(i).getQuantizerType()); + assertThat(actual.get(i).getQuantizationSizeInBytes()).isEqualTo(expected.get(i).getQuantizationSizeInBytes()); + assertThat(actual.get(i).getVectorIndexShardKeys()).isEqualTo(expected.get(i).getVectorIndexShardKeys()); } - if (Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) { - assertThat(expected.get(i).getIndexingSearchListSize()).isEqualTo(actual.get(i).getIndexingSearchListSize()); + if (Objects.equals(actual.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) { + assertThat(actual.get(i).getIndexingSearchListSize()).isEqualTo(expected.get(i).getIndexingSearchListSize()); } } } private List populateVectorIndexes() { - CosmosVectorIndexSpec cosmosVectorIndexSpec1 = new CosmosVectorIndexSpec(); - cosmosVectorIndexSpec1.setPath("/vector1"); - cosmosVectorIndexSpec1.setType(CosmosVectorIndexType.FLAT.toString()); - - CosmosVectorIndexSpec cosmosVectorIndexSpec2 = new CosmosVectorIndexSpec(); - cosmosVectorIndexSpec2.setPath("/vector2"); - cosmosVectorIndexSpec2.setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString()); - cosmosVectorIndexSpec2.setQuantizerType(QuantizerType.product); - cosmosVectorIndexSpec2.setQuantizationSizeInBytes(2); - cosmosVectorIndexSpec2.setVectorIndexShardKeys(Arrays.asList("/zipCode")); - - CosmosVectorIndexSpec cosmosVectorIndexSpec3 = new CosmosVectorIndexSpec(); - cosmosVectorIndexSpec3.setPath("/vector3"); - cosmosVectorIndexSpec3.setType(CosmosVectorIndexType.DISK_ANN.toString()); - cosmosVectorIndexSpec3.setQuantizerType(QuantizerType.product); - cosmosVectorIndexSpec3.setQuantizationSizeInBytes(2); - cosmosVectorIndexSpec3.setIndexingSearchListSize(30); - cosmosVectorIndexSpec3.setVectorIndexShardKeys(Arrays.asList("/country/city")); - - CosmosVectorIndexSpec cosmosVectorIndexSpec4 = new CosmosVectorIndexSpec(); - cosmosVectorIndexSpec4.setPath("/vector4"); - cosmosVectorIndexSpec4.setType(CosmosVectorIndexType.DISK_ANN.toString()); - cosmosVectorIndexSpec4.setQuantizerType(QuantizerType.spherical); - cosmosVectorIndexSpec4.setIndexingSearchListSize(30); - cosmosVectorIndexSpec4.setVectorIndexShardKeys(Arrays.asList("/country/city")); + CosmosVectorIndexSpec cosmosVectorIndexSpec1 = new CosmosVectorIndexSpec() + .setPath("/vector1") + .setType(CosmosVectorIndexType.FLAT.toString()); + + CosmosVectorIndexSpec cosmosVectorIndexSpec2 = new CosmosVectorIndexSpec() + .setPath("/vector2") + .setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString()) + .setQuantizerType(QuantizerType.product) + .setQuantizationSizeInBytes(2) + .setVectorIndexShardKeys(Arrays.asList("/zipCode")); + + CosmosVectorIndexSpec cosmosVectorIndexSpec3 = new CosmosVectorIndexSpec() + .setPath("/vector3") + .setType(CosmosVectorIndexType.DISK_ANN.toString()) + .setQuantizerType(QuantizerType.product) + .setQuantizationSizeInBytes(2) + .setIndexingSearchListSize(30) + .setVectorIndexShardKeys(Arrays.asList("/country/city")); + + CosmosVectorIndexSpec cosmosVectorIndexSpec4 = new CosmosVectorIndexSpec() + .setPath("/vector4") + .setType(CosmosVectorIndexType.DISK_ANN.toString()) + .setQuantizerType(QuantizerType.spherical) + .setIndexingSearchListSize(30) + .setVectorIndexShardKeys(Arrays.asList("/country/city")); return Arrays.asList(cosmosVectorIndexSpec1, cosmosVectorIndexSpec2, cosmosVectorIndexSpec3, cosmosVectorIndexSpec4); } private List populateEmbeddings() { - CosmosVectorEmbedding embedding1 = new CosmosVectorEmbedding(); - embedding1.setPath("/vector1"); - embedding1.setDataType(CosmosVectorDataType.INT8); - embedding1.setEmbeddingDimensions(3); - embedding1.setDistanceFunction(CosmosVectorDistanceFunction.COSINE); - - CosmosVectorEmbedding embedding2 = new CosmosVectorEmbedding(); - embedding2.setPath("/vector2"); - embedding2.setDataType(CosmosVectorDataType.FLOAT32); - embedding2.setEmbeddingDimensions(3); - embedding2.setDistanceFunction(CosmosVectorDistanceFunction.DOT_PRODUCT); - - CosmosVectorEmbedding embedding3 = new CosmosVectorEmbedding(); - embedding3.setPath("/vector3"); - embedding3.setDataType(CosmosVectorDataType.UINT8); - embedding3.setEmbeddingDimensions(3); - embedding3.setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); - - CosmosVectorEmbedding embedding4 = new CosmosVectorEmbedding(); - embedding4.setPath("/vector4"); - embedding4.setDataType(CosmosVectorDataType.UINT8); - embedding4.setEmbeddingDimensions(3); - embedding4.setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); + CosmosVectorEmbedding embedding1 = new CosmosVectorEmbedding() + .setPath("/vector1") + .setDataType(CosmosVectorDataType.INT8) + .setEmbeddingDimensions(3) + .setDistanceFunction(CosmosVectorDistanceFunction.COSINE); + + CosmosVectorEmbedding embedding2 = new CosmosVectorEmbedding() + .setPath("/vector2") + .setDataType(CosmosVectorDataType.FLOAT32) + .setEmbeddingDimensions(3) + .setDistanceFunction(CosmosVectorDistanceFunction.DOT_PRODUCT); + + CosmosVectorEmbedding embedding3 = new CosmosVectorEmbedding() + .setPath("/vector3") + .setDataType(CosmosVectorDataType.UINT8) + .setEmbeddingDimensions(3) + .setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); + + CosmosVectorEmbedding embedding4 = new CosmosVectorEmbedding() + .setPath("/vector4") + .setDataType(CosmosVectorDataType.UINT8) + .setEmbeddingDimensions(3) + .setDistanceFunction(CosmosVectorDistanceFunction.EUCLIDEAN); + return Arrays.asList(embedding1, embedding2, embedding3, embedding4); } From 2ad5623cf30aa691f59bda2ca8145616885892d4 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 14:53:28 -0800 Subject: [PATCH 3/8] Clean up test syntax to be more readable --- .../com/azure/cosmos/rx/VectorIndexTest.java | 29 +++++++++---------- .../models/CosmosVectorEmbeddingPolicy.java | 3 +- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index a29dbf1c66d0..ce6ecf490ab4 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -80,29 +80,26 @@ public void afterClass() { @Test(groups = {"emulator"}, timeOut = TIMEOUT*10000) public void shouldCreateVectorEmbeddingPolicy() { - PartitionKeyDefinition partitionKeyDef = new PartitionKeyDefinition(); - ArrayList paths = new ArrayList(); - paths.add("/mypk"); - partitionKeyDef.setPaths(paths); + ArrayList paths = new ArrayList<>(List.of("/mypk")); + PartitionKeyDefinition partitionKeyDef = new PartitionKeyDefinition() + .setPaths(paths); - CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(UUID.randomUUID().toString(), partitionKeyDef); - - IndexingPolicy indexingPolicy = new IndexingPolicy(); - indexingPolicy.setIndexingMode(IndexingMode.CONSISTENT); ExcludedPath excludedPath = new ExcludedPath("/*"); - indexingPolicy.setExcludedPaths(Collections.singletonList(excludedPath)); - IncludedPath includedPath1 = new IncludedPath("/name/?"); IncludedPath includedPath2 = new IncludedPath("/description/?"); - indexingPolicy.setIncludedPaths(ImmutableList.of(includedPath1, includedPath2)); - indexingPolicy.setVectorIndexes(populateVectorIndexes()); + IndexingPolicy indexingPolicy = new IndexingPolicy() + .setIndexingMode(IndexingMode.CONSISTENT) + .setExcludedPaths(Collections.singletonList(excludedPath)) + .setIncludedPaths(ImmutableList.of(includedPath1, includedPath2)) + .setVectorIndexes(populateVectorIndexes()); - CosmosVectorEmbeddingPolicy cosmosVectorEmbeddingPolicy = new CosmosVectorEmbeddingPolicy(); - cosmosVectorEmbeddingPolicy.setCosmosVectorEmbeddings(populateEmbeddings()); + CosmosVectorEmbeddingPolicy cosmosVectorEmbeddingPolicy = new CosmosVectorEmbeddingPolicy() + .setCosmosVectorEmbeddings(populateEmbeddings()); - collectionDefinition.setIndexingPolicy(indexingPolicy); - collectionDefinition.setVectorEmbeddingPolicy(cosmosVectorEmbeddingPolicy); + CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(UUID.randomUUID().toString(), partitionKeyDef) + .setIndexingPolicy(indexingPolicy) + .setVectorEmbeddingPolicy(cosmosVectorEmbeddingPolicy); database.createContainer(collectionDefinition).block(); CosmosAsyncContainer createdCollection = database.getContainer(collectionDefinition.getId()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java index 6abcc0028723..200ee2cd02c3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java @@ -44,11 +44,12 @@ public List getVectorEmbeddings() { * * @param cosmosVectorEmbeddings paths for embeddings along with path-specific settings for the item. */ - public void setCosmosVectorEmbeddings(List cosmosVectorEmbeddings) { + public CosmosVectorEmbeddingPolicy setCosmosVectorEmbeddings(List cosmosVectorEmbeddings) { cosmosVectorEmbeddings.forEach(embedding -> { checkNotNull(embedding, "Null values are not allowed in cosmosVectorEmbeddings list."); }); this.cosmosVectorEmbeddings = cosmosVectorEmbeddings; + return this; } } From 29bad2a9202cebc144e9b5f65b6d263aea166aae Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 14:53:54 -0800 Subject: [PATCH 4/8] Updated CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index ce3b76a5d722..628aa8b578c7 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.77.0-beta.1 (Unreleased) #### Features Added +* Added the `QuantizerType` to the vectorIndexSpec: `product`/`spherical`. - [PR 47566](https://github.com/Azure/azure-sdk-for-java/pull/47566) #### Breaking Changes From 61ec83a7d0539d75880c6e2911dde63de03bc14f Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 15:12:29 -0800 Subject: [PATCH 5/8] Added `QUANTIZER_TYPE` TO `IndexProperty` --- .../azure/cosmos/implementation/query/IndexProperty.java | 3 ++- .../java/com/azure/cosmos/models/CosmosVectorIndexSpec.java | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IndexProperty.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IndexProperty.java index e595d2212d80..1e1a4c1cca3d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IndexProperty.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IndexProperty.java @@ -5,5 +5,6 @@ public enum IndexProperty { INDEXING_SEARCH_LIST_SIZE, QUANTIZATION_SIZE_IN_BYTES, - VECTOR_INDEX_SHARD_KEYS; + VECTOR_INDEX_SHARD_KEYS, + QUANTIZER_TYPE, } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java index c6f1adaa8a57..193bf9faf0da 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java @@ -105,7 +105,7 @@ public QuantizerType getQuantizerType() { * @return the SpatialSpec. */ public CosmosVectorIndexSpec setQuantizerType(QuantizerType quantizerType) { - if (quantizerType != null) { + if (validateIndexType(IndexProperty.QUANTIZER_TYPE) && quantizerType != null) { this.quantizerType = quantizerType; this.jsonSerializable.set(Constants.Properties.QUANTIZER_TYPE, quantizerType); } else { @@ -223,7 +223,9 @@ JsonSerializable getJsonSerializable() { private Boolean validateIndexType(IndexProperty indexProperty) { String vectorIndexType = this.jsonSerializable.getString(Constants.Properties.VECTOR_INDEX_TYPE); - if (indexProperty.equals(IndexProperty.QUANTIZATION_SIZE_IN_BYTES) || (indexProperty.equals(IndexProperty.VECTOR_INDEX_SHARD_KEYS))) { + if (indexProperty.equals(IndexProperty.QUANTIZATION_SIZE_IN_BYTES) || + (indexProperty.equals(IndexProperty.VECTOR_INDEX_SHARD_KEYS)) || + (indexProperty.equals(IndexProperty.QUANTIZER_TYPE))) { return vectorIndexType.equals(CosmosVectorIndexType.QUANTIZED_FLAT.toString()) || vectorIndexType.equals(CosmosVectorIndexType.DISK_ANN.toString()); } From d2371518143cf4a3fb58d6a3ae1865d030bc7bde Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 15 Dec 2025 15:36:35 -0800 Subject: [PATCH 6/8] Addressed comments --- .../java/com/azure/cosmos/rx/VectorIndexTest.java | 6 +++--- .../com/azure/cosmos/implementation/Constants.java | 3 +-- .../azure/cosmos/models/CosmosVectorIndexSpec.java | 7 ++++--- .../java/com/azure/cosmos/models/QuantizerType.java | 13 ++++++++----- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index ce6ecf490ab4..9bd17afaebda 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -341,14 +341,14 @@ private List populateVectorIndexes() { CosmosVectorIndexSpec cosmosVectorIndexSpec2 = new CosmosVectorIndexSpec() .setPath("/vector2") .setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString()) - .setQuantizerType(QuantizerType.product) + .setQuantizerType(QuantizerType.PRODUCT) .setQuantizationSizeInBytes(2) .setVectorIndexShardKeys(Arrays.asList("/zipCode")); CosmosVectorIndexSpec cosmosVectorIndexSpec3 = new CosmosVectorIndexSpec() .setPath("/vector3") .setType(CosmosVectorIndexType.DISK_ANN.toString()) - .setQuantizerType(QuantizerType.product) + .setQuantizerType(QuantizerType.PRODUCT) .setQuantizationSizeInBytes(2) .setIndexingSearchListSize(30) .setVectorIndexShardKeys(Arrays.asList("/country/city")); @@ -356,7 +356,7 @@ private List populateVectorIndexes() { CosmosVectorIndexSpec cosmosVectorIndexSpec4 = new CosmosVectorIndexSpec() .setPath("/vector4") .setType(CosmosVectorIndexType.DISK_ANN.toString()) - .setQuantizerType(QuantizerType.spherical) + .setQuantizerType(QuantizerType.SPHERICAL) .setIndexingSearchListSize(30) .setVectorIndexShardKeys(Arrays.asList("/country/city")); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java index 340ef505f151..42d1d4a40bf2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java @@ -3,7 +3,6 @@ package com.azure.cosmos.implementation; -import java.util.List; /** * Used internally. Constants in the Azure Cosmos DB database service Java SDK. @@ -121,7 +120,7 @@ public static final class Properties { public static final String ORDER = "order"; public static final String SPATIAL_INDEXES = "spatialIndexes"; public static final String TYPES = "types"; - public static final String QUANTIZER_TYPE = "QuantizerType"; + public static final String QUANTIZER_TYPE = "quantizerType"; // Full text search public static final String FULL_TEXT_INDEXES = "fullTextIndexes"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java index 193bf9faf0da..870897595438 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java @@ -87,7 +87,7 @@ public CosmosVectorIndexSpec setType(String type) { } /** - * Gets quantizer type. + * Gets the quantizer type. * * @return the quantizer type. */ @@ -99,10 +99,10 @@ public QuantizerType getQuantizerType() { } /** - * Set quantizer type. + * Set the quantizer type. * * @param quantizerType The quantizer type - * @return the SpatialSpec. + * @return the CosmosVectorIndexSpec. */ public CosmosVectorIndexSpec setQuantizerType(QuantizerType quantizerType) { if (validateIndexType(IndexProperty.QUANTIZER_TYPE) && quantizerType != null) { @@ -110,6 +110,7 @@ public CosmosVectorIndexSpec setQuantizerType(QuantizerType quantizerType) { this.jsonSerializable.set(Constants.Properties.QUANTIZER_TYPE, quantizerType); } else { this.quantizerType = null; + this.jsonSerializable.remove(Constants.Properties.QUANTIZER_TYPE); } return this; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java index eeebe27c4541..03ca19ea893b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/QuantizerType.java @@ -3,19 +3,21 @@ package com.azure.cosmos.models; +import com.fasterxml.jackson.annotation.JsonValue; + /** - * Defines the quantizer type of vector index path specification in the Azure Cosmos DB service. + * Defines quantizer types for vector index specifications in the Azure Cosmos DB service. */ public enum QuantizerType { /** - * Represent a product quantizer type. + * Represents a product quantizer type. */ - product("product"), + PRODUCT("product"), /** - * Represent a spherical quantizer type. + * Represents a spherical quantizer type. */ - spherical("spherical"); + SPHERICAL("spherical"); QuantizerType(String overWireValue) { @@ -24,6 +26,7 @@ public enum QuantizerType { private final String overWireValue; + @JsonValue @Override public String toString() { return this.overWireValue; From 201916f52e5c6aaa113e5f8d84387ee7f73ce420 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Tue, 16 Dec 2025 14:14:32 -0800 Subject: [PATCH 7/8] Fix compile error with the usage of `List.of` --- .../src/test/java/com/azure/cosmos/rx/VectorIndexTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index 9bd17afaebda..12916398ddc8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -80,7 +80,7 @@ public void afterClass() { @Test(groups = {"emulator"}, timeOut = TIMEOUT*10000) public void shouldCreateVectorEmbeddingPolicy() { - ArrayList paths = new ArrayList<>(List.of("/mypk")); + ArrayList paths = new ArrayList<>(Arrays.asList("/mypk")); PartitionKeyDefinition partitionKeyDef = new PartitionKeyDefinition() .setPaths(paths); From 08ee4a4c5f07cc867ff728dbe50087eda048860c Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Wed, 17 Dec 2025 11:04:02 -0800 Subject: [PATCH 8/8] Revert return type change to avoid linting failure --- .../src/test/java/com/azure/cosmos/rx/VectorIndexTest.java | 4 ++-- .../com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java index 12916398ddc8..c2aef763329f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java @@ -94,8 +94,8 @@ public void shouldCreateVectorEmbeddingPolicy() { .setIncludedPaths(ImmutableList.of(includedPath1, includedPath2)) .setVectorIndexes(populateVectorIndexes()); - CosmosVectorEmbeddingPolicy cosmosVectorEmbeddingPolicy = new CosmosVectorEmbeddingPolicy() - .setCosmosVectorEmbeddings(populateEmbeddings()); + CosmosVectorEmbeddingPolicy cosmosVectorEmbeddingPolicy = new CosmosVectorEmbeddingPolicy(); + cosmosVectorEmbeddingPolicy.setCosmosVectorEmbeddings(populateEmbeddings()); CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(UUID.randomUUID().toString(), partitionKeyDef) .setIndexingPolicy(indexingPolicy) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java index 200ee2cd02c3..6abcc0028723 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbeddingPolicy.java @@ -44,12 +44,11 @@ public List getVectorEmbeddings() { * * @param cosmosVectorEmbeddings paths for embeddings along with path-specific settings for the item. */ - public CosmosVectorEmbeddingPolicy setCosmosVectorEmbeddings(List cosmosVectorEmbeddings) { + public void setCosmosVectorEmbeddings(List cosmosVectorEmbeddings) { cosmosVectorEmbeddings.forEach(embedding -> { checkNotNull(embedding, "Null values are not allowed in cosmosVectorEmbeddings list."); }); this.cosmosVectorEmbeddings = cosmosVectorEmbeddings; - return this; } }