chore: add Feature init/get

PiperOrigin-RevId: 635498993

google/cloud/aiplatform/utils/__init__.py

@@ -640,16 +640,72 @@ class FeatureOnlineStoreClientWithOverride(ClientWithOverride):
 
 
 class FeatureRegistryClientWithOverride(ClientWithOverride):
+    """Adds function override for client classes to support new Feature Store.
+
+    `feature_path()` and `parse_feature_path()` are overriden here to compensate
+    for the auto-generated GAPIC class which only supports Feature Store
+    Legacy's feature paths.
+    """
+
+    @staticmethod
+    def feature_path(
+        project: str,
+        location: str,
+        feature_group: str,
+        feature: str,
+    ) -> str:
+        return "projects/{project}/locations/{location}/featureGroups/{feature_group}/features/{feature}".format(
+            project=project,
+            location=location,
+            feature_group=feature_group,
+            feature=feature,
+        )
+
+    @staticmethod
+    def parse_feature_path(path: str) -> Dict[str, str]:
+        """Parses a feature path into its component segments."""
+        m = re.match(
+            r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/featureGroups/(?P<feature_group>.+?)/features/(?P<feature>.+?)$",
+            path,
+        )
+        return m.groupdict() if m else {}
+
+    class FeatureRegistryServiceClientV1(
+        feature_registry_service_client_v1.FeatureRegistryServiceClient
+    ):
+        @staticmethod
+        def feature_path(project: str, location: str, feature_group: str, feature: str):
+            return FeatureRegistryClientWithOverride.feature_path(
+                project, location, feature_group, feature
+            )
+
+        @staticmethod
+        def parse_feature_path(path: str) -> Dict[str, str]:
+            return FeatureRegistryClientWithOverride.parse_feature_path(path)
+
+    class FeatureRegistryServiceClientV1Beta1(
+        feature_registry_service_client_v1beta1.FeatureRegistryServiceClient
+    ):
+        @staticmethod
+        def feature_path(project: str, location: str, feature_group: str, feature: str):
+            return FeatureRegistryClientWithOverride.feature_path(
+                project, location, feature_group, feature
+            )
+
+        @staticmethod
+        def parse_feature_path(path: str) -> Dict[str, str]:
+            return FeatureRegistryClientWithOverride.parse_feature_path(path)
+
     _is_temporary = True
     _default_version = compat.DEFAULT_VERSION
     _version_map = (
         (
             compat.V1,
-            feature_registry_service_client_v1.FeatureRegistryServiceClient,
+            FeatureRegistryServiceClientV1,
         ),
         (
             compat.V1BETA1,
-            feature_registry_service_client_v1beta1.FeatureRegistryServiceClient,
+            FeatureRegistryServiceClientV1Beta1,
         ),
     )
 

tests/unit/vertexai/conftest.py

@@ -51,6 +51,9 @@
 from google.cloud.aiplatform.compat.services import (
     feature_online_store_admin_service_client,
 )
+from google.cloud.aiplatform.compat.services import (
+    feature_registry_service_client,
+)
 from feature_store_constants import (
     _TEST_BIGTABLE_FOS1,
     _TEST_EMBEDDING_FV1,
@@ -61,6 +64,7 @@
     _TEST_OPTIMIZED_FV2,
     _TEST_PSC_OPTIMIZED_FOS,
     _TEST_OPTIMIZED_EMBEDDING_FV,
+    _TEST_FG1_F1,
 )
 
 _TEST_PROJECT = "test-project"
@@ -496,3 +500,13 @@ def get_optimized_fv_no_endpointmock():
     ) as get_optimized_fv_no_endpointmock:
         get_optimized_fv_no_endpointmock.return_value = _TEST_OPTIMIZED_FV2
         yield get_optimized_fv_no_endpointmock
+
+
+@pytest.fixture
+def get_feature_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "get_feature",
+    ) as get_fg_mock:
+        get_fg_mock.return_value = _TEST_FG1_F1
+        yield get_fg_mock

tests/unit/vertexai/feature_store_constants.py

@@ -314,3 +314,17 @@
 )
 
 _TEST_FG_LIST = [_TEST_FG1, _TEST_FG2, _TEST_FG3]
+
+_TEST_FG1_F1_ID = "my_fg1_f1"
+_TEST_FG1_F1_PATH = (
+    f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/features/{_TEST_FG1_F1_ID}"
+)
+_TEST_FG1_F1_DESCRIPTION = "My feature 1 in feature group 1"
+_TEST_FG1_F1_LABELS = {"my_fg1_feature": "f1"}
+_TEST_FG1_F1_POINT_OF_CONTACT = "fg1-f1-announce-list"
+_TEST_FG1_F1 = types.feature.Feature(
+    name=_TEST_FG1_F1_PATH,
+    description=_TEST_FG1_F1_DESCRIPTION,
+    labels=_TEST_FG1_F1_LABELS,
+    point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
+)

tests/unit/vertexai/test_feature.py

@@ -0,0 +1,130 @@
+# -*- 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.
+#
+
+import re
+from typing import Dict
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import base
+from vertexai.resources.preview import (
+    Feature,
+)
+import pytest
+
+
+from feature_store_constants import (
+    _TEST_PROJECT,
+    _TEST_LOCATION,
+    _TEST_FG1_ID,
+    _TEST_FG1_F1_ID,
+    _TEST_FG1_F1_PATH,
+    _TEST_FG1_F1_DESCRIPTION,
+    _TEST_FG1_F1_LABELS,
+    _TEST_FG1_F1_POINT_OF_CONTACT,
+)
+
+
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+
+
+def feature_eq(
+    feature_to_check: Feature,
+    name: str,
+    resource_name: str,
+    project: str,
+    location: str,
+    description: str,
+    labels: Dict[str, str],
+    point_of_contact: str,
+):
+    """Check if a Feature has the appropriate values set."""
+    assert feature_to_check.name == name
+    assert feature_to_check.resource_name == resource_name
+    assert feature_to_check.project == project
+    assert feature_to_check.location == location
+    assert feature_to_check.description == description
+    assert feature_to_check.labels == labels
+    assert feature_to_check.point_of_contact == point_of_contact
+
+
+def test_init_with_feature_id_and_no_fg_id_raises_error(get_feature_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Since feature is not provided as a path, please specify"
+            + " feature_group_id."
+        ),
+    ):
+        Feature(_TEST_FG1_F1_ID)
+
+
+def test_init_with_feature_path_and_fg_id_raises_error(get_feature_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Since feature is provided as a path, feature_group_id should not be specified."
+        ),
+    ):
+        Feature(_TEST_FG1_F1_PATH, feature_group_id=_TEST_FG1_ID)
+
+
+def test_init_with_feature_id(get_feature_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature = Feature(_TEST_FG1_F1_ID, feature_group_id=_TEST_FG1_ID)
+
+    get_feature_mock.assert_called_once_with(
+        name=_TEST_FG1_F1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F1_ID,
+        resource_name=_TEST_FG1_F1_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_F1_DESCRIPTION,
+        labels=_TEST_FG1_F1_LABELS,
+        point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
+    )
+
+
+def test_init_with_feature_path(get_feature_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature = Feature(_TEST_FG1_F1_PATH)
+
+    get_feature_mock.assert_called_once_with(
+        name=_TEST_FG1_F1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F1_ID,
+        resource_name=_TEST_FG1_F1_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_F1_DESCRIPTION,
+        labels=_TEST_FG1_F1_LABELS,
+        point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
+    )

tests/unit/vertexai/test_feature_group.py

@@ -60,7 +60,13 @@
     _TEST_FG3_ENTITY_ID_COLUMNS,
     _TEST_FG3_LABELS,
     _TEST_FG_LIST,
+    _TEST_FG1_F1_ID,
+    _TEST_FG1_F1_PATH,
+    _TEST_FG1_F1_DESCRIPTION,
+    _TEST_FG1_F1_LABELS,
+    _TEST_FG1_F1_POINT_OF_CONTACT,
 )
+from test_feature import feature_eq
 
 
 pytestmark = pytest.mark.usefixtures("google_auth_mock")
@@ -334,3 +340,26 @@ def test_delete(force, delete_fg_mock, get_fg_mock, fg_logger_mock, sync=True):
             ),
         ]
     )
+
+
+def test_get_feature(get_fg_mock, get_feature_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    fg = FeatureGroup(_TEST_FG1_ID)
+    feature = fg.get_feature(_TEST_FG1_F1_ID)
+
+    get_feature_mock.assert_called_once_with(
+        name=_TEST_FG1_F1_PATH,
+        retry=base._DEFAULT_RETRY,
+    )
+
+    feature_eq(
+        feature,
+        name=_TEST_FG1_F1_ID,
+        resource_name=_TEST_FG1_F1_PATH,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        description=_TEST_FG1_F1_DESCRIPTION,
+        labels=_TEST_FG1_F1_LABELS,
+        point_of_contact=_TEST_FG1_F1_POINT_OF_CONTACT,
+    )

vertexai/resources/preview/__init__.py

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

vertexai/resources/preview/feature_store/__init__.py

@@ -16,6 +16,10 @@
 #
 """The vertexai resources preview module."""
 
+from vertexai.resources.preview.feature_store.feature import (
+    Feature,
+)
+
 from vertexai.resources.preview.feature_store.feature_group import (
     FeatureGroup,
 )
@@ -41,6 +45,7 @@
 )
 
 __all__ = (
+    Feature,
     FeatureGroup,
     FeatureGroupBigQuerySource,
     FeatureOnlineStoreType,