feat(vertexai): support shard_size in google_vertex_ai_index (#8293) (#15133)

Signed-off-by: Modular Magician <[email protected]>

Author: modular-magician

.changelog/8293.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+vertexai: supported`shard_size` in `google_vertex_ai_index`
+```

google/resource_vertex_ai_index_generated_test.go

@@ -88,6 +88,7 @@ resource "google_vertex_ai_index" "index" {
     config {
       dimensions = 2
       approximate_neighbors_count = 150
+      shard_size = "SHARD_SIZE_SMALL"
       distance_measure_type = "DOT_PRODUCT_DISTANCE"
       algorithm_config {
         tree_ah_config {
@@ -158,6 +159,7 @@ resource "google_vertex_ai_index" "index" {
     contents_delta_uri = "gs://${google_storage_bucket.bucket.name}/contents"
     config {
       dimensions = 2
+      shard_size = "SHARD_SIZE_LARGE"
       distance_measure_type = "COSINE_DISTANCE"
       feature_norm_type = "UNIT_L2_NORM"
       algorithm_config {

google/resource_vertex_ai_index_test.go

@@ -91,6 +91,7 @@ resource "google_vertex_ai_index" "index" {
     config {
       dimensions = 2
       approximate_neighbors_count = 150
+      shard_size = "SHARD_SIZE_SMALL"
       distance_measure_type = "DOT_PRODUCT_DISTANCE"
       algorithm_config {
         tree_ah_config {
@@ -147,6 +148,7 @@ resource "google_vertex_ai_index" "index" {
     config {
       dimensions = 2
       approximate_neighbors_count = 150
+      shard_size = "SHARD_SIZE_SMALL"
       distance_measure_type = "DOT_PRODUCT_DISTANCE"
       algorithm_config {
         tree_ah_config {

google/services/vertexai/resource_vertex_ai_index.go

@@ -165,6 +165,17 @@ Required if tree-AH algorithm is used.`,
 * NONE: No normalization type is specified.`,
 										Default: "NONE",
 									},
+									"shard_size": {
+										Type:     schema.TypeString,
+										Computed: true,
+										Optional: true,
+										ForceNew: true,
+										Description: `Index data is split into equal parts to be processed. These are called "shards".
+The shard size must be specified when creating an index. The value must be one of the followings:
+* SHARD_SIZE_SMALL: Small (2GB)
+* SHARD_SIZE_MEDIUM: Medium (20GB)
+* SHARD_SIZE_LARGE: Large (50GB)`,
+									},
 								},
 							},
 						},
@@ -706,6 +717,8 @@ func flattenVertexAIIndexMetadataConfig(v interface{}, d *schema.ResourceData, c
 		flattenVertexAIIndexMetadataConfigDimensions(original["dimensions"], d, config)
 	transformed["approximate_neighbors_count"] =
 		flattenVertexAIIndexMetadataConfigApproximateNeighborsCount(original["approximateNeighborsCount"], d, config)
+	transformed["shard_size"] =
+		flattenVertexAIIndexMetadataConfigShardSize(original["shardSize"], d, config)
 	transformed["distance_measure_type"] =
 		flattenVertexAIIndexMetadataConfigDistanceMeasureType(original["distanceMeasureType"], d, config)
 	transformed["feature_norm_type"] =
@@ -748,6 +761,10 @@ func flattenVertexAIIndexMetadataConfigApproximateNeighborsCount(v interface{},
 	return v // let terraform core handle it otherwise
 }
 
+func flattenVertexAIIndexMetadataConfigShardSize(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
+	return v
+}
+
 func flattenVertexAIIndexMetadataConfigDistanceMeasureType(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
 	return v
 }
@@ -983,6 +1000,13 @@ func expandVertexAIIndexMetadataConfig(v interface{}, d tpgresource.TerraformRes
 		transformed["approximateNeighborsCount"] = transformedApproximateNeighborsCount
 	}
 
+	transformedShardSize, err := expandVertexAIIndexMetadataConfigShardSize(original["shard_size"], d, config)
+	if err != nil {
+		return nil, err
+	} else if val := reflect.ValueOf(transformedShardSize); val.IsValid() && !tpgresource.IsEmptyValue(val) {
+		transformed["shardSize"] = transformedShardSize
+	}
+
 	transformedDistanceMeasureType, err := expandVertexAIIndexMetadataConfigDistanceMeasureType(original["distance_measure_type"], d, config)
 	if err != nil {
 		return nil, err
@@ -1015,6 +1039,10 @@ func expandVertexAIIndexMetadataConfigApproximateNeighborsCount(v interface{}, d
 	return v, nil
 }
 
+func expandVertexAIIndexMetadataConfigShardSize(v interface{}, d tpgresource.TerraformResourceData, config *transport_tpg.Config) (interface{}, error) {
+	return v, nil
+}
+
 func expandVertexAIIndexMetadataConfigDistanceMeasureType(v interface{}, d tpgresource.TerraformResourceData, config *transport_tpg.Config) (interface{}, error) {
 	return v, nil
 }

website/docs/r/vertex_ai_index.html.markdown

@@ -59,6 +59,7 @@ resource "google_vertex_ai_index" "index" {
     config {
       dimensions = 2
       approximate_neighbors_count = 150
+      shard_size = "SHARD_SIZE_SMALL"
       distance_measure_type = "DOT_PRODUCT_DISTANCE"
       algorithm_config {
         tree_ah_config {
@@ -103,6 +104,7 @@ resource "google_vertex_ai_index" "index" {
     contents_delta_uri = "gs://${google_storage_bucket.bucket.name}/contents"
     config {
       dimensions = 2
+      shard_size = "SHARD_SIZE_LARGE"
       distance_measure_type = "COSINE_DISTANCE"
       feature_norm_type = "UNIT_L2_NORM"
       algorithm_config {
@@ -189,6 +191,14 @@ The following arguments are supported:
   approximate search algorithm are reordered via a more expensive distance computation.
   Required if tree-AH algorithm is used.
 
+* `shard_size` -
+  (Optional)
+  Index data is split into equal parts to be processed. These are called "shards".
+  The shard size must be specified when creating an index. The value must be one of the followings:
+  * SHARD_SIZE_SMALL: Small (2GB)
+  * SHARD_SIZE_MEDIUM: Medium (20GB)
+  * SHARD_SIZE_LARGE: Large (50GB)
+
 * `distance_measure_type` -
   (Optional)
   The distance measure used in nearest neighbor search. The value must be one of the followings: