samples: add feature store samples

Author: morgandu

samples/model-builder/conftest.py

@@ -364,3 +364,48 @@ def mock_endpoint_explain(mock_endpoint):
  with patch.object(mock_endpoint, "explain") as mock_endpoint_explain:
  mock_get_endpoint.return_value = mock_endpoint
  yield mock_endpoint_explain
+
+
+"""
+----------------------------------------------------------------------------
+Featurestore Fixtures
+----------------------------------------------------------------------------
+"""
+
+
+@pytest.fixture
+def mock_featurestore():
+ mock = MagicMock(aiplatform.featurestore.Featurestore)
+ yield mock
+
+
+@pytest.fixture
+def mock_entity_type():
+ mock = MagicMock(aiplatform.featurestore.EntityType)
+ yield mock
+
+
+@pytest.fixture
+def mock_get_featurestore(mock_featurestore):
+ with patch.object(aiplatform, "Featurestore") as mock_get_featurestore:
+ mock_get_featurestore.return_value = mock_featurestore
+ yield mock_get_featurestore
+
+
+@pytest.fixture
+def mock_batch_serve_to_df(mock_featurestore):
+ with patch.object(mock_featurestore, "batch_serve_to_df") as mock_batch_serve_to_df:
+ yield mock_batch_serve_to_df
+
+
+@pytest.fixture
+def mock_get_entity_type(mock_entity_type):
+ with patch.object(aiplatform, "EntityType") as mock_get_entity_type:
+ mock_get_entity_type.return_value = mock_entity_type
+ yield mock_get_entity_type
+
+
+@pytest.fixture
+def mock_ingest_from_df(mock_entity_type):
+ with patch.object(mock_entity_type, "ingest_from_df") as mock_ingest_from_df:
+ yield mock_ingest_from_df

samples/model-builder/entity_type_ingest_from_df_sample.py

@@ -0,0 +1,58 @@
+# 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 pandas as pd
+
+from google.cloud import aiplatform
+
+# [START aiplatform_sdk_entity_type_ingest_feature_values_from_df_sample]
+def entity_type_ingest_feature_values_from_df_with_feature_time_field_sample(
+ project: str, location: str, entity_type_name: str,
+):
+ aiplatform.init(project=project, location=location)
+
+ et = aiplatform.EntityType(entity_type_name=entity_type_name)
+
+ df_source = pd.DataFrame(
+ data=[
+ {
+ "movie_id": "movie_01",
+ "average_rating": 4.9,
+ "title": "The Shawshank Redemption",
+ "genres": "Drama",
+ "update_time": "2021-08-20 20:44:11.094375+00:00",
+ },
+ {
+ "movie_id": "movie_02",
+ "average_rating": 4.2,
+ "title": "The Shining",
+ "genres": "Horror",
+ "update_time": "2021-08-20 20:44:11.094375+00:00",
+ },
+ ],
+ columns=["movie_id", "average_rating", "title", "genres", "update_time"],
+ )
+
+ et.ingest_from_df(
+ feature_ids=["movid_id", "average_rating", "title", "genres"],
+ feature_time="update_time",
+ df_source=df_source,
+ entity_id_field="movie_id",
+ )
+
+ return et
+
+
+# [END aiplatform_sdk_entity_type_ingest_feature_values_from_df_sample]

samples/model-builder/entity_type_ingest_from_df_sample_test.py

@@ -0,0 +1,33 @@
+# 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 entity_type_ingest_from_df_sample
+
+import test_constants as constants
+
+
+def test_image_dataset_create_classification_sample(
+ mock_sdk_init, mock_get_entity_type, mock_ingest_from_df
+):
+ entity_type_ingest_from_df_sample.entity_type_ingest_feature_values_from_df_with_feature_time_field_sample(
+ project=constants.PROJECT,
+ location=constants.LOCATION,
+ )
+
+ mock_sdk_init.assert_called_once_with(
+ project=constants.PROJECT, location=constants.LOCATION
+ )
+ mock_ingest_from_df.assert_called_once_with(
+ )

samples/model-builder/featurestore_batch_serve_feature_values_to_df_sample.py

@@ -0,0 +1,60 @@
+# 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 pandas as pd
+
+from google.cloud import aiplatform
+
+# [START aiplatform_sdk_featurestore_batch_serve_feature_values_to_df_sample]
+def featurestore_batch_serve_feature_values_to_df_sample(
+ project: str, location: str, featurestore_name: str,
+):
+ aiplatform.init(project=project, location=location)
+
+ fs = aiplatform.Featurestore(featurestore_name=featurestore_name)
+
+ read_instances_df = pd.DataFrame(
+ data=[
+ ["alice", "movie_01", "2021-09-15T08:28:14Z"],
+ ["bob", "movie_02", "2021-09-15T08:28:14Z"],
+ ["dav", "movie_03", "2021-09-15T08:28:14Z"],
+ ["eve", "movie_04", "2021-09-15T08:28:14Z"],
+ ["alice", "movie_03", "2021-09-14T09:35:15Z"],
+ ["bob", "movie_04", "2020-02-14T09:35:15Z"],
+ ],
+ columns=["users_entity_type_id", "movies_entity_type_id", "timestamp"],
+ )
+ read_instances_df = read_instances_df.astype({"timestamp": "datetime64"})
+
+ df = fs.batch_serve_to_df(
+ serving_feature_ids={
+ "users_entity_type_id": [
+ "user_age_feature_id",
+ "user_gender_feature_id",
+ "user_liked_genres_feature_id",
+ ],
+ "movies_entity_type_id": [
+ "movie_title_feature_id",
+ "movie_genres_feature_id",
+ "movie_average_rating_feature_id",
+ ],
+ },
+ read_instances_df=read_instances_df,
+ )
+
+ return df
+
+
+# [END aiplatform_sdk_featurestore_batch_serve_feature_values_to_df_sample]