Skip to content

Commit 4bc5fb6

Browse files
authored
feature: Add methods for feature group to list feature metadata parameters and tags (#3654)
1 parent c4c4e83 commit 4bc5fb6

File tree

4 files changed

+92
-3
lines changed

4 files changed

+92
-3
lines changed

doc/api/prep_data/feature_store.rst

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
Feature Store APIs
22
------------------
33

4-
Feature group
4+
Feature Group
55
*************
66

77
.. autoclass:: sagemaker.feature_store.feature_group.FeatureGroup
@@ -18,7 +18,7 @@ Feature group
1818
:show-inheritance:
1919

2020

21-
Feature definition
21+
Feature Definition
2222
******************
2323

2424
.. autoclass:: sagemaker.feature_store.feature_definition.FeatureDefinition
@@ -77,10 +77,46 @@ Inputs
7777
:members:
7878
:show-inheritance:
7979

80+
.. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
81+
:members:
82+
:show-inheritance:
83+
84+
.. autoclass:: sagemaker.feature_store.inputs.SearchOperatorEnum
85+
:members:
86+
:show-inheritance:
87+
88+
.. autoclass:: sagemaker.feature_store.inputs.SortOrderEnum
89+
:members:
90+
:show-inheritance:
91+
92+
.. autoclass:: sagemaker.feature_store.inputs.FilterOperatorEnum
93+
:members:
94+
:show-inheritance:
95+
96+
.. autoclass:: sagemaker.feature_store.inputs.Filter
97+
:members:
98+
:show-inheritance:
99+
100+
.. autoclass:: sagemaker.feature_store.inputs.Identifier
101+
:members:
102+
:show-inheritance:
103+
104+
.. autoclass:: sagemaker.feature_store.inputs.FeatureParameter
105+
:members:
106+
:show-inheritance:
107+
80108

81109
Dataset Builder
82110
***************
83111

84112
.. autoclass:: sagemaker.feature_store.dataset_builder.DatasetBuilder
85113
:members:
86114
:show-inheritance:
115+
116+
117+
Feature Store
118+
*************
119+
120+
.. autoclass:: sagemaker.feature_store.feature_store.FeatureStore
121+
:members:
122+
:show-inheritance:

src/sagemaker/feature_store/feature_group.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -690,6 +690,32 @@ def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]:
690690
feature_group_name=self.name, feature_name=feature_name
691691
)
692692

693+
def list_tags(self) -> Sequence[Dict[str, str]]:
694+
"""List all tags for a feature group.
695+
696+
Returns:
697+
list of key, value pair of the tags.
698+
"""
699+
700+
feature_group_arn = self.sagemaker_session.describe_feature_group(
701+
feature_group_name=self.name
702+
).get("FeatureGroupArn")
703+
704+
return self.sagemaker_session.list_tags(resource_arn=feature_group_arn)
705+
706+
def list_parameters_for_feature_metadata(self, feature_name: str) -> Sequence[Dict[str, str]]:
707+
"""List all parameters for a feature metadata.
708+
709+
Args:
710+
feature_name (str): name of the feature.
711+
Returns:
712+
list of key, value pair of the parameters.
713+
"""
714+
715+
return self.sagemaker_session.describe_feature_metadata(
716+
feature_group_name=self.name, feature_name=feature_name
717+
).get("Parameters")
718+
693719
def load_feature_definitions(
694720
self,
695721
data_frame: DataFrame,

tests/integ/test_feature_store.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,9 +207,10 @@ def test_create_feature_store_online_only(
207207
event_time_feature_name="feature3",
208208
role_arn=role,
209209
enable_online_store=True,
210+
tags=[{"Key": "key1", "Value": "value1"}],
210211
)
211212
_wait_for_feature_group_create(feature_group)
212-
213+
assert feature_group.list_tags() == [{"Key": "key1", "Value": "value1"}]
213214
assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
214215

215216

@@ -520,6 +521,10 @@ def test_feature_metadata(
520521
print(describe_feature_metadata)
521522
assert description == describe_feature_metadata.get("Description")
522523
assert 2 == len(describe_feature_metadata.get("Parameters"))
524+
assert [
525+
{"Key": "key1", "Value": "value1"},
526+
{"Key": "key2", "Value": "value2"},
527+
] == feature_group.list_parameters_for_feature_metadata(feature_name=feature_name)
523528

524529
parameter_removals = ["key1"]
525530
feature_group.update_feature_metadata(
@@ -530,6 +535,9 @@ def test_feature_metadata(
530535
)
531536
assert description == describe_feature_metadata.get("Description")
532537
assert 1 == len(describe_feature_metadata.get("Parameters"))
538+
assert [
539+
{"Key": "key2", "Value": "value2"}
540+
] == feature_group.list_parameters_for_feature_metadata(feature_name=feature_name)
533541

534542

535543
def test_search(feature_store_session, role, feature_group_name, pandas_data_frame):

tests/unit/sagemaker/feature_store/test_feature_group.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,25 @@ def test_load_feature_definition_unsupported_types(sagemaker_session_mock):
333333
assert "Failed to infer Feature type based on dtype bool for column bool." in str(error)
334334

335335

336+
def test_list_tags(sagemaker_session_mock):
337+
feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
338+
sagemaker_session_mock.describe_feature_group.return_value = {"FeatureGroupArn": "test-arn"}
339+
feature_group.list_tags()
340+
sagemaker_session_mock.list_tags.assert_called_with(resource_arn="test-arn")
341+
342+
343+
def test_list_parameters_for_feature_metadata(sagemaker_session_mock):
344+
feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
345+
test_feature_metadata = {"Parameters": [{"Key": "k", "Value": "y"}]}
346+
sagemaker_session_mock.describe_feature_metadata.return_value = test_feature_metadata
347+
assert feature_group.list_parameters_for_feature_metadata(feature_name="feature") == [
348+
{"Key": "k", "Value": "y"}
349+
]
350+
sagemaker_session_mock.describe_feature_metadata.assert_called_with(
351+
feature_group_name="MyFeatureGroup", feature_name="feature"
352+
)
353+
354+
336355
def test_ingest_zero_processes():
337356
feature_group = FeatureGroup(name="MyGroup", sagemaker_session=sagemaker_session_mock)
338357
df = Mock()

0 commit comments

Comments
 (0)