addding shardSize parameter to MatchingEngineIndexConfig

google/cloud/aiplatform/matching_engine/matching_engine_index.py

@@ -408,6 +408,7 @@ def create_tree_ah_index(
         approximate_neighbors_count: int = None,
         leaf_node_embedding_count: Optional[int] = None,
         leaf_nodes_to_search_percent: Optional[float] = None,
+        shard_type: Optional[matching_engine_index_config.ShardSizeType] = None,
         distance_measure_type: Optional[
             matching_engine_index_config.DistanceMeasureType
         ] = None,
@@ -461,6 +462,10 @@ def create_tree_ah_index(
             leaf_nodes_to_search_percent (float):
                 Optional. The default percentage of leaf nodes that any query may be searched. Must be in
                 range 1-100, inclusive. The default value is 10 (means 10%) if not set.
+            shard_type (ShardSizeType):
+                Optional. When you create an index, you must specify the size of the shards to use. The machine types
+                that you can use to deploy your index depends on the shard size of the index. Default value is
+                SHARD_SIZE_MEDIUM.
             distance_measure_type (matching_engine_index_config.DistanceMeasureType):
                 Optional. The distance measure used in nearest neighbor search.
             description (str):
@@ -523,6 +528,7 @@ def create_tree_ah_index(
 
         config = matching_engine_index_config.MatchingEngineIndexConfig(
             dimensions=dimensions,
+            shardSize=shard_type,
             algorithm_config=algorithm_config,
             approximate_neighbors_count=approximate_neighbors_count,
             distance_measure_type=distance_measure_type,

google/cloud/aiplatform/matching_engine/matching_engine_index_config.py

@@ -35,6 +35,16 @@ class DistanceMeasureType(enum.Enum):
     COSINE_DISTANCE = "COSINE_DISTANCE"
 
 
+# Index data is split into equal parts called shards for processing. When you create an index, you must specify the size of the shards to use.
+class ShardSizeType(enum.Enum):
+    # 2 GiB per shard
+    SHARD_SIZE_SMALL = "SHARD_SIZE_SMALL"
+    # 20 GiB per shard
+    SHARD_SIZE_MEDIUM = "SHARD_SIZE_MEDIUM"
+    # 50 GiB per shard
+    SHARD_SIZE_LARGE = "SHARD_SIZE_LARGE"
+
+
 class FeatureNormType(enum.Enum):
     """Type of normalization to be carried out on each vector."""
 
@@ -112,6 +122,10 @@ class MatchingEngineIndexConfig:
             Required. The number of dimensions of the input vectors.
         algorithm_config (AlgorithmConfig):
             Required. The configuration with regard to the algorithms used for efficient search.
+        shardSize (ShardSizeType):
+            Optional. When you create an index, you must specify the size of the shards to use. The machine types
+            that you can use to deploy your index depends on the shard size of the index. Default value is
+            SHARD_SIZE_MEDIUM.
         approximate_neighbors_count (int):
             Optional. The default number of neighbors to find via approximate search before exact reordering is
             performed. Exact reordering is a procedure where results returned by an
@@ -124,6 +138,7 @@ class MatchingEngineIndexConfig:
 
     dimensions: int
     algorithm_config: AlgorithmConfig
+    shardSize: ShardSizeType = ShardSizeType.SHARD_SIZE_MEDIUM
     approximate_neighbors_count: Optional[int] = None
     distance_measure_type: Optional[DistanceMeasureType] = None
 
@@ -136,6 +151,7 @@ def as_dict(self) -> Dict[str, Any]:
 
         return {
             "dimensions": self.dimensions,
+            "shardSize": self.shardSize,
             "algorithmConfig": self.algorithm_config.as_dict(),
             "approximateNeighborsCount": self.approximate_neighbors_count,
             "distanceMeasureType": self.distance_measure_type,