Add methods for feature group to list feature metadata parameters and tags

Author: can-sun

doc/api/prep_data/feature_store.rst

@@ -1,7 +1,7 @@
 Feature Store APIs
 ------------------
 
-Feature group
+Feature Group
 *************
 
 .. autoclass:: sagemaker.feature_store.feature_group.FeatureGroup
@@ -18,7 +18,7 @@ Feature group
     :show-inheritance:
 
 
-Feature definition
+Feature Definition
 ******************
 
 .. autoclass:: sagemaker.feature_store.feature_definition.FeatureDefinition
@@ -77,10 +77,46 @@ Inputs
     :members:
     :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.SearchOperatorEnum
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.SortOrderEnum
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.FilterOperatorEnum
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.Filter
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.Identifier
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.FeatureParameter
+    :members:
+    :show-inheritance:
+
 
 Dataset Builder
 ***************
 
 .. autoclass:: sagemaker.feature_store.dataset_builder.DatasetBuilder
     :members:
     :show-inheritance:
+
+
+Feature Store
+*************
+
+.. autoclass:: sagemaker.feature_store.feature_store.FeatureStore
+    :members:
+    :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -664,6 +664,32 @@ def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]:
             feature_group_name=self.name, feature_name=feature_name
         )
 
+    def list_tags(self) -> Sequence[Dict[str, str]]:
+        """List all tags for a feature group.
+
+        Returns:
+            list of key, value pair of the tags.
+        """
+
+        feature_group_arn = self.sagemaker_session.describe_feature_group(
+            feature_group_name=self.name
+        ).get("FeatureGroupArn")
+
+        return self.sagemaker_session.list_tags(resource_arn=feature_group_arn)
+
+    def list_parameters_for_feature_metadata(self, feature_name: str) -> Sequence[Dict[str, str]]:
+        """List all parameters for a feature metadata.
+
+        Args:
+            feature_name (str): name of the feature.
+        Returns:
+            list of key, value pair of the parameters.
+        """
+
+        return self.sagemaker_session.describe_feature_metadata(
+            feature_group_name=self.name, feature_name=feature_name
+        ).get("Parameters")
+
     def load_feature_definitions(
         self,
         data_frame: DataFrame,

tests/integ/test_feature_store.py

@@ -203,9 +203,10 @@ def test_create_feature_store_online_only(
             event_time_feature_name="feature3",
             role_arn=role,
             enable_online_store=True,
+            tags=[{"Key": "key1", "Value": "value1"}],
         )
         _wait_for_feature_group_create(feature_group)
-
+        assert feature_group.list_tags() == [{"Key": "key1", "Value": "value1"}]
     assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
 
 
@@ -516,6 +517,10 @@ def test_feature_metadata(
         print(describe_feature_metadata)
         assert description == describe_feature_metadata.get("Description")
         assert 2 == len(describe_feature_metadata.get("Parameters"))
+        assert [
+            {"Key": "key1", "Value": "value1"},
+            {"Key": "key2", "Value": "value2"},
+        ] == feature_group.list_parameters_for_feature_metadata(feature_name=feature_name)
 
         parameter_removals = ["key1"]
         feature_group.update_feature_metadata(
@@ -526,6 +531,9 @@ def test_feature_metadata(
         )
         assert description == describe_feature_metadata.get("Description")
         assert 1 == len(describe_feature_metadata.get("Parameters"))
+        assert [
+            {"Key": "key2", "Value": "value2"}
+        ] == feature_group.list_parameters_for_feature_metadata(feature_name=feature_name)
 
 
 def test_search(feature_store_session, role, feature_group_name, pandas_data_frame):

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -274,6 +274,25 @@ def test_load_feature_definition_unsupported_types(sagemaker_session_mock):
     assert "Failed to infer Feature type based on dtype bool for column bool." in str(error)
 
 
+def test_list_tags(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    sagemaker_session_mock.describe_feature_group.return_value = {"FeatureGroupArn": "test-arn"}
+    feature_group.list_tags()
+    sagemaker_session_mock.list_tags.assert_called_with(resource_arn="test-arn")
+
+
+def test_list_parameters_for_feature_metadata(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    test_feature_metadata = {"Parameters": [{"Key": "k", "Value": "y"}]}
+    sagemaker_session_mock.describe_feature_metadata.return_value = test_feature_metadata
+    assert feature_group.list_parameters_for_feature_metadata(feature_name="feature") == [
+        {"Key": "k", "Value": "y"}
+    ]
+    sagemaker_session_mock.describe_feature_metadata.assert_called_with(
+        feature_group_name="MyFeatureGroup", feature_name="feature"
+    )
+
+
 def test_ingest_zero_processes():
     feature_group = FeatureGroup(name="MyGroup", sagemaker_session=sagemaker_session_mock)
     df = Mock()