docs(samples): Feature Store: Streaming ingestion code sample and test

vertex-sdk-bot · copybara-github · commit bc9e2cf55f6b · 2022-12-09T13:51:32.000-08:00
PiperOrigin-RevId: 494255355
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -114,7 +114,7 @@ def mock_create_image_dataset(mock_image_dataset):
 @pytest.fixture
 def mock_create_tabular_dataset(mock_tabular_dataset):
     with patch.object(
-            aiplatform.TabularDataset, "create"
+        aiplatform.TabularDataset, "create"
     ) as mock_create_tabular_dataset:
         mock_create_tabular_dataset.return_value = mock_tabular_dataset
         yield mock_create_tabular_dataset
@@ -123,7 +123,7 @@ def mock_create_tabular_dataset(mock_tabular_dataset):
 @pytest.fixture
 def mock_create_time_series_dataset(mock_time_series_dataset):
     with patch.object(
-            aiplatform.TimeSeriesDataset, "create"
+        aiplatform.TimeSeriesDataset, "create"
     ) as mock_create_time_series_dataset:
         mock_create_time_series_dataset.return_value = mock_time_series_dataset
         yield mock_create_time_series_dataset
@@ -251,7 +251,9 @@ def mock_run_automl_forecasting_training_job(mock_forecasting_training_job):
 
 @pytest.fixture
 def mock_get_automl_forecasting_seq2seq_training_job(mock_forecasting_training_job):
-    with patch.object(aiplatform, "SequenceToSequencePlusForecastingTrainingJob") as mock:
+    with patch.object(
+        aiplatform, "SequenceToSequencePlusForecastingTrainingJob"
+    ) as mock:
         mock.return_value = mock_forecasting_training_job
         yield mock
 
@@ -445,6 +447,7 @@ def mock_endpoint_explain(mock_endpoint):
         mock_get_endpoint.return_value = mock_endpoint
         yield mock_endpoint_explain
 
+
 # ----------------------------------------------------------------------------
 # Hyperparameter Tuning Job Fixtures
 # ----------------------------------------------------------------------------
@@ -471,7 +474,9 @@ def mock_run_hyperparameter_tuning_job(mock_hyperparameter_tuning_job):
 
 @pytest.fixture
 def mock_hyperparameter_tuning_job_get(mock_hyperparameter_tuning_job):
-    with patch.object(aiplatform.HyperparameterTuningJob, "get") as mock_hyperparameter_tuning_job_get:
+    with patch.object(
+        aiplatform.HyperparameterTuningJob, "get"
+    ) as mock_hyperparameter_tuning_job_get:
         mock_hyperparameter_tuning_job_get.return_value = mock_hyperparameter_tuning_job
         yield mock_hyperparameter_tuning_job_get
 
@@ -530,7 +535,7 @@ def mock_get_entity_type(mock_entity_type):
 @pytest.fixture
 def mock_create_featurestore(mock_featurestore):
     with patch.object(
-            aiplatform.featurestore.Featurestore, "create"
+        aiplatform.featurestore.Featurestore, "create"
     ) as mock_create_featurestore:
         mock_create_featurestore.return_value = mock_featurestore
         yield mock_create_featurestore
@@ -539,7 +544,7 @@ def mock_create_featurestore(mock_featurestore):
 @pytest.fixture
 def mock_create_entity_type(mock_entity_type):
     with patch.object(
-            aiplatform.featurestore.EntityType, "create"
+        aiplatform.featurestore.EntityType, "create"
     ) as mock_create_entity_type:
         mock_create_entity_type.return_value = mock_entity_type
         yield mock_create_entity_type
@@ -567,7 +572,7 @@ def mock_batch_serve_to_bq(mock_featurestore):
 @pytest.fixture
 def mock_batch_create_features(mock_entity_type):
     with patch.object(
-            mock_entity_type, "batch_create_features"
+        mock_entity_type, "batch_create_features"
     ) as mock_batch_create_features:
         yield mock_batch_create_features
 
@@ -581,11 +586,19 @@ def mock_read_feature_values(mock_entity_type):
 @pytest.fixture
 def mock_import_feature_values(mock_entity_type):
     with patch.object(
-            mock_entity_type, "ingest_from_gcs"
+        mock_entity_type, "ingest_from_gcs"
     ) as mock_import_feature_values:
         yield mock_import_feature_values
 
 
+@pytest.fixture
+def mock_write_feature_values(mock_entity_type):
+    with patch.object(
+        mock_entity_type.preview, "write_feature_values"
+    ) as mock_write_feature_values:
+        yield mock_write_feature_values
+
+
 """
 ----------------------------------------------------------------------------
 Experiment Tracking Fixtures
@@ -724,7 +737,7 @@ def mock_context_list(mock_context):
 @pytest.fixture
 def mock_create_schema_base_context(mock_context):
     with patch.object(
-            aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
+        aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
     ) as mock_create_schema_base_context:
         mock_create_schema_base_context.return_value = mock_context
         yield mock_create_schema_base_context
@@ -782,7 +795,7 @@ def mock_create_artifact(mock_artifact):
 @pytest.fixture
 def mock_create_schema_base_artifact(mock_artifact):
     with patch.object(
-            aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
+        aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
     ) as mock_create_schema_base_artifact:
         mock_create_schema_base_artifact.return_value = mock_artifact
         yield mock_create_schema_base_artifact
@@ -791,7 +804,7 @@ def mock_create_schema_base_artifact(mock_artifact):
 @pytest.fixture
 def mock_create_schema_base_execution(mock_execution):
     with patch.object(
-            aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
+        aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
     ) as mock_create_schema_base_execution:
         mock_create_schema_base_execution.return_value = mock_execution
         yield mock_create_schema_base_execution
@@ -837,7 +850,7 @@ def mock_log_metrics():
 @pytest.fixture
 def mock_log_time_series_metrics():
     with patch.object(
-            aiplatform, "log_time_series_metrics"
+        aiplatform, "log_time_series_metrics"
     ) as mock_log_time_series_metrics:
         mock_log_time_series_metrics.return_value = None
         yield mock_log_time_series_metrics
@@ -909,24 +922,24 @@ def mock_get_params(mock_params, mock_experiment_run):
 @pytest.fixture
 def mock_get_time_series_metrics(mock_time_series_metrics, mock_experiment_run):
     with patch.object(
-            mock_experiment_run, "get_time_series_data_frame"
+        mock_experiment_run, "get_time_series_data_frame"
     ) as mock_get_time_series_metrics:
         mock_get_time_series_metrics.return_value = mock_time_series_metrics
         yield mock_get_time_series_metrics
 
 
 @pytest.fixture
 def mock_get_classification_metrics(mock_classification_metrics, mock_experiment_run):
-    with patch.object(mock_experiment_run, "get_classification_metrics") as mock_get_classification_metrics:
+    with patch.object(
+        mock_experiment_run, "get_classification_metrics"
+    ) as mock_get_classification_metrics:
         mock_get_classification_metrics.return_value = mock_classification_metrics
         yield mock_get_classification_metrics
 
 
 @pytest.fixture
 def mock_get_artifacts(mock_artifacts, mock_experiment_run):
-    with patch.object(
-        mock_experiment_run, "get_artifacts"
-    ) as mock_get_artifacts:
+    with patch.object(mock_experiment_run, "get_artifacts") as mock_get_artifacts:
         mock_get_artifacts.return_value = mock_artifacts
         yield mock_get_artifacts
 
@@ -966,7 +979,9 @@ def mock_get_model(mock_model_registry):
 
 @pytest.fixture
 def mock_get_model_version_info(mock_model_registry):
-    with patch.object(mock_model_registry, "get_version_info") as mock_get_model_version_info:
+    with patch.object(
+        mock_model_registry, "get_version_info"
+    ) as mock_get_model_version_info:
         mock_get_model_version_info.return_value = mock_version_info
         yield mock_get_model_version_info
 
@@ -987,13 +1002,17 @@ def mock_delete_version(mock_model_registry):
 
 @pytest.fixture
 def mock_add_version_aliases(mock_model_registry):
-    with patch.object(mock_model_registry, "add_version_aliases") as mock_add_version_aliases:
+    with patch.object(
+        mock_model_registry, "add_version_aliases"
+    ) as mock_add_version_aliases:
         mock_add_version_aliases.return_value = None
         yield mock_add_version_aliases
 
 
 @pytest.fixture
 def mock_remove_version_aliases(mock_model_registry):
-    with patch.object(mock_model_registry, "remove_version_aliases") as mock_remove_version_aliases:
+    with patch.object(
+        mock_model_registry, "remove_version_aliases"
+    ) as mock_remove_version_aliases:
         mock_remove_version_aliases.return_value = None
         yield mock_remove_version_aliases
diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -213,6 +213,18 @@
 ENTITY_TYPE_ID = "users"
 ENTITY_IDS = ["alice", "bob"]
 ENTITY_TYPE_NAME = f"projects/{PROJECT}/locations/{LOCATION}/featurestores/{FEATURESTORE_ID}/entityTypes/{ENTITY_TYPE_ID}"
+ENTITY_INSTANCES = {
+    "movie_01": {
+        "title": "The Shawshank Redemption",
+        "average_rating": 4.7,
+        "genre": "Drama",
+    },
+    "movie_02": {
+        "title": "Everything Everywhere All At Once",
+        "average_rating": 4.4,
+        "genre": "Adventure",
+    },
+}
 FEATURE_ID = "liked_genres"
 FEATURE_IDS = ["age", "gender", "liked_genres"]
 FEATURE_NAME = f"projects/{PROJECT}/locations/{LOCATION}/featurestores/{FEATURESTORE_ID}/entityTypes/{ENTITY_TYPE_ID}/features/{FEATURE_ID}"
@@ -290,10 +302,10 @@
 # Hyperparameter tuning job
 HYPERPARAMETER_TUNING_JOB_DISPLAY_NAME = "hpt_job"
 HYPERPARAMETER_TUNING_JOB_ID = "4447046521673744384"
-HYPERPARAMETER_TUNING_JOB_METRIC_SPEC = {'loss': 'minimize'}
+HYPERPARAMETER_TUNING_JOB_METRIC_SPEC = {"loss": "minimize"}
 HYPERPARAMETER_TUNING_JOB_MAX_TRIAL_COUNT = 128
 HYPERPARAMETER_TUNING_JOB_PARALLEL_TRIAL_COUNT = 8
-HYPERPARAMETER_TUNING_JOB_LABELS = {'my_key': 'my_value'}
+HYPERPARAMETER_TUNING_JOB_LABELS = {"my_key": "my_value"}
 
 # Custom job
 CUSTOM_JOB_DISPLAY_NAME = "custom_job"
diff --git a/samples/model-builder/write_feature_values_sample.py b/samples/model-builder/write_feature_values_sample.py
@@ -0,0 +1,49 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Serve feature values from a single entity for a particular entity type.
+# See https://cloud.google.com/vertex-ai/docs/featurestore/setup before running
+# the code snippet
+
+# [START aiplatform_write_feature_values_sample]
+from google.cloud import aiplatform
+
+
+def write_feature_values_sample(
+    project: str, location: str, entity_type_id: str, featurestore_id: str
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_entity_type = aiplatform.featurestore.EntityType(
+        entity_type_name=entity_type_id, featurestore_id=featurestore_id
+    )
+
+    my_data = {
+        "movie_01": {
+            "title": "The Shawshank Redemption",
+            "average_rating": 4.7,
+            "genre": "Drama",
+        },
+        "movie_02": {
+            "title": "Everything Everywhere All At Once",
+            "average_rating": 4.4,
+            "genre": "Adventure",
+        },
+    }
+
+    my_entity_type.preview.write_feature_values(instances=my_data)
+
+
+# [END aiplatform_write_feature_values_sample]
diff --git a/samples/model-builder/write_feature_values_sample_test.py b/samples/model-builder/write_feature_values_sample_test.py
@@ -0,0 +1,41 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import test_constants as constants
+import write_feature_values_sample
+
+
+def test_write_feature_values_sample(
+    mock_sdk_init, mock_get_entity_type, mock_write_feature_values
+):
+
+    write_feature_values_sample.write_feature_values_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_get_entity_type.assert_called_once_with(
+        entity_type_name=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+    )
+
+    mock_write_feature_values.assert_called_once_with(
+        instances=constants.ENTITY_INSTANCES
+    )
