feat: Feature Store - Support returning stats in get feature

PiperOrigin-RevId: 705152184

Author: vertex-sdk-bot

tests/unit/vertexai/feature_store_constants.py

@@ -459,3 +459,19 @@
  labels=_TEST_FG1_FMJ2_LABELS,
 )
 _TEST_FG1_FMJ_LIST = [_TEST_FG1_FMJ1, _TEST_FG1_FMJ2]
+
+_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY = types.feature_monitor.FeatureStatsAndAnomaly(
+ feature_id="my_fg1_f1",
+ distribution_deviation=0.5,
+ drift_detection_threshold=0.4,
+ drift_detected=True,
+ feature_monitor_job_id="1234567890",
+ feature_monitor_id="1234567891",
+)
+_TEST_FG1_F1_WITH_STATS = types.feature_v1beta1.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,
+ feature_stats_and_anomaly=[_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY],
+)

tests/unit/vertexai/test_feature.py

@@ -16,22 +16,28 @@
 #
 
 import re
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 from unittest import mock
 from unittest.mock import call, patch
 
 from google.api_core import operation as ga_operation
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
+from google.cloud.aiplatform.compat import types
 from google.cloud.aiplatform.compat.services import (
  feature_registry_service_client,
 )
+from google.cloud.aiplatform_v1beta1.services.feature_registry_service import (
+ FeatureRegistryServiceClient,
+)
 from feature_store_constants import (
  _TEST_FG1_F1_DESCRIPTION,
+ _TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY,
  _TEST_FG1_F1_ID,
  _TEST_FG1_F1_LABELS,
  _TEST_FG1_F1_PATH,
  _TEST_FG1_F1_POINT_OF_CONTACT,
+ _TEST_FG1_F1_WITH_STATS,
  _TEST_FG1_F2_DESCRIPTION,
  _TEST_FG1_F2_ID,
  _TEST_FG1_F2_LABELS,
@@ -60,6 +66,16 @@ def delete_feature_mock():
  yield delete_feature_mock
 
 
+@pytest.fixture
+def get_feature_with_stats_and_anomalies_mock():
+ with patch.object(
+ FeatureRegistryServiceClient,
+ "get_feature",
+ ) as get_feature_with_stats_and_anomalies_mock:
+ get_feature_with_stats_and_anomalies_mock.return_value = _TEST_FG1_F1_WITH_STATS
+ yield get_feature_with_stats_and_anomalies_mock
+
+
 pytestmark = pytest.mark.usefixtures("google_auth_mock")
 
 
@@ -73,6 +89,9 @@ def feature_eq(
  labels: Dict[str, str],
  point_of_contact: str,
  version_column_name: Optional[str] = None,
+ feature_stats_and_anomalies: Optional[
+ List[types.feature_monitor.FeatureStatsAndAnomaly]
+ ] = None,
 ):
  """Check if a Feature has the appropriate values set."""
  assert feature_to_check.name == name
@@ -85,6 +104,10 @@ def feature_eq(
 
  if version_column_name:
  assert feature_to_check.version_column_name == version_column_name
+ if feature_stats_and_anomalies:
+ assert (
+ feature_to_check.feature_stats_and_anomalies == feature_stats_and_anomalies
+ )
 
 
 def test_init_with_feature_id_and_no_fg_id_raises_error(get_feature_mock):
@@ -206,6 +229,33 @@ def test_init_with_feature_path_for_explicit_version_column(
  )
 
 
+def test_init_with_latest_stats_count(get_feature_with_stats_and_anomalies_mock):
+ aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+ feature = Feature(name=_TEST_FG1_F1_PATH, latest_stats_count=1)
+
+ get_feature_with_stats_and_anomalies_mock.assert_called_once_with(
+ request=types.featurestore_service_v1beta1.GetFeatureRequest(
+ name=_TEST_FG1_F1_PATH,
+ feature_stats_and_anomaly_spec=types.feature_monitor.FeatureStatsAndAnomalySpec(
+ latest_stats_count=1
+ ),
+ )
+ )
+
+ 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,
+ feature_stats_and_anomalies=[_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY],
+ )
+
+
 @pytest.mark.parametrize("sync", [True])
 def test_delete_feature(
  get_fg_mock, get_feature_mock, delete_feature_mock, base_logger_mock, sync

tests/unit/vertexai/test_feature_group.py

@@ -70,6 +70,8 @@
  _TEST_FG1_F1_DESCRIPTION,
  _TEST_FG1_F1_LABELS,
  _TEST_FG1_F1_POINT_OF_CONTACT,
+ _TEST_FG1_F1_WITH_STATS,
+ _TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY,
  _TEST_FG1_F2,
  _TEST_FG1_F2_ID,
  _TEST_FG1_F2_PATH,
@@ -201,6 +203,16 @@ def list_feature_monitors_mock():
  yield list_feature_monitors_mock
 
 
+@pytest.fixture
+def get_feature_with_stats_and_anomalies_mock():
+ with patch.object(
+ FeatureRegistryServiceClient,
+ "get_feature",
+ ) as get_feature_with_stats_and_anomalies_mock:
+ get_feature_with_stats_and_anomalies_mock.return_value = _TEST_FG1_F1_WITH_STATS
+ yield get_feature_with_stats_and_anomalies_mock
+
+
 @pytest.fixture()
 def mock_base_instantiate_client():
  with patch.object(
@@ -452,6 +464,36 @@ def test_get_feature(get_fg_mock, get_feature_mock):
  )
 
 
+def test_get_feature_with_latest_stats_count(
+ get_fg_mock, get_feature_with_stats_and_anomalies_mock
+):
+ aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+ fg = FeatureGroup(_TEST_FG1_ID)
+ feature = fg.get_feature(_TEST_FG1_F1_ID, latest_stats_count=1)
+
+ get_feature_with_stats_and_anomalies_mock.assert_called_once_with(
+ request=types.featurestore_service_v1beta1.GetFeatureRequest(
+ name=_TEST_FG1_F1_PATH,
+ feature_stats_and_anomaly_spec=types.feature_monitor.FeatureStatsAndAnomalySpec(
+ latest_stats_count=1
+ ),
+ )
+ )
+
+ 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,
+ feature_stats_and_anomalies=[_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY],
+ )
+
+
 def test_get_feature_credentials_set_in_init(mock_base_instantiate_client):
  credentials = mock.MagicMock(spec=auth_credentials.Credentials)
  aiplatform.init(

vertexai/resources/preview/feature_store/feature.py

@@ -16,12 +16,15 @@
 #
 
 import re
-from typing import Optional
+from typing import List, Optional
 from google.auth import credentials as auth_credentials
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform.compat.types import (
  feature as gca_feature,
+ feature_monitor_v1beta1 as gca_feature_monitor,
+ feature_v1beta1 as gca_feature_v1beta1,
+ featurestore_service_v1beta1 as gca_featurestore_service_v1beta1,
 )
 
 
@@ -44,6 +47,7 @@ def __init__(
  feature_group_id: Optional[str] = None,
  project: Optional[str] = None,
  location: Optional[str] = None,
+ latest_stats_count: Optional[int] = None,
  credentials: Optional[auth_credentials.Credentials] = None,
  ):
  """Retrieves an existing managed feature.
@@ -62,6 +66,9 @@ def __init__(
  location:
  Location to retrieve feature from. If not set, the location set
  in aiplatform.init will be used.
+ gca_feature_arg:
+ The GCA feature object.
+ Only set when calling from get_feature with latest_stats_count set.
  credentials:
  Custom credentials to use to retrieve this feature. Overrides
  credentials set in aiplatform.init.
@@ -102,7 +109,24 @@ def __init__(
 
  feature = f"{feature_group_path}/features/{name}"
 
- self._gca_resource = self._get_gca_resource(resource_name=feature)
+ if latest_stats_count is not None:
+ api_client = self.__class__._instantiate_client(
+ location=location, credentials=credentials
+ )
+
+ feature_obj: gca_feature_v1beta1.Feature = api_client.select_version(
+ "v1beta1"
+ ).get_feature(
+ request=gca_featurestore_service_v1beta1.GetFeatureRequest(
+ name=f"{feature}",
+ feature_stats_and_anomaly_spec=gca_feature_monitor.FeatureStatsAndAnomalySpec(
+ latest_stats_count=latest_stats_count
+ ),
+ )
+ )
+ self._gca_resource = feature_obj
+ else:
+ self._gca_resource = self._get_gca_resource(resource_name=feature)
 
  @property
  def version_column_name(self) -> str:
@@ -118,3 +142,10 @@ def description(self) -> str:
  def point_of_contact(self) -> str:
  """The point of contact for the feature."""
  return self._gca_resource.point_of_contact
+
+ @property
+ def feature_stats_and_anomalies(
+ self,
+ ) -> List[gca_feature_monitor.FeatureStatsAndAnomaly]:
+ """The number of latest stats to return. Only present when gca_feature is set."""
+ return self._gca_resource.feature_stats_and_anomaly

vertexai/resources/preview/feature_store/feature_group.py

@@ -239,6 +239,7 @@ def delete(self, force: bool = False, sync: bool = True) -> None:
  def get_feature(
  self,
  feature_id: str,
+ latest_stats_count: Optional[int] = None,
  credentials: Optional[auth_credentials.Credentials] = None,
  ) -> Feature:
  """Retrieves an existing managed feature.
@@ -257,6 +258,12 @@ def get_feature(
  credentials = (
  credentials or self.credentials or initializer.global_config.credentials
  )
+ if latest_stats_count is not None:
+ return Feature(
+ name=f"{self.resource_name}/features/{feature_id}",
+ latest_stats_count=latest_stats_count,
+ credentials=credentials,
+ )
  return Feature(
  f"{self.resource_name}/features/{feature_id}", credentials=credentials
  )