Add list_feature_groups API (aws#647)

Author: imingtsou

src/sagemaker/feature_store/feature_store.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+"""Feature Store.
+
+Amazon SageMaker Feature Store is a fully managed, purpose-built repository to store, share, and
+manage features for machine learning (ML) models.
+"""
+from __future__ import absolute_import
+
+import datetime
+from typing import Dict, Any
+
+import attr
+
+from sagemaker import Session
+
+
+@attr.s
+class FeatureStore:
+    """FeatureStore definition.
+
+    This class instantiates a FeatureStore object that comprises a SageMaker session instance.
+
+    Attributes:
+        sagemaker_session (Session): session instance to perform boto calls.
+    """
+
+    sagemaker_session: Session = attr.ib(default=Session)
+
+    def list_feature_groups(
+        self,
+        name_contains: str = None,
+        feature_group_status_equals: str = None,
+        offline_store_status_equals: str = None,
+        creation_time_after: datetime.datetime = None,
+        creation_time_before: datetime.datetime = None,
+        sort_order: str = None,
+        sort_by: str = None,
+        max_results: int = None,
+        next_token: str = None,
+    ) -> Dict[str, Any]:
+        """List all FeatureGroups satisfying given filters.
+
+        Args:
+            name_contains (str): A string that partially matches one or more FeatureGroups' names.
+                Filters FeatureGroups by name.
+            feature_group_status_equals (str): A FeatureGroup status.
+                Filters FeatureGroups by FeatureGroup status.
+            offline_store_status_equals (str): An OfflineStore status.
+                Filters FeatureGroups by OfflineStore status.
+            creation_time_after (datetime.datetime): Use this parameter to search for FeatureGroups
+                created after a specific date and time.
+            creation_time_before (datetime.datetime): Use this parameter to search for FeatureGroups
+                created before a specific date and time.
+            sort_order (str): The order in which FeatureGroups are listed.
+            sort_by (str): The value on which the FeatureGroup list is sorted.
+            max_results (int): The maximum number of results returned by ListFeatureGroups.
+            next_token (str): A token to resume pagination of ListFeatureGroups results.
+        Returns:
+            Response dict from service.
+        """
+        return self.sagemaker_session.list_feature_groups(
+            name_contains=name_contains,
+            feature_group_status_equals=feature_group_status_equals,
+            offline_store_status_equals=offline_store_status_equals,
+            creation_time_after=creation_time_after,
+            creation_time_before=creation_time_before,
+            sort_order=sort_order,
+            sort_by=sort_by,
+            max_results=max_results,
+            next_token=next_token,
+        )

src/sagemaker/session.py

@@ -4332,6 +4332,61 @@ def update_feature_group(
             FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
         )
 
+    def list_feature_groups(
+        self,
+        name_contains,
+        feature_group_status_equals,
+        offline_store_status_equals,
+        creation_time_after,
+        creation_time_before,
+        sort_order,
+        sort_by,
+        max_results,
+        next_token,
+    ) -> Dict[str, Any]:
+        """List all FeatureGroups satisfying given filters.
+
+        Args:
+            name_contains (str): A string that partially matches one or more FeatureGroups' names.
+                Filters FeatureGroups by name.
+            feature_group_status_equals (str): A FeatureGroup status.
+                Filters FeatureGroups by FeatureGroup status.
+            offline_store_status_equals (str): An OfflineStore status.
+                Filters FeatureGroups by OfflineStore status.
+            creation_time_after (datetime.datetime): Use this parameter to search for FeatureGroups
+                created after a specific date and time.
+            creation_time_before (datetime.datetime): Use this parameter to search for FeatureGroups
+                created before a specific date and time.
+            sort_order (str): The order in which FeatureGroups are listed.
+            sort_by (str): The value on which the FeatureGroup list is sorted.
+            max_results (int): The maximum number of results returned by ListFeatureGroups.
+            next_token (str): A token to resume pagination of ListFeatureGroups results.
+        Returns:
+            Response dict from service.
+        """
+        list_feature_groups_args = {}
+
+        if name_contains:
+            list_feature_groups_args["NameContains"] = name_contains
+        if feature_group_status_equals:
+            list_feature_groups_args["FeatureGroupStatusEquals"] = feature_group_status_equals
+        if offline_store_status_equals:
+            list_feature_groups_args["OfflineStoreStatusEquals"] = offline_store_status_equals
+        if creation_time_after:
+            list_feature_groups_args["CreationTimeAfter"] = creation_time_after
+        if creation_time_before:
+            list_feature_groups_args["CreationTimeBefore"] = creation_time_before
+        if sort_order:
+            list_feature_groups_args["SortOrder"] = sort_order
+        if sort_by:
+            list_feature_groups_args["SortBy"] = sort_by
+        if max_results:
+            list_feature_groups_args["MaxResults"] = max_results
+        if next_token:
+            list_feature_groups_args["NextToken"] = next_token
+
+        return self.sagemaker_client.list_feature_groups(**list_feature_groups_args)
+
     def update_feature_metadata(
         self,
         feature_group_name: str,

tests/integ/test_feature_store.py

@@ -24,6 +24,7 @@
 
 from sagemaker.feature_store.feature_definition import FractionalFeatureDefinition
 from sagemaker.feature_store.feature_group import FeatureGroup
+from sagemaker.feature_store.feature_store import FeatureStore
 from sagemaker.feature_store.inputs import FeatureValue, FeatureParameter, TableFormatEnum
 from sagemaker.session import get_execution_role, Session
 from tests.integ.timeout import timeout
@@ -316,6 +317,25 @@ def test_update_feature_group(
         assert any([True for elem in feature_definitions if new_feature_name in elem.values()])
 
 
+def test_list_feature_groups(feature_store_session, role, feature_group_name, pandas_data_frame):
+    feature_store = FeatureStore(sagemaker_session=feature_store_session)
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=False,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+        )
+        _wait_for_feature_group_create(feature_group)
+        output = feature_store.list_feature_groups(name_contains=feature_group_name)
+
+    assert output["FeatureGroupSummaries"][0]["FeatureGroupName"] == feature_group_name
+
+
 def test_feature_metadata(
     feature_store_session,
     role,