feat: add vertex_registry_source to create_feature_view method

PiperOrigin-RevId: 724463720

samples/model-builder/feature_store/create_feature_view_from_registry_source.py

@@ -0,0 +1,42 @@
+# 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_view_from_registry_source]
+
+from google.cloud import aiplatform
+from vertexai.resources.preview import feature_store
+from typing import List, Optional
+
+
+def create_feature_view_from_registry_source(
+    project: str,
+    location: str,
+    existing_feature_online_store_id: str,
+    feature_view_id: str,
+    features: List[str],
+    project_number: Optional[int],
+):
+    aiplatform.init(project=project, location=location)
+    fos = feature_store.FeatureOnlineStore(existing_feature_online_store_id)
+    fv = fos.create_feature_view(
+        name=feature_view_id,
+        source=feature_store.utils.FeatureViewRegistrySource(
+            features=features,
+            project_number=project_number,
+        ),
+    )
+    return fv
+
+
+#  [END aiplatform_sdk_create_feature_view_from_registry_source]

samples/model-builder/feature_store/create_feature_view_from_registry_source_test.py

@@ -0,0 +1,45 @@
+# 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_view_from_registry_source
+import test_constants as constants
+
+
+def test_create_feature_view_from_registry_source_sample(
+    mock_sdk_init,
+    mock_get_feature_online_store,
+):
+    create_feature_view_from_registry_source.create_feature_view_from_registry_source(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        existing_feature_online_store_id=constants.FEATURE_ONLINE_STORE_ID,
+        feature_view_id=constants.FEATURE_VIEW_ID,
+        features=[
+            constants.FR_FEATURE_ID,
+            constants.FR_FEATURE_ID_2,
+            constants.FR_FEATURE_ID_3,
+            constants.FR_FEATURE_ID_4,
+        ],
+        project_number=constants.PROJECT_NUMBER,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_get_feature_online_store.assert_called_once()
+    mock_get_feature_online_store.return_value.create_feature_view.assert_called_once_with(
+        name=constants.FEATURE_VIEW_ID,
+        source=constants.FEATURE_VIEW_REGISTRY_SOURCE,
+    )

samples/model-builder/test_constants.py

@@ -21,6 +21,7 @@
 import vertexai
 
 PROJECT = "abc"
+PROJECT_NUMBER = 123
 LOCATION = "us-central1"
 LOCATION_EUROPE = "europe-west4"
 LOCATION_ASIA = "asia-east1"
@@ -280,6 +281,7 @@
     )
 )
 FEATURE_GROUP_ID = "sample_feature_group"
+FEATURE_GROUP_ID_2 = "sample_feature_group_2"
 FEATURE_GROUP_BQ_URI = "bq://my_proj.my_dataset.my_table"
 FEATURE_GROUP_BQ_ENTITY_ID_COLUMNS = ["id"]
 FEATURE_GROUP_SOURCE = (
@@ -289,8 +291,28 @@
     )
 )
 REGISTRY_FEATURE_ID = "sample_feature"
+REGISTRY_FEATURE_ID_2 = "sample_feature_2"
+FG_2_REGISTRY_FEATURE_ID_3 = "sample_feature_3"
+FG_2_REGISTRY_FEATURE_ID_4 = "sample_feature_4"
 VERSION_COLUMN_NAME = "feature_column"
 PROJECT_ALLOWLISTED = ["test-project"]
+FR_FEATURE_ID = ".".join([FEATURE_GROUP_ID, REGISTRY_FEATURE_ID])
+FR_FEATURE_ID_2 = ".".join([FEATURE_GROUP_ID, REGISTRY_FEATURE_ID_2])
+FR_FEATURE_ID_3 = ".".join([FEATURE_GROUP_ID_2, FG_2_REGISTRY_FEATURE_ID_3])
+FR_FEATURE_ID_4 = ".".join([FEATURE_GROUP_ID_2, FG_2_REGISTRY_FEATURE_ID_4])
+
+FEATURE_GROUPS_MAPPING = {
+    FEATURE_GROUP_ID: [REGISTRY_FEATURE_ID, REGISTRY_FEATURE_ID_2],
+    FEATURE_GROUP_ID_2: [FG_2_REGISTRY_FEATURE_ID_3, FG_2_REGISTRY_FEATURE_ID_4],
+}
+
+
+FEATURE_VIEW_REGISTRY_SOURCE = (
+    vertexai.resources.preview.feature_store.utils.FeatureViewRegistrySource(
+        features=[FR_FEATURE_ID, FR_FEATURE_ID_2, FR_FEATURE_ID_3, FR_FEATURE_ID_4],
+        project_number=PROJECT_NUMBER,
+    )
+)
 
 TABULAR_TARGET_COLUMN = "target_column"
 FORECASTNG_TIME_COLUMN = "date"
@@ -382,8 +404,12 @@
 # Vector Search
 VECTOR_SEARCH_INDEX = "123"
 VECTOR_SEARCH_INDEX_DATAPOINTS = [
-    aiplatform.compat.types.index_v1beta1.IndexDatapoint(datapoint_id="datapoint_id_1", feature_vector=[0.1, 0.2]),
-    aiplatform.compat.types.index_v1beta1.IndexDatapoint(datapoint_id="datapoint_id_2", feature_vector=[0.3, 0.4]),
+    aiplatform.compat.types.index_v1beta1.IndexDatapoint(
+        datapoint_id="datapoint_id_1", feature_vector=[0.1, 0.2]
+    ),
+    aiplatform.compat.types.index_v1beta1.IndexDatapoint(
+        datapoint_id="datapoint_id_2", feature_vector=[0.3, 0.4]
+    ),
 ]
 VECTOR_SEARCH_INDEX_DATAPOINT_IDS = ["datapoint_id_1", "datapoint_id_2"]
 VECTOR_SEARCH_INDEX_ENDPOINT = "456"

tests/unit/vertexai/conftest.py

@@ -66,6 +66,7 @@
     _TEST_FV_LIST,
     _TEST_FV1,
     _TEST_FV3,
+    _TEST_FV4,
     _TEST_OPTIMIZED_EMBEDDING_FV,
     _TEST_OPTIMIZED_FV1,
     _TEST_OPTIMIZED_FV2,
@@ -444,6 +445,16 @@ def get_rag_fv_mock():
         yield get_rag_fv_mock
 
 
+@pytest.fixture
+def get_registry_fv_mock():
+    with patch.object(
+        feature_online_store_admin_service_client.FeatureOnlineStoreAdminServiceClient,
+        "get_feature_view",
+    ) as get_rag_fv_mock:
+        get_rag_fv_mock.return_value = _TEST_FV4
+        yield get_rag_fv_mock
+
+
 @pytest.fixture
 def list_fv_mock():
     with patch.object(
@@ -478,6 +489,18 @@ def create_rag_fv_mock():
         yield create_rag_fv_mock
 
 
+@pytest.fixture
+def create_registry_fv_mock():
+    with patch.object(
+        feature_online_store_admin_service_client.FeatureOnlineStoreAdminServiceClient,
+        "create_feature_view",
+    ) as create_registry_fv_mock:
+        create_registry_fv_lro_mock = mock.Mock(ga_operation.Operation)
+        create_registry_fv_lro_mock.result.return_value = _TEST_FV4
+        create_registry_fv_mock.return_value = create_registry_fv_lro_mock
+        yield create_registry_fv_mock
+
+
 @pytest.fixture
 def create_embedding_fv_from_bq_mock():
     with patch.object(

tests/unit/vertexai/feature_store_constants.py

@@ -169,8 +169,6 @@
 )
 
 
-_TEST_FV_LIST = [_TEST_FV1, _TEST_FV2, _TEST_FV3]
-
 # Test feature view sync 1
 _TEST_FV_SYNC1_ID = "my_fv_sync1"
 _TEST_FV_SYNC1_PATH = f"{_TEST_FV1_PATH}/featureViewSyncs/my_fv_sync1"
@@ -300,6 +298,8 @@
 
 
 _TEST_FG2_ID = "my_fg2"
+_TEST_FG2_F1_ID = "my_fg2_f1"
+_TEST_FG2_F2_ID = "my_fg2_f2"
 _TEST_FG2_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG2_ID}"
 _TEST_FG2_BQ_URI = f"bq://{_TEST_PROJECT}.my_dataset.my_table_for_fg2"
 _TEST_FG2_ENTITY_ID_COLUMNS = ["entity_id1", "entity_id2"]
@@ -475,3 +475,26 @@
     point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
     feature_stats_and_anomaly=[_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY],
 )
+
+# Test feature view 4
+_TEST_FV4_ID = "my_fv4"
+_TEST_FV4_PATH = f"{_TEST_BIGTABLE_FOS1_PATH}/featureViews/my_fv4"
+_TEST_FV4_LABELS = {"my_key": "my_fv4"}
+_TEST_FV4 = types.feature_view.FeatureView(
+    name=_TEST_FV4_PATH,
+    feature_registry_source=types.feature_view.FeatureView.FeatureRegistrySource(
+        feature_groups=[
+            types.feature_view.FeatureView.FeatureRegistrySource.FeatureGroup(
+                feature_group_id=_TEST_FG1_ID,
+                feature_ids=[_TEST_FG1_F1_ID, _TEST_FG1_F2_ID],
+            ),
+            types.feature_view.FeatureView.FeatureRegistrySource.FeatureGroup(
+                feature_group_id=_TEST_FG2_ID,
+                feature_ids=[_TEST_FG2_F1_ID, _TEST_FG2_F2_ID],
+            ),
+        ],
+    ),
+    labels=_TEST_FV4_LABELS,
+)
+
+_TEST_FV_LIST = [_TEST_FV1, _TEST_FV2, _TEST_FV3, _TEST_FV4]