Resolving comments

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java

@@ -12,6 +12,7 @@
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.DirectConnectionConfig;
 import com.azure.cosmos.implementation.TestConfigurations;
+import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosVectorDataType;
@@ -25,6 +26,8 @@
 import com.azure.cosmos.models.VectorEmbeddingPolicy;
 import com.azure.cosmos.models.VectorIndexSpec;
 import com.azure.cosmos.models.VectorIndexType;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.annotations.AfterClass;
@@ -48,6 +51,7 @@ public class VectorIndexTest extends TestSuiteBase {
     protected static final int SHUTDOWN_TIMEOUT = 20000;
 
     protected static Logger logger = LoggerFactory.getLogger(VectorIndexTest.class.getSimpleName());
+    private final ObjectMapper simpleObjectMapper = Utils.getSimpleObjectMapper();
     private final String databaseId = CosmosDatabaseForTest.generateId();
     private CosmosAsyncClient client;
     private CosmosAsyncDatabase database;
@@ -161,9 +165,9 @@ public void shouldFailOnWrongVectorEmbeddingPolicy() {
 
         CosmosVectorEmbedding embedding = new CosmosVectorEmbedding(
             "/vector1",
-            CosmosVectorDistanceFunction.COSINE.toString(),
+            "String",
             3L,
-            "String");
+            CosmosVectorDistanceFunction.COSINE.toString());
 
         try {
             VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
@@ -173,7 +177,7 @@ public void shouldFailOnWrongVectorEmbeddingPolicy() {
             assertThat(ex.getMessage()).isEqualTo("Invalid vector data type for the vector embedding policy.");
         }
 
-        embedding.setVectorDataType("");
+        embedding.setCosmosVectorDataType("");
         try {
             VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
             collectionDefinition.setVectorEmbeddingPolicy(vectorEmbeddingPolicy);
@@ -182,8 +186,8 @@ public void shouldFailOnWrongVectorEmbeddingPolicy() {
             assertThat(ex.getMessage()).isEqualTo("Vector data type cannot be empty for the vector embedding policy.");
         }
 
-        embedding.setVectorDataType(CosmosVectorDataType.FLOAT32.toString());
-        embedding.setDistanceFunction("COS");
+        embedding.setCosmosVectorDataType(CosmosVectorDataType.FLOAT32.toString());
+        embedding.setCosmosVectorDistanceFunction("COS");
         try {
             VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
             collectionDefinition.setVectorEmbeddingPolicy(vectorEmbeddingPolicy);
@@ -192,7 +196,7 @@ public void shouldFailOnWrongVectorEmbeddingPolicy() {
             assertThat(ex.getMessage()).isEqualTo("Invalid distance function for the vector embedding policy.");
         }
 
-        embedding.setDistanceFunction("");
+        embedding.setCosmosVectorDistanceFunction("");
         try {
             VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
             collectionDefinition.setVectorEmbeddingPolicy(vectorEmbeddingPolicy);
@@ -201,7 +205,7 @@ public void shouldFailOnWrongVectorEmbeddingPolicy() {
             assertThat(ex.getMessage()).isEqualTo("Distance function cannot be empty for the vector embedding policy.");
         }
 
-        embedding.setDistanceFunction(CosmosVectorDistanceFunction.COSINE.toString());
+        embedding.setCosmosVectorDistanceFunction(CosmosVectorDistanceFunction.COSINE.toString());
         embedding.setDimensions(-1L);
         try {
             VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
@@ -239,7 +243,7 @@ public void shouldFailOnWrongVectorIndex() {
             "/vector1",
             CosmosVectorDistanceFunction.COSINE.toString(),
             3L,
-            CosmosVectorDataType.INT8.toString());
+            CosmosVectorDistanceFunction.COSINE.toString());
         VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(ImmutableList.of(embedding));
         collectionDefinition.setVectorEmbeddingPolicy(vectorEmbeddingPolicy);
 
@@ -287,22 +291,47 @@ public void shouldCreateVectorIndexSimilarPathDifferentVectorType() {
         validateCollectionProperties(collectionDefinition, collectionProperties);
     }
 
+    @Test(groups = {"emulator"}, timeOut = TIMEOUT)
+    public void shouldValidateVectorEmbeddingPolicySerializationAndDeserialization() throws JsonProcessingException {
+        IndexingPolicy indexingPolicy = new IndexingPolicy();
+        indexingPolicy.setVectorIndexes(populateVectorIndexes());
+
+        VectorEmbeddingPolicy vectorEmbeddingPolicy = new VectorEmbeddingPolicy(populateEmbeddings());
+        String vectorEmbeddingPolicyJson = getVectorEmbeddingPolicyAsString();
+        String expectedVectorEmbeddingPolicyJson = simpleObjectMapper.writeValueAsString(vectorEmbeddingPolicy);
+        assertThat(vectorEmbeddingPolicyJson).isEqualTo(expectedVectorEmbeddingPolicyJson);
+
+        VectorEmbeddingPolicy expectedVectorEmbeddingPolicy = simpleObjectMapper.readValue(expectedVectorEmbeddingPolicyJson, VectorEmbeddingPolicy.class);
+        validateVectorEmbeddingPolicy(vectorEmbeddingPolicy, expectedVectorEmbeddingPolicy);
+    }
+
     private void validateCollectionProperties(CosmosContainerProperties collectionDefinition, CosmosContainerProperties collectionProperties) {
         assertThat(collectionProperties.getVectorEmbeddingPolicy()).isNotNull();
-        assertThat(collectionProperties.getVectorEmbeddingPolicy().getEmbeddings()).isNotNull();
-        List<CosmosVectorEmbedding> embeddings = collectionProperties.getVectorEmbeddingPolicy().getEmbeddings();
-        assertThat(embeddings).hasSameSizeAs(collectionDefinition.getVectorEmbeddingPolicy().getEmbeddings());
-        for (int i = 0; i < embeddings.size(); i++) {
-            assertThat(embeddings.get(0).getPath()).isEqualTo(
-                collectionDefinition.getVectorEmbeddingPolicy().getEmbeddings().get(0).getPath());
-        }
+        assertThat(collectionProperties.getVectorEmbeddingPolicy().getCosmosVectorEmbeddings()).isNotNull();
+        validateVectorEmbeddingPolicy(collectionProperties.getVectorEmbeddingPolicy(),
+            collectionDefinition.getVectorEmbeddingPolicy());
 
         assertThat(collectionProperties.getIndexingPolicy().getVectorIndexes()).isNotNull();
-        List<VectorIndexSpec> vectorIndexes = collectionProperties.getIndexingPolicy().getVectorIndexes();
-        assertThat(vectorIndexes).hasSameSizeAs(collectionDefinition.getIndexingPolicy().getVectorIndexes());
-        for (int i = 0; i < vectorIndexes.size(); i++) {
-            assertThat(vectorIndexes.get(0).getPath()).isEqualTo(
-                collectionDefinition.getIndexingPolicy().getVectorIndexes().get(0).getPath());
+        validateVectorIndexes(collectionDefinition.getIndexingPolicy().getVectorIndexes(), collectionProperties.getIndexingPolicy().getVectorIndexes());
+    }
+
+    private void validateVectorEmbeddingPolicy(VectorEmbeddingPolicy actual, VectorEmbeddingPolicy expected) {
+        List<CosmosVectorEmbedding> actualEmbeddings = actual.getCosmosVectorEmbeddings();
+        List<CosmosVectorEmbedding> expectedEmbeddings = expected.getCosmosVectorEmbeddings();
+        assertThat(expectedEmbeddings).hasSameSizeAs(actualEmbeddings);
+        for (int i = 0; i < expectedEmbeddings.size(); i++) {
+            assertThat(expectedEmbeddings.get(i).getPath()).isEqualTo(actualEmbeddings.get(i).getPath());
+            assertThat(expectedEmbeddings.get(i).getCosmosVectorDataType()).isEqualTo(actualEmbeddings.get(i).getCosmosVectorDataType());
+            assertThat(expectedEmbeddings.get(i).getDimensions()).isEqualTo(actualEmbeddings.get(i).getDimensions());
+            assertThat(expectedEmbeddings.get(i).getCosmosVectorDistanceFunction()).isEqualTo(actualEmbeddings.get(i).getCosmosVectorDistanceFunction());
+        }
+    }
+
+    private void validateVectorIndexes(List<VectorIndexSpec> actual, List<VectorIndexSpec> expected) {
+        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());
         }
     }
 
@@ -322,21 +351,29 @@ private List<VectorIndexSpec> populateVectorIndexes() {
     private List<CosmosVectorEmbedding> populateEmbeddings() {
         CosmosVectorEmbedding embedding1 = new CosmosVectorEmbedding(
             "/vector1",
-            CosmosVectorDistanceFunction.COSINE.toString(),
+            CosmosVectorDataType.FLOAT32.toString(),
             3L,
-            CosmosVectorDataType.FLOAT32.toString());
+            CosmosVectorDistanceFunction.COSINE.toString());
 
         CosmosVectorEmbedding embedding2 = new CosmosVectorEmbedding(
             "/vector2",
-            CosmosVectorDistanceFunction.DOT_PRODUCT.toString(),
+            CosmosVectorDataType.INT8.toString(),
             3L,
-            CosmosVectorDataType.INT8.toString());
+            CosmosVectorDistanceFunction.DOT_PRODUCT.toString());
 
         CosmosVectorEmbedding embedding3 = new CosmosVectorEmbedding(
             "/vector3",
-            CosmosVectorDistanceFunction.EUCLIDEAN.toString(),
+            CosmosVectorDataType.UINT8.toString(),
             3L,
-            CosmosVectorDataType.UINT8.toString());
+            CosmosVectorDistanceFunction.EUCLIDEAN.toString());
         return Arrays.asList(embedding1, embedding2, embedding3);
     }
+
+    private String getVectorEmbeddingPolicyAsString() {
+        return "{\"vectorEmbeddings\":[" +
+            "{\"path\":\"/vector1\",\"dataType\":\"float32\",\"dimensions\":3,\"distanceFunction\":\"cosine\"}," +
+            "{\"path\":\"/vector2\",\"dataType\":\"int8\",\"dimensions\":3,\"distanceFunction\":\"dotproduct\"}," +
+            "{\"path\":\"/vector3\",\"dataType\":\"uint8\",\"dimensions\":3,\"distanceFunction\":\"euclidean\"}" +
+            "]}";
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorEmbedding.java

@@ -4,6 +4,7 @@
 package com.azure.cosmos.models;
 
 import com.azure.cosmos.implementation.Constants;
+import com.azure.cosmos.implementation.JsonSerializable;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 /**
@@ -13,11 +14,12 @@ public final class CosmosVectorEmbedding {
     @JsonProperty(Constants.Properties.PATH)
     private String path;
     @JsonProperty(Constants.Properties.VECTOR_DATA_TYPE)
-    private String vectorDataType;
+    private String cosmosVectorDataType;
     @JsonProperty(Constants.Properties.VECTOR_DIMENSIONS)
     private Long dimensions;
     @JsonProperty(Constants.Properties.DISTANCE_FUNCTION)
-    private String distanceFunction;
+    private String cosmosVectorDistanceFunction;
+    private JsonSerializable jsonSerializable;
 
     /**
      * Constructor
@@ -29,9 +31,16 @@ public final class CosmosVectorEmbedding {
      */
     public CosmosVectorEmbedding(String path, String vectorDataType, Long dimensions, String distanceFunction) {
         this.path = path;
-        this.vectorDataType = vectorDataType;
+        this.cosmosVectorDataType = vectorDataType;
         this.dimensions = dimensions;
-        this.distanceFunction = distanceFunction;
+        this.cosmosVectorDistanceFunction = distanceFunction;
+    }
+
+    /**
+     * Constructor
+     */
+    public CosmosVectorEmbedding() {
+        this.jsonSerializable = new JsonSerializable();
     }
 
     /**
@@ -57,20 +66,20 @@ public CosmosVectorEmbedding setPath(String path) {
     /**
      * Gets the data type for the cosmosVectorEmbedding.
      *
-     * @return vectorDataType
+     * @return cosmosVectorDataType
      */
-    public String getVectorDataType() {
-        return vectorDataType;
+    public String getCosmosVectorDataType() {
+        return cosmosVectorDataType;
     }
 
     /**
      * Sets the data type for the cosmosVectorEmbedding.
      *
-     * @param vectorDataType the data type for the cosmosVectorEmbedding
+     * @param cosmosVectorDataType the data type for the cosmosVectorEmbedding
      * @return CosmosVectorEmbedding
      */
-    public CosmosVectorEmbedding setVectorDataType(String vectorDataType) {
-        this.vectorDataType = vectorDataType;
+    public CosmosVectorEmbedding setCosmosVectorDataType(String cosmosVectorDataType) {
+        this.cosmosVectorDataType = cosmosVectorDataType;
         return this;
     }
 
@@ -97,20 +106,20 @@ public CosmosVectorEmbedding setDimensions(Long dimensions) {
     /**
      * Gets the distanceFunction for the cosmosVectorEmbedding.
      *
-     * @return distanceFunction
+     * @return cosmosVectorDistanceFunction
      */
-    public String getDistanceFunction() {
-        return distanceFunction;
+    public String getCosmosVectorDistanceFunction() {
+        return cosmosVectorDistanceFunction;
     }
 
     /**
      * Sets the distanceFunction for the cosmosVectorEmbedding.
      *
-     * @param distanceFunction the distanceFunction for the cosmosVectorEmbedding
+     * @param cosmosVectorDistanceFunction the distanceFunction for the cosmosVectorEmbedding
      * @return CosmosVectorEmbedding
      */
-    public CosmosVectorEmbedding setDistanceFunction(String distanceFunction) {
-        this.distanceFunction = distanceFunction;
+    public CosmosVectorEmbedding setCosmosVectorDistanceFunction(String cosmosVectorDistanceFunction) {
+        this.cosmosVectorDistanceFunction = cosmosVectorDistanceFunction;
         return this;
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/VectorEmbeddingPolicy.java

@@ -22,16 +22,16 @@ public final class VectorEmbeddingPolicy {
      * Paths for embeddings along with path-specific settings for the item.
      */
     @JsonProperty(Constants.Properties.VECTOR_EMBEDDINGS)
-    private List<CosmosVectorEmbedding> embeddings;
+    private List<CosmosVectorEmbedding> cosmosVectorEmbeddings;
 
     /**
      * Constructor
      *
-     * @param embeddings list of path for embeddings along with path-specific settings for the item.
+     * @param cosmosVectorEmbeddings list of path for embeddings along with path-specific settings for the item.
      */
-    public VectorEmbeddingPolicy(List<CosmosVectorEmbedding> embeddings) {
-        validateEmbeddings(embeddings);
-        this.embeddings = embeddings;
+    public VectorEmbeddingPolicy(List<CosmosVectorEmbedding> cosmosVectorEmbeddings) {
+        validateEmbeddings(cosmosVectorEmbeddings);
+        this.cosmosVectorEmbeddings = cosmosVectorEmbeddings;
     }
 
     /**
@@ -41,15 +41,15 @@ public VectorEmbeddingPolicy() {
         this.jsonSerializable = new JsonSerializable();
     }
 
-    private static void validateEmbeddings(List<CosmosVectorEmbedding> embeddings) {
-        embeddings.forEach(embedding -> {
+    private static void validateEmbeddings(List<CosmosVectorEmbedding> cosmosVectorEmbeddings) {
+        cosmosVectorEmbeddings.forEach(embedding -> {
             if (embedding == null) {
                 throw new IllegalArgumentException("Embedding cannot be null.");
             }
             validateEmbeddingPath(embedding.getPath());
             validateEmbeddingDimensions(embedding.getDimensions());
-            validateEmbeddingVectorDataType(embedding.getVectorDataType());
-            validateEmbeddingDistanceFunction(embedding.getDistanceFunction());
+            validateEmbeddingVectorDataType(embedding.getCosmosVectorDataType());
+            validateEmbeddingDistanceFunction(embedding.getCosmosVectorDistanceFunction());
         });
     }
 
@@ -93,7 +93,7 @@ private static void validateEmbeddingDistanceFunction(String value) {
      *
      * @return the paths for embeddings along with path-specific settings for the item.
      */
-    public List<CosmosVectorEmbedding> getEmbeddings() {
-        return this.embeddings;
+    public List<CosmosVectorEmbedding> getCosmosVectorEmbeddings() {
+        return this.cosmosVectorEmbeddings;
     }
 }