feat: add FeatureGroup init/get

PiperOrigin-RevId: 631486385

google/cloud/aiplatform/compat/__init__.py

@@ -36,6 +36,9 @@
     services.feature_online_store_service_client = (
         services.feature_online_store_service_client_v1beta1
     )
+    services.feature_registry_service_client = (
+        services.feature_registry_service_client_v1beta1
+    )
     services.featurestore_online_serving_service_client = (
         services.featurestore_online_serving_service_client_v1beta1
     )
@@ -91,6 +94,7 @@
     types.feature_online_store_admin_service = (
         types.feature_online_store_admin_service_v1beta1
     )
+    types.feature_registry_service = types.feature_registry_service_v1beta1
     types.feature_online_store_service = types.feature_online_store_service_v1beta1
     types.feature_selector = types.feature_selector_v1beta1
     types.feature_view = types.feature_view_v1beta1
@@ -157,6 +161,9 @@
     services.feature_online_store_admin_service_client = (
         services.feature_online_store_admin_service_client_v1
     )
+    services.feature_registry_service_client = (
+        services.feature_registry_service_client_v1
+    )
     services.feature_online_store_service_client = (
         services.feature_online_store_service_client_v1
     )
@@ -208,6 +215,7 @@
     types.feature_online_store_admin_service = (
         types.feature_online_store_admin_service_v1
     )
+    types.feature_registry_service = types.feature_registry_service_v1
     types.feature_online_store_service = types.feature_online_store_service_v1
     types.feature_selector = types.feature_selector_v1
     types.feature_view = types.feature_view_v1

google/cloud/aiplatform/compat/services/__init__.py

@@ -36,6 +36,9 @@
 from google.cloud.aiplatform_v1beta1.services.feature_online_store_admin_service import (
     client as feature_online_store_admin_service_client_v1beta1,
 )
+from google.cloud.aiplatform_v1beta1.services.feature_registry_service import (
+    client as feature_registry_service_client_v1beta1,
+)
 from google.cloud.aiplatform_v1beta1.services.featurestore_online_serving_service import (
     client as featurestore_online_serving_service_client_v1beta1,
 )
@@ -119,6 +122,9 @@
 from google.cloud.aiplatform_v1.services.feature_online_store_admin_service import (
     client as feature_online_store_admin_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.feature_registry_service import (
+    client as feature_registry_service_client_v1,
+)
 from google.cloud.aiplatform_v1.services.featurestore_online_serving_service import (
     client as featurestore_online_serving_service_client_v1,
 )
@@ -174,6 +180,7 @@
     endpoint_service_client_v1,
     feature_online_store_service_client_v1,
     feature_online_store_admin_service_client_v1,
+    feature_registry_service_client_v1,
     featurestore_online_serving_service_client_v1,
     featurestore_service_client_v1,
     index_service_client_v1,
@@ -196,6 +203,7 @@
     endpoint_service_client_v1beta1,
     feature_online_store_service_client_v1beta1,
     feature_online_store_admin_service_client_v1beta1,
+    feature_registry_service_client_v1beta1,
     featurestore_online_serving_service_client_v1beta1,
     featurestore_service_client_v1beta1,
     index_service_client_v1beta1,

google/cloud/aiplatform/compat/types/__init__.py

@@ -47,6 +47,7 @@
     feature_online_store as feature_online_store_v1beta1,
     feature_online_store_admin_service as feature_online_store_admin_service_v1beta1,
     feature_online_store_service as feature_online_store_service_v1beta1,
+    feature_registry_service as feature_registry_service_v1beta1,
     feature_selector as feature_selector_v1beta1,
     feature_view as feature_view_v1beta1,
     feature_view_sync as feature_view_sync_v1beta1,
@@ -136,6 +137,7 @@
     feature_online_store as feature_online_store_v1,
     feature_online_store_admin_service as feature_online_store_admin_service_v1,
     feature_online_store_service as feature_online_store_service_v1,
+    feature_registry_service as feature_registry_service_v1,
     feature_selector as feature_selector_v1,
     feature_view as feature_view_v1,
     feature_view_sync as feature_view_sync_v1,

google/cloud/aiplatform/utils/__init__.py

@@ -43,6 +43,7 @@
     extension_registry_service_client_v1beta1,
     feature_online_store_admin_service_client_v1beta1,
     feature_online_store_service_client_v1beta1,
+    feature_registry_service_client_v1beta1,
     featurestore_online_serving_service_client_v1beta1,
     featurestore_service_client_v1beta1,
     index_service_client_v1beta1,
@@ -71,6 +72,7 @@
     endpoint_service_client_v1,
     feature_online_store_admin_service_client_v1,
     feature_online_store_service_client_v1,
+    feature_registry_service_client_v1,
     featurestore_online_serving_service_client_v1,
     featurestore_service_client_v1,
     index_service_client_v1,
@@ -100,6 +102,7 @@
     endpoint_service_client_v1beta1.EndpointServiceClient,
     feature_online_store_admin_service_client_v1beta1.FeatureOnlineStoreAdminServiceClient,
     feature_online_store_service_client_v1beta1.FeatureOnlineStoreServiceClient,
+    feature_registry_service_client_v1beta1.FeatureRegistryServiceClient,
     featurestore_online_serving_service_client_v1beta1.FeaturestoreOnlineServingServiceClient,
     featurestore_service_client_v1beta1.FeaturestoreServiceClient,
     index_service_client_v1beta1.IndexServiceClient,
@@ -120,6 +123,7 @@
     endpoint_service_client_v1.EndpointServiceClient,
     feature_online_store_admin_service_client_v1.FeatureOnlineStoreAdminServiceClient,
     feature_online_store_service_client_v1.FeatureOnlineStoreServiceClient,
+    feature_registry_service_client_v1.FeatureRegistryServiceClient,
     featurestore_online_serving_service_client_v1.FeaturestoreOnlineServingServiceClient,
     featurestore_service_client_v1.FeaturestoreServiceClient,
     metadata_service_client_v1.MetadataServiceClient,
@@ -635,6 +639,21 @@ class FeatureOnlineStoreClientWithOverride(ClientWithOverride):
     )
 
 
+class FeatureRegistryClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.DEFAULT_VERSION
+    _version_map = (
+        (
+            compat.V1,
+            feature_registry_service_client_v1.FeatureRegistryServiceClient,
+        ),
+        (
+            compat.V1BETA1,
+            feature_registry_service_client_v1beta1.FeatureRegistryServiceClient,
+        ),
+    )
+
+
 class FeaturestoreClientWithOverride(ClientWithOverride):
     _is_temporary = True
     _default_version = compat.DEFAULT_VERSION

tests/unit/vertexai/feature_store_constants.py

@@ -262,3 +262,19 @@
         ]
     )
 )
+
+_TEST_FG1_ID = "my_fg1"
+_TEST_FG1_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}"
+_TEST_FG1_BQ_URI = f"bq://{_TEST_PROJECT}.my_dataset.my_table_for_fg1"
+_TEST_FG1_ENTITY_ID_COLUMNS = ["entity_id"]
+_TEST_FG1_LABELS = {"my_key": "my_fg1"}
+_TEST_FG1 = types.feature_group.FeatureGroup(
+    name=_TEST_FG1_PATH,
+    big_query=types.feature_group.FeatureGroup.BigQuery(
+        big_query_source=types.io.BigQuerySource(
+            input_uri=_TEST_FG1_BQ_URI,
+        ),
+        entity_id_columns=_TEST_FG1_ENTITY_ID_COLUMNS,
+    ),
+    labels=_TEST_FG1_LABELS,
+)

tests/unit/vertexai/test_feature_group.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+
+# 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
+#
+#     http://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 typing import Dict, List
+from unittest.mock import patch
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import base
+from vertexai.resources.preview import (
+    FeatureGroup,
+)
+import vertexai.resources.preview.feature_store.utils as fs_utils
+import pytest
+from google.cloud.aiplatform.compat.services import (
+    feature_registry_service_client,
+)
+
+
+from feature_store_constants import (
+    _TEST_PROJECT,
+    _TEST_LOCATION,
+    _TEST_FG1,
+    _TEST_FG1_ID,
+    _TEST_FG1_PATH,
+    _TEST_FG1_BQ_URI,
+    _TEST_FG1_ENTITY_ID_COLUMNS,
+    _TEST_FG1_LABELS,
+)
+
+
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+
+
+@pytest.fixture
+def get_fg_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "get_feature_group",
+    ) as get_fg_mock:
+        get_fg_mock.return_value = _TEST_FG1
+        yield get_fg_mock
+
+
+def fg_eq(
+    fg_to_check: FeatureGroup,
+    name: str,
+    resource_name: str,
+    source_uri: str,
+    entity_id_columns: List[str],
+    project: str,
+    location: str,
+    labels: Dict[str, str],
+):
+    """Check if a FeatureGroup has the appropriate values set."""
+    assert fg_to_check.name == name
+    assert fg_to_check.resource_name == resource_name
+    assert fg_to_check.source == fs_utils.FeatureGroupBigQuerySource(
+        uri=source_uri,
+        entity_id_columns=entity_id_columns,
+    )
+    assert fg_to_check.project == project
+    assert fg_to_check.location == location
+    assert fg_to_check.labels == labels
+
+
+@pytest.mark.parametrize(
+    "feature_group_name",
+    [_TEST_FG1_ID, _TEST_FG1_PATH],
+)
+def test_init(feature_group_name, get_fg_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    fg = FeatureGroup(feature_group_name)
+
+    get_fg_mock.assert_called_once_with(
+        name=_TEST_FG1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    fg_eq(
+        fg,
+        name=_TEST_FG1_ID,
+        resource_name=_TEST_FG1_PATH,
+        source_uri=_TEST_FG1_BQ_URI,
+        entity_id_columns=_TEST_FG1_ENTITY_ID_COLUMNS,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FG1_LABELS,
+    )

tests/unit/vertexai/test_feature_view.py

@@ -26,7 +26,7 @@
 from vertexai.resources.preview import (
     FeatureView,
 )
-import vertexai.resources.preview.feature_store.utils as fv_utils
+import vertexai.resources.preview.feature_store.utils as fs_utils
 import pytest
 from google.cloud.aiplatform.compat.services import (
     feature_online_store_admin_service_client,
@@ -418,7 +418,7 @@ def test_fetch_feature_values_optimized_no_endpoint(
 ):
     """Tests that the public endpoint is not created for the optimized online store."""
     with pytest.raises(
-        fv_utils.PublicEndpointNotFoundError,
+        fs_utils.PublicEndpointNotFoundError,
         match=re.escape(
             "Public endpoint is not created yet for the optimized online "
             "store:my_esf_optimised_fos2. Please run sync and wait for it "
@@ -498,8 +498,8 @@ def test_search_nearest_entities_no_endpoint(
     try:
         FeatureView(_TEST_OPTIMIZED_FV2_PATH).search(entity_id="key1").to_dict()
         assert not fetch_feature_values_mock.called
-    except fv_utils.PublicEndpointNotFoundError as e:
-        assert isinstance(e, fv_utils.PublicEndpointNotFoundError)
+    except fs_utils.PublicEndpointNotFoundError as e:
+        assert isinstance(e, fs_utils.PublicEndpointNotFoundError)
         error_msg = (
             "Public endpoint is not created yet for the optimized online "
             "store:my_esf_optimised_fos2. Please run sync and wait for it "

vertexai/resources/preview/__init__.py

@@ -35,6 +35,7 @@
 )
 
 from vertexai.resources.preview.feature_store import (
+    FeatureGroup,
     FeatureOnlineStore,
     FeatureOnlineStoreType,
     FeatureView,
@@ -62,6 +63,7 @@
     "PersistentResource",
     "EntityType",
     "PipelineJobSchedule",
+    "FeatureGroup",
     "FeatureOnlineStoreType",
     "FeatureOnlineStore",
     "FeatureView",

vertexai/resources/preview/feature_store/__init__.py

@@ -16,6 +16,10 @@
 #
 """The vertexai resources preview module."""
 
+from vertexai.resources.preview.feature_store.feature_group import (
+    FeatureGroup,
+)
+
 from vertexai.resources.preview.feature_store.feature_online_store import (
     FeatureOnlineStore,
     FeatureOnlineStoreType,
@@ -36,6 +40,7 @@
 )
 
 __all__ = (
+    FeatureGroup,
     FeatureOnlineStoreType,
     FeatureOnlineStore,
     FeatureView,