diff --git a/samples/snippets/feature_store_service/batch_create_features_sample.py b/samples/snippets/feature_store_service/batch_create_features_sample.py index ad090f1aff6..dcc76c19d9a 100644 --- a/samples/snippets/feature_store_service/batch_create_features_sample.py +++ b/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,34 @@ 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 ) diff --git a/samples/snippets/feature_store_service/batch_create_features_sample_test.py b/samples/snippets/feature_store_service/batch_create_features_sample_test.py index 5b31ba3e6d9..8864bd57417 100644 --- a/samples/snippets/feature_store_service/batch_create_features_sample_test.py +++ b/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 diff --git a/samples/snippets/feature_store_service/batch_read_feature_values_sample.py b/samples/snippets/feature_store_service/batch_read_feature_values_sample.py index aa2b75d6552..bd93bdb54bc 100644 --- a/samples/snippets/feature_store_service/batch_read_feature_values_sample.py +++ b/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, diff --git a/samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py b/samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py index 8093020cdef..b51d0349ca6 100644 --- a/samples/snippets/feature_store_service/batch_read_feature_values_sample_test.py +++ b/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 diff --git a/samples/snippets/feature_store_service/import_feature_values_sample.py b/samples/snippets/feature_store_service/import_feature_values_sample.py index 4ad6e073ee2..c4bb9068181 100644 --- a/samples/snippets/feature_store_service/import_feature_values_sample.py +++ b/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, diff --git a/samples/snippets/feature_store_service/import_feature_values_sample_test.py b/samples/snippets/feature_store_service/import_feature_values_sample_test.py index 9f8ca698b95..c7f927cc97f 100644 --- a/samples/snippets/feature_store_service/import_feature_values_sample_test.py +++ b/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, diff --git a/samples/snippets/feature_store_service/read_feature_values_sample.py b/samples/snippets/feature_store_service/read_feature_values_sample.py index 96ce47c4384..582f1d02c9c 100644 --- a/samples/snippets/feature_store_service/read_feature_values_sample.py +++ b/samples/snippets/feature_store_service/read_feature_values_sample.py @@ -21,7 +21,6 @@ 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, @@ -34,6 +33,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 )