feat: sample code for Vertex AI Feature Store

PiperOrigin-RevId: 646188701

Author: vertex-sdk-bot

samples/model-builder/conftest.py

@@ -731,6 +731,21 @@ def mock_create_feature_group(mock_feature_group):
  yield mock_create_feature_group
 
 
+@pytest.fixture
+def mock_registry_feature():
+ mock = MagicMock(preview_resources.Feature)
+ yield mock
+
+
+@pytest.fixture
+def mock_create_registry_feature(mock_feature_group, mock_registry_feature):
+ with patch.object(
+ mock_feature_group, "create_feature"
+ ) as mock_create_registry_feature:
+ mock_create_registry_feature.return_value = mock_registry_feature
+ yield mock_create_registry_feature
+
+
 @pytest.fixture
 def mock_create_optimized_private_online_store(mock_feature_online_store):
  with patch.object(

samples/model-builder/feature_store/create_feature_sample.py

@@ -0,0 +1,35 @@
+# Copyright 2024 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_create_feature_sample]
+
+from google.cloud import aiplatform
+from vertexai.resources.preview import feature_store
+
+
+def create_feature_sample(
+ project: str,
+ location: str,
+ feature_group_id: str,
+ feature_id: str,
+):
+ aiplatform.init(project=project, location=location)
+ feature_group = feature_store.FeatureGroup.create(
+ feature_group_id
+ )
+ feature = feature_group.create_feature(feature_id)
+ return feature
+
+
+# [END aiplatform_sdk_create_feature_sample]

samples/model-builder/feature_store/create_registry_feature_sample_test.py

@@ -0,0 +1,40 @@
+# Copyright 2024 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.
+
+from feature_store import create_feature_sample
+
+import test_constants as constants
+
+
+def test_create_feature_sample(
+ mock_sdk_init, mock_create_registry_feature, mock_create_feature_group
+):
+ create_feature_sample.create_feature_sample(
+ project=constants.PROJECT,
+ location=constants.LOCATION,
+ feature_group_id=constants.FEATURE_GROUP_ID,
+ feature_id=constants.REGISTRY_FEATURE_ID,
+ )
+
+ mock_sdk_init.assert_called_once_with(
+ project=constants.PROJECT, location=constants.LOCATION
+ )
+
+ mock_create_feature_group.assert_called_once_with(
+ constants.FEATURE_GROUP_ID
+ )
+
+ mock_create_registry_feature.assert_called_once_with(
+ constants.REGISTRY_FEATURE_ID
+ )

samples/model-builder/test_constants.py

@@ -274,6 +274,7 @@
  )
 )
 FEATURE_GROUP_ID = "sample_feature_group"
+REGISTRY_FEATURE_ID = "sample_feature"
 PROJECT_ALLOWLISTED = ["test-project"]
 
 TABULAR_TARGET_COLUMN = "target_column"