fix: construct the request directly in sample code

samples/snippets/feature_store_service/batch_create_features_sample.py

@@ -20,7 +20,6 @@ def batch_create_features_sample(
     project: str,
     featurestore_id: str,
     entity_type_id: str,
-    requests: list,
     location: str = "us-central1",
     api_endpoint: str = "us-central1-aiplatform.googleapis.com",
     timeout: int = 300,
@@ -31,6 +30,33 @@ def batch_create_features_sample(
     # This client only needs to be created once, and can be reused for multiple requests.
     client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
     parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
+    requests = [
+        aiplatform.CreateFeatureRequest(
+            feature=aiplatform.Feature(
+                value_type=aiplatform.Feature.ValueType.INT64, description="User age",
+            ),
+            feature_id="age",
+        ),
+        aiplatform.CreateFeatureRequest(
+            feature=aiplatform.Feature(
+                value_type=aiplatform.Feature.ValueType.STRING,
+                description="User gender",
+                monitoring_config=aiplatform.FeaturestoreMonitoringConfig(
+                    snapshot_analysis=aiplatform.FeaturestoreMonitoringConfig.SnapshotAnalysis(
+                        disabled=True,
+                    ),
+                ),
+            ),
+            feature_id="gender",
+        ),
+        aiplatform.CreateFeatureRequest(
+            feature=aiplatform.Feature(
+                value_type=aiplatform.Feature.ValueType.STRING_ARRAY,
+                description="An array of genres that this user liked",
+            ),
+            feature_id="liked_genres",
+        ),
+    ]
     batch_create_features_request = aiplatform.BatchCreateFeaturesRequest(
         parent=parent, requests=requests
     )

samples/snippets/feature_store_service/batch_create_features_sample_test.py

@@ -16,49 +16,43 @@
 from uuid import uuid4
 
 import batch_create_features_sample
-from google.cloud import aiplatform_v1beta1 as aiplatform
+import create_entity_type_sample
 
 import pytest
 
+import helpers
+
 PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT")
 
 
 @pytest.fixture(scope="function", autouse=True)
-def teardown(teardown_features):
+def teardown(teardown_entity_type):
     yield
 
 
+def setup_temp_entity_type(featurestore_id, entity_type_id, capsys):
+    create_entity_type_sample.create_entity_type_sample(
+        project=PROJECT_ID,
+        featurestore_id=featurestore_id,
+        entity_type_id=entity_type_id,
+    )
+    out, _ = capsys.readouterr()
+    assert "create_entity_type_response" in out
+    return helpers.get_featurestore_resource_name(out)
+
+
 def test_ucaip_generated_batch_create_features_sample_vision(capsys, shared_state):
     featurestore_id = "perm_sample_featurestore"
-    entity_type_id = "perm_sample_entity_type"
-    requests = [
-        aiplatform.CreateFeatureRequest(
-            feature=aiplatform.Feature(
-                value_type=aiplatform.Feature.ValueType.STRING,
-                description="lorem ipsum",
-            ),
-            feature_id=f"gender_{uuid4()}".replace("-", "_")[:60],
-        ),
-        aiplatform.CreateFeatureRequest(
-            feature=aiplatform.Feature(
-                value_type=aiplatform.Feature.ValueType.STRING_ARRAY,
-                description="lorem ipsum",
-            ),
-            feature_id=f"liked_genres_{uuid4()}".replace("-", "_")[:60],
-        ),
-    ]
+    entity_type_id = f"users_{uuid4()}".replace("-", "_")[:60]
+    entity_type_name = setup_temp_entity_type(featurestore_id, entity_type_id, capsys)
     location = "us-central1"
     batch_create_features_sample.batch_create_features_sample(
         project=PROJECT_ID,
         featurestore_id=featurestore_id,
         entity_type_id=entity_type_id,
-        requests=requests,
         location=location,
     )
     out, _ = capsys.readouterr()
     assert "batch_create_features_response" in out
 
-    parent = f"projects/{PROJECT_ID}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}/features/"
-    shared_state["feature_names"] = []
-    for request in requests:
-        shared_state["feature_names"].append(parent + request.feature_id)
+    shared_state["entity_type_name"] = entity_type_name

samples/snippets/feature_store_service/batch_read_feature_values_sample.py

@@ -19,9 +19,8 @@
 def batch_read_feature_values_sample(
     project: str,
     featurestore_id: str,
-    csv_read_instances: aiplatform.CsvSource,
-    destination: aiplatform.FeatureValueDestination,
-    entity_type_specs: list,
+    input_csv_file: str,
+    destination_table_uri: str,
     location: str = "us-central1",
     api_endpoint: str = "us-central1-aiplatform.googleapis.com",
     timeout: int = 300,
@@ -34,6 +33,38 @@ def batch_read_feature_values_sample(
     featurestore = (
         f"projects/{project}/locations/{location}/featurestores/{featurestore_id}"
     )
+    csv_read_instances = aiplatform.CsvSource(
+        gcs_source=aiplatform.GcsSource(uris=[input_csv_file])
+    )
+    destination = aiplatform.FeatureValueDestination(
+        bigquery_destination=aiplatform.BigQueryDestination(
+            # Output to BigQuery table created earlier
+            output_uri=destination_table_uri
+        )
+    )
+    entity_type_specs = [
+        aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
+            # Read the 'age', 'gender' and 'liked_genres' features from the 'perm_users' entity
+            entity_type_id="perm_users",
+            feature_selector=aiplatform.FeatureSelector(
+                id_matcher=aiplatform.IdMatcher(
+                    ids=[
+                        # features, use "*" if you want to select all features within this entity type
+                        "age",
+                        "gender",
+                        "liked_genres",
+                    ]
+                )
+            ),
+        ),
+        aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
+            # Read the 'average_rating' and 'genres' features from the 'perm_movies' entity
+            entity_type_id="perm_movies",
+            feature_selector=aiplatform.FeatureSelector(
+                id_matcher=aiplatform.IdMatcher(ids=["average_rating", "genres"])
+            ),
+        ),
+    ]
     # Batch serving request from CSV
     batch_read_feature_values_request = aiplatform.BatchReadFeatureValuesRequest(
         featurestore=featurestore,

samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py

@@ -16,7 +16,6 @@
 import os
 
 import batch_read_feature_values_sample
-from google.cloud import aiplatform_v1beta1 as aiplatform
 from google.cloud import bigquery
 
 import pytest
@@ -36,14 +35,12 @@ def setup_test():
     destination_data_set = "movie_predictions_" + datetime.now().strftime(
         "%Y%m%d%H%M%S"
     )
-
     # Output table. Make sure that the table does NOT already exist; the BatchReadFeatureValues API cannot overwrite an existing table
     destination_table_name = "training_data"
     DESTINATION_PATTERN = "bq://{project}.{dataset}.{table}"
     destination_table_uri = DESTINATION_PATTERN.format(
         project=PROJECT_ID, dataset=destination_data_set, table=destination_table_name
     )
-
     # Create dataset
     bq_client = bigquery.Client(project=PROJECT_ID)
     dataset_id = "{}.{}".format(bq_client.project, destination_data_set)
@@ -56,47 +53,13 @@ def setup_test():
 
 def test_ucaip_generated_batch_read_feature_values_sample_vision(capsys, shared_state):
     destination_data_set, destination_table_uri = setup_test()
-
     featurestore_id = "perm_sample_featurestore"
-    csv_read_instances = aiplatform.CsvSource(
-        gcs_source=aiplatform.GcsSource(uris=[INPUT_CSV_FILE])
-    )
-    destination = aiplatform.FeatureValueDestination(
-        bigquery_destination=aiplatform.BigQueryDestination(
-            # Output to BigQuery table created earlier
-            output_uri=destination_table_uri
-        )
-    )
-    entity_type_specs = [
-        aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
-            # Read the 'age', 'gender' and 'liked_genres' features from the 'perm_users' entity
-            entity_type_id="perm_users",
-            feature_selector=aiplatform.FeatureSelector(
-                id_matcher=aiplatform.IdMatcher(
-                    ids=[
-                        # features, use "*" if you want to select all features within this entity type
-                        "age",
-                        "gender",
-                        "liked_genres",
-                    ]
-                )
-            ),
-        ),
-        aiplatform.BatchReadFeatureValuesRequest.EntityTypeSpec(
-            # Read the 'average_rating' and 'genres' features from the 'perm_movies' entity
-            entity_type_id="perm_movies",
-            feature_selector=aiplatform.FeatureSelector(
-                id_matcher=aiplatform.IdMatcher(ids=["average_rating", "genres"])
-            ),
-        ),
-    ]
 
     batch_read_feature_values_sample.batch_read_feature_values_sample(
         project=PROJECT_ID,
         featurestore_id=featurestore_id,
-        csv_read_instances=csv_read_instances,
-        destination=destination,
-        entity_type_specs=entity_type_specs,
+        input_csv_file=INPUT_CSV_FILE,
+        destination_table_uri=destination_table_uri,
     )
     out, _ = capsys.readouterr()
     assert "batch_read_feature_values_response" in out

samples/snippets/feature_store_service/import_feature_values_sample.py

@@ -20,8 +20,7 @@ def import_feature_values_sample(
     project: str,
     featurestore_id: str,
     entity_type_id: str,
-    avro_source: aiplatform.AvroSource,
-    feature_specs: list,
+    avro_gcs_uri: str,
     entity_id_field: str,
     feature_time_field: str,
     worker_count: int = 2,
@@ -35,6 +34,14 @@ def import_feature_values_sample(
     # This client only needs to be created once, and can be reused for multiple requests.
     client = aiplatform.FeaturestoreServiceClient(client_options=client_options)
     entity_type = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
+    avro_source = aiplatform.AvroSource(
+        gcs_source=aiplatform.GcsSource(uris=[avro_gcs_uri])
+    )
+    feature_specs = [
+        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="age"),
+        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="gender"),
+        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="liked_genres"),
+    ]
     import_feature_values_request = aiplatform.ImportFeatureValuesRequest(
         entity_type=entity_type,
         avro_source=avro_source,

samples/snippets/feature_store_service/import_feature_values_sample_test.py

@@ -17,7 +17,6 @@
 
 import batch_create_features_sample
 import create_entity_type_sample
-from google.cloud import aiplatform_v1beta1 as aiplatform
 import import_feature_values_sample
 import pytest
 
@@ -35,30 +34,10 @@ def teardown(teardown_entity_type):
 
 
 def setup_features(featurestore_id, entity_type_id, capsys):
-    requests = [
-        {
-            "feature_id": "age",
-            "feature": {"value_type": "INT64", "description": "User age"},
-        },
-        {
-            "feature_id": "gender",
-            "feature": {"value_type": "STRING", "description": "User gender"},
-        },
-        {
-            "feature_id": "liked_genres",
-            "feature": {
-                "value_type": "STRING_ARRAY",
-                "description": "An array of genres that this user liked",
-            },
-        },
-    ]
-    location = "us-central1"
     batch_create_features_sample.batch_create_features_sample(
         project=PROJECT_ID,
         featurestore_id=featurestore_id,
         entity_type_id=entity_type_id,
-        requests=requests,
-        location=location,
     )
     out, _ = capsys.readouterr()
     assert "batch_create_features_response" in out
@@ -81,20 +60,11 @@ def test_ucaip_generated_import_feature_values_sample_vision(capsys, shared_stat
     entity_type_name = setup_temp_entity_type(featurestore_id, entity_type_id, capsys)
     setup_features(featurestore_id, entity_type_id, capsys)
 
-    avro_source = aiplatform.AvroSource(
-        gcs_source=aiplatform.GcsSource(uris=[AVRO_GCS_URI])
-    )
-    feature_specs = [
-        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="age"),
-        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="gender"),
-        aiplatform.ImportFeatureValuesRequest.FeatureSpec(id="liked_genres"),
-    ]
     import_feature_values_sample.import_feature_values_sample(
         project=PROJECT_ID,
         featurestore_id=featurestore_id,
         entity_type_id=entity_type_id,
-        avro_source=avro_source,
-        feature_specs=feature_specs,
+        avro_gcs_uri=AVRO_GCS_URI,
         entity_id_field="user_id",
         feature_time_field="update_time",
         worker_count=2,

samples/snippets/feature_store_service/read_feature_values_sample.py

@@ -21,10 +21,8 @@ def read_feature_values_sample(
     featurestore_id: str,
     entity_type_id: str,
     entity_id: str,
-    feature_selector: str,
     location: str = "us-central1",
     api_endpoint: str = "us-central1-aiplatform.googleapis.com",
-    timeout: int = 300,
 ):
     # The AI Platform services require regional API endpoints.
     client_options = {"api_endpoint": api_endpoint}
@@ -34,6 +32,9 @@ def read_feature_values_sample(
         client_options=client_options
     )
     entity_type = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}"
+    feature_selector = aiplatform.FeatureSelector(
+        id_matcher=aiplatform.IdMatcher(ids=["age", "gender", "liked_genres"])
+    )
     read_feature_values_request = aiplatform.ReadFeatureValuesRequest(
         entity_type=entity_type, entity_id=entity_id, feature_selector=feature_selector
     )

samples/snippets/feature_store_service/read_feature_values_sample_test.py

@@ -14,7 +14,6 @@
 
 import os
 
-from google.cloud import aiplatform_v1beta1 as aiplatform
 import pytest
 import read_feature_values_sample
 
@@ -29,15 +28,11 @@ def teardown():
 def test_ucaip_generated_read_feature_values_sample_vision(capsys, shared_state):
     featurestore_id = "perm_sample_featurestore"
     entity_type_id = "perm_users"
-    feature_selector = aiplatform.FeatureSelector(
-        id_matcher=aiplatform.IdMatcher(ids=["age", "gender", "liked_genres"],)
-    )
     read_feature_values_sample.read_feature_values_sample(
         project=PROJECT_ID,
         featurestore_id=featurestore_id,
         entity_type_id=entity_type_id,
         entity_id="alice",
-        feature_selector=feature_selector,
     )
     out, _ = capsys.readouterr()
     assert "int64_value: 55" in out