[Cosmos][VectorIndex]Adding changes for vectorIndex and vectorEmbeddingPolicy

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,8 @@
 ### 4.58.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added public APIs `setMaxMicroBatchSize` and `getMaxMicroBatchSize` in `CosmosBulkExecutionOptions` - See [PR 39335](https://github.com/Azure/azure-sdk-for-java/pull/39335)
+* Added `vectorEmbeddingPolicy` in `cosmosContainerProperties` and `vectorIndexes` in `indexPolicy` to support vector search in CosmosDB - See[39379](https://github.com/Azure/azure-sdk-for-java/pull/39379)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java

@@ -120,6 +120,15 @@ public static final class Properties {
         public static final String SPATIAL_INDEXES = "spatialIndexes";
         public static final String TYPES = "types";
 
+        // Vector Embedding Policy
+        public static final String VECTOR_EMBEDDING_POLICY = "vectorEmbeddingPolicy";
+        public static final String VECTOR_INDEXES = "vectorIndexes";
+        public static final String VECTOR_EMBEDDINGS = "vectorEmbeddings";
+        public static final String VECTOR_INDEX_TYPE = "type";
+        public static final String VECTOR_DATA_TYPE = "dataType";
+        public static final String VECTOR_DIMENSIONS = "dimensions";
+        public static final String DISTANCE_FUNCTION = "distanceFunction";
+
         // Unique index.
         public static final String UNIQUE_KEY_POLICY = "uniqueKeyPolicy";
         public static final String UNIQUE_KEYS = "uniqueKeys";

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentCollection.java

@@ -13,6 +13,7 @@
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKeyDefinition;
 import com.azure.cosmos.models.UniqueKeyPolicy;
+import com.azure.cosmos.models.VectorEmbeddingPolicy;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.TextNode;
@@ -41,6 +42,7 @@ public final class DocumentCollection extends Resource {
     private UniqueKeyPolicy uniqueKeyPolicy;
     private PartitionKeyDefinition partitionKeyDefinition;
     private ClientEncryptionPolicy clientEncryptionPolicyInternal;
+    private VectorEmbeddingPolicy vectorEmbeddingPolicy;
 
     /**
      * Constructor.
@@ -411,6 +413,36 @@ public void setClientEncryptionPolicy(ClientEncryptionPolicy value) {
         setProperty(this, Constants.Properties.CLIENT_ENCRYPTION_POLICY, value);
     }
 
+    /**
+     * Gets the Vector Embedding Policy containing paths for embeddings along with path-specific settings for the item
+     * used in performing vector search on the items in a collection in the Azure CosmosDB database service.
+     *
+     * @return the Vector Embedding Policy.
+     */
+    public VectorEmbeddingPolicy getVectorEmbeddingPolicy() {
+        if (this.vectorEmbeddingPolicy == null) {
+            if (super.has(Constants.Properties.VECTOR_EMBEDDING_POLICY)) {
+                this.vectorEmbeddingPolicy = super.getObject(Constants.Properties.VECTOR_EMBEDDING_POLICY,
+                    VectorEmbeddingPolicy.class);
+            }
+        }
+        return this.vectorEmbeddingPolicy;
+    }
+
+    /**
+     * Sets the Vector Embedding Policy containing paths for embeddings along with path-specific settings for the item
+     * used in performing vector search on the items in a collection in the Azure CosmosDB database service.
+     *
+     * @param value the Vector Embedding Policy.
+     */
+    public void setVectorEmbeddingPolicy(VectorEmbeddingPolicy value) {
+        if (value == null) {
+            throw new IllegalArgumentException("VectorEmbeddingPolicy cannot be null.");
+        }
+        this.vectorEmbeddingPolicy = value;
+        setProperty(this, Constants.Properties.VECTOR_EMBEDDING_POLICY, value);
+    }
+
     public void populatePropertyBag() {
         super.populatePropertyBag();
         if (this.indexingPolicy == null) {

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

@@ -92,7 +92,7 @@ public CompositePathSortOrder getOrder() {
     }
 
     /**
-     * Gets the sort order for the composite path.
+     * Sets the sort order for the composite path.
      * <p>
      * For example if you want to run the query "SELECT * FROM c ORDER BY c.age asc, c.height desc",
      * then you need to make the order for "/age" "ascending" and the order for "/height" "descending".

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

@@ -347,6 +347,28 @@ public CosmosContainerProperties setClientEncryptionPolicy(ClientEncryptionPolic
         return this;
     }
 
+    /**
+     * Gets the Vector Embedding Policy containing paths for embeddings along with path-specific settings for the item
+     * used in performing vector search on the items in a collection in the Azure CosmosDB database service.
+     *
+     * @return the Vector Embedding Policy.
+     */
+    public VectorEmbeddingPolicy getVectorEmbeddingPolicy() {
+        return this.documentCollection.getVectorEmbeddingPolicy();
+    }
+
+    /**
+     * Sets the Vector Embedding Policy containing paths for embeddings along with path-specific settings for the item
+     * used in performing vector search on the items in a collection in the Azure CosmosDB database service.
+     *
+     * @param value the Vector Embedding Policy.
+     * @return the CosmosContainerProperties.
+     */
+    public CosmosContainerProperties setVectorEmbeddingPolicy(VectorEmbeddingPolicy value) {
+        this.documentCollection.setVectorEmbeddingPolicy(value);
+        return this;
+    }
+
     Resource getResource() {
         return this.documentCollection;
     }

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

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.models;
+
+/**
+ * Distance Function for the embeddings in the Cosmos DB database service.
+ */
+public enum DistanceFunction {
+    /**
+     * Represents the euclidean distance function.
+     */
+    EUCLIDEAN("euclidean"),
+
+    /**
+     * Represents the cosine distance function.
+     */
+    COSINE("cosine"),
+
+    /**
+     * Represents the dot product distance function.
+     */
+    DOT_PRODUCT("dotproduct");
+
+    private final String overWireValue;
+
+    DistanceFunction(String overWireValue) {
+        this.overWireValue = overWireValue;
+    }
+
+    @Override
+    public String toString() {
+        return this.overWireValue;
+    }
+
+    /**
+     *
+     * @return value for the enum
+     */
+    public String getValue() {
+        return this.overWireValue;
+    }
+
+    /**
+     *
+     * @return if the value for the enum is empty or not.
+     */
+    public boolean isEmpty() {
+        return this.overWireValue.isEmpty();
+    }
+}

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

@@ -0,0 +1,101 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.models;
+
+import com.azure.cosmos.implementation.Constants;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * Embedding settings within {@link VectorEmbeddingPolicy}
+ */
+public final class Embedding {
+    @JsonProperty(Constants.Properties.PATH)
+    private String path;
+    @JsonProperty(Constants.Properties.VECTOR_DATA_TYPE)
+    private String vectorDataType;
+    @JsonProperty(Constants.Properties.VECTOR_DIMENSIONS)
+    private Long dimensions;
+    @JsonProperty(Constants.Properties.DISTANCE_FUNCTION)
+    private String distanceFunction;
+
+    /**
+     * Gets the path for the embedding.
+     *
+     * @return path
+     */
+    public String getPath() {
+        return path;
+    }
+
+    /**
+     * Sets the path for the embedding.
+     *
+     * @param path the path for the embedding
+     * @return Embedding
+     */
+    public Embedding setPath(String path) {
+        this.path = path;
+        return this;
+    }
+
+    /**
+     * Gets the data type for the embedding.
+     *
+     * @return vectorDataType
+     */
+    public String getVectorDataType() {
+        return vectorDataType;
+    }
+
+    /**
+     * Sets the data type for the embedding.
+     *
+     * @param vectorDataType the data type for the embedding
+     * @return Embedding
+     */
+    public Embedding setVectorDataType(String vectorDataType) {
+        this.vectorDataType = vectorDataType;
+        return this;
+    }
+
+    /**
+     * Gets the dimensions for the embedding.
+     *
+     * @return dimensions
+     */
+    public Long getDimensions() {
+        return dimensions;
+    }
+
+    /**
+     * Sets the dimensions for the embedding.
+     *
+     * @param dimensions the dimensions for the embedding
+     * @return Embedding
+     */
+    public Embedding setDimensions(Long dimensions) {
+        this.dimensions = dimensions;
+        return this;
+    }
+
+    /**
+     * Gets the distanceFunction for the embedding.
+     *
+     * @return distanceFunction
+     */
+    public String getDistanceFunction() {
+        return distanceFunction;
+    }
+
+    /**
+     * Sets the distanceFunction for the embedding.
+     *
+     * @param distanceFunction the distanceFunction for the embedding
+     * @return Embedding
+     */
+    public Embedding setDistanceFunction(String distanceFunction) {
+        this.distanceFunction = distanceFunction;
+        return this;
+    }
+}

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

@@ -19,13 +19,12 @@
  */
 public final class IndexingPolicy {
     private static final String DEFAULT_PATH = "/*";
-
+    private final JsonSerializable jsonSerializable;
     private List<IncludedPath> includedPaths;
     private List<ExcludedPath> excludedPaths;
     private List<List<CompositePath>> compositeIndexes;
     private List<SpatialSpec> spatialIndexes;
-
-    private JsonSerializable jsonSerializable;
+    private List<VectorIndexSpec> vectorIndexes;
 
     /**
      * Constructor.
@@ -53,7 +52,7 @@ public IndexingPolicy() {
      * </pre>
      *
      * @param defaultIndexOverrides comma separated set of indexes that serve as default index specifications for the
-     * root path.
+     *                              root path.
      * @throws IllegalArgumentException throws when defaultIndexOverrides is null
      */
     IndexingPolicy(Index[] defaultIndexOverrides) {
@@ -234,7 +233,7 @@ public IndexingPolicy setCompositeIndexes(List<List<CompositePath>> compositeInd
     }
 
     /**
-     * Sets the spatial indexes for additional indexes.
+     * Gets the spatial indexes for additional indexes.
      *
      * @return the spatial indexes.
      */
@@ -262,11 +261,55 @@ public IndexingPolicy setSpatialIndexes(List<SpatialSpec> spatialIndexes) {
         return this;
     }
 
+    /**
+     * Gets the vector indexes.
+     *
+     * @return the vector indexes
+     */
+    public List<VectorIndexSpec> getVectorIndexes() {
+        if (this.vectorIndexes == null) {
+            this.vectorIndexes = this.jsonSerializable.getList(Constants.Properties.VECTOR_INDEXES, VectorIndexSpec.class);
+
+            if (this.vectorIndexes == null) {
+                this.vectorIndexes = new ArrayList<VectorIndexSpec>();
+            }
+        }
+
+        return this.vectorIndexes;
+    }
+
+    /**
+     * Sets the vector indexes.
+     *
+     * Example of the vectorIndexes:
+     * "vectorIndexes": [
+     *      {
+     *          "path": "/vector1",
+     *          "type": "diskANN"
+     *      },
+     *      {
+     *          "path": "/vector1",
+     *          "type": "flat"
+     *      },
+     *      {
+     *          "path": "/vector2",
+     *          "type": "quantizedFlat"
+     *      }]
+     *
+     * @param vectorIndexes the vector indexes
+     * @return the Indexing Policy.
+     */
+    public IndexingPolicy setVectorIndexes(List<VectorIndexSpec> vectorIndexes) {
+        this.vectorIndexes = vectorIndexes;
+        this.jsonSerializable.set(Constants.Properties.VECTOR_INDEXES,this.vectorIndexes);
+        return this;
+    }
+
     void populatePropertyBag() {
         this.jsonSerializable.populatePropertyBag();
         // If indexing mode is not 'none' and not paths are set, set them to the defaults
         if (this.getIndexingMode() != IndexingMode.NONE && this.getIncludedPaths().size() == 0
-                && this.getExcludedPaths().size() == 0) {
+            && this.getExcludedPaths().size() == 0) {
             IncludedPath includedPath = new IncludedPath(IndexingPolicy.DEFAULT_PATH);
             this.getIncludedPaths().add(includedPath);
         }
@@ -286,5 +329,7 @@ void populatePropertyBag() {
         }
     }
 
-    JsonSerializable getJsonSerializable() { return this.jsonSerializable; }
+    JsonSerializable getJsonSerializable() {
+        return this.jsonSerializable;
+    }
 }

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

@@ -598,6 +598,8 @@ public static <T> void populatePropertyBag(T t) {
             ((PartitionKeyDefinition) t).populatePropertyBag();
         } else if (t instanceof SpatialSpec) {
             ((SpatialSpec) t).populatePropertyBag();
+        } else if (t instanceof VectorIndexSpec) {
+            ((VectorIndexSpec) t).populatePropertyBag();
         } else if (t instanceof SqlParameter) {
             ((SqlParameter) t).populatePropertyBag();
         } else if (t instanceof SqlQuerySpec) {
@@ -631,6 +633,8 @@ public static <T> JsonSerializable getJsonSerializable(T t) {
             return ((PartitionKeyDefinition) t).getJsonSerializable();
         } else if (t instanceof SpatialSpec) {
             return ((SpatialSpec) t).getJsonSerializable();
+        } else if (t instanceof VectorIndexSpec) {
+            return ((VectorIndexSpec) t).getJsonSerializable();
         } else if (t instanceof SqlParameter) {
             return ((SqlParameter) t).getJsonSerializable();
         } else if (t instanceof SqlQuerySpec) {