docs(samples): read, import, batch_serve, batch_create features (#1046)

* feat: SDK feature store samples (create/delete fs) * feat: adding to conftest.py * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples) added changes * docs(samples): style issues * adding create entity * docs(samples): added create feature and entity type * docs(samples): edited test * docs(samples) edited style * moving constants * fixed dates * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * docs(samples): add samples to create/delete featurestore (#980) * feat: SDK feature store samples (create/delete fs) * feat: adding to conftest.py * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples) added changes * docs(samples): style issues * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * moving constants * added variables, made fixes, fixed spelling * batch_create_features_sample and test * added batch serve to bq and test * working on batch samples * fixed batch_serve_to_bq test * added read sample * moving constants * added featurestore samples * lint issues * adding revisions * adding revisions, fixing tests * editing spacing for readability Co-authored-by: Morgan Du <[email protected]>
googleapis · Mar 8, 2022 · 80dd40d · 80dd40d
1 parent 6383d4f
commit 80dd40d
Show file tree

Hide file tree

Showing 10 changed files with 440 additions and 10 deletions.
diff --git a/samples/model-builder/batch_create_features_sample.py b/samples/model-builder/batch_create_features_sample.py
@@ -0,0 +1,52 @@
+# 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.
+
+
+#  [START aiplatform_sdk_batch_create_features_sample]
+from google.cloud import aiplatform
+
+
+def batch_create_features_sample(
+    project: str,
+    location: str,
+    entity_type_id: str,
+    featurestore_id: str,
+    sync: bool = True,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_entity_type = aiplatform.featurestore.EntityType(
+        entity_type_name=entity_type_id, featurestore_id=featurestore_id
+    )
+
+    FEATURE_CONFIGS = {
+        "age": {
+            "value_type": "INT64",
+            "description": "User age"
+        },
+        "gender": {
+            "value_type": "STRING",
+            "description": "User gender"
+        },
+        "liked_genres": {
+            "value_type": "STRING_ARRAY",
+            "description": "An array of genres this user liked",
+        },
+    }
+
+    my_entity_type.batch_create_features(feature_configs=FEATURE_CONFIGS, sync=sync)
+
+
+#  [END aiplatform_sdk_batch_create_features_sample]
diff --git a/samples/model-builder/batch_create_features_sample_test.py b/samples/model-builder/batch_create_features_sample_test.py
@@ -0,0 +1,42 @@
+# 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 batch_create_features_sample
+import test_constants as constants
+
+
+def test_batch_create_features_sample(
+    mock_sdk_init, mock_get_entity_type, mock_batch_create_features
+):
+
+    batch_create_features_sample.batch_create_features_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+        sync=constants.SYNC,
+    )
+
+    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_batch_create_features.assert_called_once_with(
+        feature_configs=constants.FEATURE_CONFIGS, sync=constants.SYNC
+    )
diff --git a/samples/model-builder/batch_serve_features_to_bq_sample.py b/samples/model-builder/batch_serve_features_to_bq_sample.py
@@ -0,0 +1,46 @@
+# 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.
+
+
+#  [START aiplatform_sdk_batch_serve_features_to_bq_sample]
+from google.cloud import aiplatform
+
+
+def batch_serve_features_to_bq_sample(
+    project: str,
+    location: str,
+    featurestore_name: str,
+    bq_destination_output_uri: str,
+    read_instances_uri: str,
+    sync: bool = True,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    fs = aiplatform.featurestore.Featurestore(featurestore_name=featurestore_name)
+
+    SERVING_FEATURE_IDS = {
+        "users": ["age", "gender", "liked_genres"],
+        "movies": ["title", "average_rating", "genres"],
+    }
+
+    fs.batch_serve_to_bq(
+        bq_destination_output_uri=bq_destination_output_uri,
+        serving_feature_ids=SERVING_FEATURE_IDS,
+        read_instances_uri=read_instances_uri,
+        sync=sync,
+    )
+
+
+#  [END aiplatform_sdk_batch_serve_features_to_bq_sample]
diff --git a/samples/model-builder/batch_serve_features_to_bq_sample_test.py b/samples/model-builder/batch_serve_features_to_bq_sample_test.py
@@ -0,0 +1,45 @@
+# 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 batch_serve_features_to_bq_sample
+import test_constants as constants
+
+
+def test_batch_serve_features_to_bq_sample(
+    mock_sdk_init, mock_get_featurestore, mock_batch_serve_to_bq
+):
+
+    batch_serve_features_to_bq_sample.batch_serve_features_to_bq_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        featurestore_name=constants.FEATURESTORE_NAME,
+        bq_destination_output_uri=constants.BQ_DESTINATION_OUTPUT_URI,
+        read_instances_uri=constants.INPUT_CSV_FILE,
+        sync=constants.SYNC,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_get_featurestore.assert_called_once_with(
+        featurestore_name=constants.FEATURESTORE_NAME
+    )
+
+    mock_batch_serve_to_bq.assert_called_once_with(
+        bq_destination_output_uri=constants.BQ_DESTINATION_OUTPUT_URI,
+        serving_feature_ids=constants.SERVING_FEATURE_IDS,
+        read_instances_uri=constants.INPUT_CSV_FILE,
+        sync=constants.SYNC,
+    )
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -368,7 +368,7 @@ def mock_endpoint_explain(mock_endpoint):
 
 """
 ----------------------------------------------------------------------------
-FeatureStore Fixtures
+Feature Store Fixtures
 ----------------------------------------------------------------------------
 """
 
@@ -398,6 +398,13 @@ def mock_get_featurestore(mock_featurestore):
         yield mock_get_featurestore
 
 
+@pytest.fixture
+def mock_get_entity_type(mock_entity_type):
+    with patch.object(aiplatform.featurestore, "EntityType") as mock_get_entity_type:
+        mock_get_entity_type.return_value = mock_entity_type
+        yield mock_get_entity_type
+
+
 @pytest.fixture
 def mock_create_featurestore(mock_featurestore):
     with patch.object(
@@ -418,9 +425,7 @@ def mock_create_entity_type(mock_entity_type):
 
 @pytest.fixture
 def mock_create_feature(mock_feature):
-    with patch.object(
-        aiplatform.featurestore.Feature, "create"
-    ) as mock_create_feature:
+    with patch.object(aiplatform.featurestore.Feature, "create") as mock_create_feature:
         mock_create_feature.return_value = mock_feature
         yield mock_create_feature
 
@@ -429,3 +434,31 @@ def mock_create_feature(mock_feature):
 def mock_delete_featurestore(mock_featurestore):
     with patch.object(mock_featurestore, "delete") as mock_delete_featurestore:
         yield mock_delete_featurestore
+
+
+@pytest.fixture
+def mock_batch_serve_to_bq(mock_featurestore):
+    with patch.object(mock_featurestore, "batch_serve_to_bq") as mock_batch_serve_to_bq:
+        yield mock_batch_serve_to_bq
+
+
+@pytest.fixture
+def mock_batch_create_features(mock_entity_type):
+    with patch.object(
+        mock_entity_type, "batch_create_features"
+    ) as mock_batch_create_features:
+        yield mock_batch_create_features
+
+
+@pytest.fixture
+def mock_read_feature_values(mock_entity_type):
+    with patch.object(mock_entity_type, "read") as mock_read_feature_values:
+        yield mock_read_feature_values
+
+
+@pytest.fixture
+def mock_import_feature_values(mock_entity_type):
+    with patch.object(
+        mock_entity_type, "ingest_from_gcs"
+    ) as mock_import_feature_values:
+        yield mock_import_feature_values
diff --git a/samples/model-builder/import_feature_values_sample.py b/samples/model-builder/import_feature_values_sample.py
@@ -0,0 +1,48 @@
+# 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.
+
+
+#  [START aiplatform_sdk_import_feature_values_sample]
+import datetime
+from typing import List, Union
+
+from google.cloud import aiplatform
+
+
+def import_feature_values_sample(
+    project: str,
+    location: str,
+    entity_type_id: str,
+    featurestore_id: str,
+    feature_ids: List[str],
+    feature_time: Union[str, datetime.datetime],
+    gcs_source_uris: Union[str, List[str]],
+    gcs_source_type: str,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_entity_type = aiplatform.featurestore.EntityType(
+        entity_type_name=entity_type_id, featurestore_id=featurestore_id
+    )
+
+    my_entity_type.ingest_from_gcs(
+        feature_ids=feature_ids,
+        feature_time=feature_time,
+        gcs_source_uris=gcs_source_uris,
+        gcs_source_type=gcs_source_type,
+    )
+
+
+#  [END aiplatform_sdk_import_feature_values_sample]
diff --git a/samples/model-builder/import_feature_values_sample_test.py b/samples/model-builder/import_feature_values_sample_test.py
@@ -0,0 +1,48 @@
+# 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 import_feature_values_sample
+import test_constants as constants
+
+
+def test_import_feature_values_sample(
+    mock_sdk_init, mock_get_entity_type, mock_import_feature_values
+):
+
+    import_feature_values_sample.import_feature_values_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+        feature_ids=constants.FEATURE_IDS,
+        feature_time=constants.USERS_FEATURE_TIME,
+        gcs_source_uris=constants.USERS_GCS_SOURCE_URI,
+        gcs_source_type=constants.GCS_SOURCE_TYPE,
+    )
+
+    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_import_feature_values.assert_called_once_with(
+        feature_ids=constants.FEATURE_IDS,
+        feature_time=constants.USERS_FEATURE_TIME,
+        gcs_source_uris=constants.USERS_GCS_SOURCE_URI,
+        gcs_source_type=constants.GCS_SOURCE_TYPE,
+    )
diff --git a/samples/model-builder/read_feature_values_sample.py b/samples/model-builder/read_feature_values_sample.py
@@ -0,0 +1,42 @@
+# 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.
+
+
+#  [START aiplatform_sdk_read_features_sample]
+from typing import List, Union
+
+from google.cloud import aiplatform
+
+
+def read_feature_values_sample(
+    project: str,
+    location: str,
+    entity_type_id: str,
+    featurestore_id: str,
+    entity_ids: Union[str, List[str]],
+    feature_ids: Union[str, List[str]] = "*",
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_entity_type = aiplatform.featurestore.EntityType(
+        entity_type_name=entity_type_id, featurestore_id=featurestore_id
+    )
+
+    my_dataframe = my_entity_type.read(entity_ids=entity_ids, feature_ids=feature_ids)
+
+    return my_dataframe
+
+
+#  [END aiplatform_sdk_read_features_sample]