Revert "Add batch_get_record and search API for FeatureStore (#3580)"

This reverts commit 27ebde6.

Author: trajanikant

src/sagemaker/feature_store/feature_store.py

@@ -26,13 +26,6 @@
 from sagemaker import Session
 from sagemaker.feature_store.dataset_builder import DatasetBuilder
 from sagemaker.feature_store.feature_group import FeatureGroup
-from sagemaker.feature_store.inputs import (
-    Filter,
-    ResourceEnum,
-    SearchOperatorEnum,
-    SortOrderEnum,
-    Identifier,
-)
 
 
 @attr.s
@@ -121,7 +114,6 @@ def list_feature_groups(
             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.
         """
@@ -136,61 +128,3 @@ def list_feature_groups(
             max_results=max_results,
             next_token=next_token,
         )
-
-    def batch_get_record(self, identifiers: Sequence[Identifier]) -> Dict[str, Any]:
-        """Get record in batch from FeatureStore
-
-        Args:
-            identifiers (Sequence[Identifier]): A list of identifiers to uniquely identify records
-                in FeatureStore.
-
-        Returns:
-            Response dict from service.
-        """
-        batch_get_record_identifiers = [identifier.to_dict() for identifier in identifiers]
-        return self.sagemaker_session.batch_get_record(identifiers=batch_get_record_identifiers)
-
-    def search(
-        self,
-        resource: ResourceEnum,
-        filters: Sequence[Filter] = None,
-        operator: SearchOperatorEnum = None,
-        sort_by: str = None,
-        sort_order: SortOrderEnum = None,
-        next_token: str = None,
-        max_results: int = None,
-    ) -> Dict[str, Any]:
-        """Search for FeatureGroups or FeatureMetadata satisfying given filters.
-
-        Args:
-            resource (ResourceEnum): The name of the Amazon SageMaker resource to search for.
-                Valid values are ``FeatureGroup`` or ``FeatureMetadata``.
-            filters (Sequence[Filter]): A list of filter objects (Default: None).
-            operator (SearchOperatorEnum): A Boolean operator used to evaluate the filters.
-                Valid values are ``And`` or ``Or``. The default is ``And`` (Default: None).
-            sort_by (str): The name of the resource property used to sort the ``SearchResults``.
-                The default is ``LastModifiedTime``.
-            sort_order (SortOrderEnum): How ``SearchResults`` are ordered.
-                Valid values are ``Ascending`` or ``Descending``. The default is ``Descending``.
-            next_token (str): If more than ``MaxResults`` resources match the specified
-                filters, the response includes a ``NextToken``. The ``NextToken`` can be passed to
-                the next ``SearchRequest`` to continue retrieving results (Default: None).
-            max_results (int): The maximum number of results to return (Default: None).
-
-        Returns:
-            Response dict from service.
-        """
-        search_expression = {}
-        if filters:
-            search_expression["Filters"] = [filter.to_dict() for filter in filters]
-        if operator:
-            search_expression["Operator"] = str(operator)
-
-        return self.sagemaker_session.search(
-            resource=str(resource),
-            search_expression=search_expression,
-            sort_by=sort_by,
-            sort_order=None if not sort_order else str(sort_order),
-            next_token=next_token,
-            max_results=max_results,
-        )

src/sagemaker/feature_store/inputs.py

@@ -29,7 +29,7 @@
 from __future__ import absolute_import
 
 import abc
-from typing import Dict, Any, List
+from typing import Dict, Any
 from enum import Enum
 
 import attr
@@ -245,127 +245,3 @@ def to_dict(self) -> Dict[str, Any]:
             Key=self.key,
             Value=self.value,
         )
-
-
-class ResourceEnum(Enum):
-    """Enum of resources.
-
-    The data type of resource can be ``FeatureGroup`` or ``FeatureMetadata``.
-    """
-
-    def __str__(self):
-        """Override str method to return enum value."""
-        return str(self.value)
-
-    FEATURE_GROUP = "FeatureGroup"
-    FEATURE_METADATA = "FeatureMetadata"
-
-
-class SearchOperatorEnum(Enum):
-    """Enum of search operators.
-
-    The data type of search operator can be ``And`` or ``Or``.
-    """
-
-    def __str__(self):
-        """Override str method to return enum value."""
-        return str(self.value)
-
-    AND = "And"
-    OR = "Or"
-
-
-class SortOrderEnum(Enum):
-    """Enum of sort orders.
-
-    The data type of sort order can be ``Ascending`` or ``Descending``.
-    """
-
-    def __str__(self):
-        """Override str method to return enum value."""
-        return str(self.value)
-
-    ASCENDING = "Ascending"
-    DESCENDING = "Descending"
-
-
-class FilterOperatorEnum(Enum):
-    """Enum of filter operators.
-
-    The data type of filter operator can be ``Equals``, ``NotEquals``, ``GreaterThan``,
-    ``GreaterThanOrEqualTo``, ``LessThan``, ``LessThanOrEqualTo``, ``Contains``, ``Exists``,
-    ``NotExists``, or ``In``.
-    """
-
-    def __str__(self):
-        """Override str method to return enum value."""
-        return str(self.value)
-
-    EQUALS = "Equals"
-    NOT_EQUALS = "NotEquals"
-    GREATER_THAN = "GreaterThan"
-    GREATER_THAN_OR_EQUAL_TO = "GreaterThanOrEqualTo"
-    LESS_THAN = "LessThan"
-    LESS_THAN_OR_EQUAL_TO = "LessThanOrEqualTo"
-    CONTAINS = "Contains"
-    EXISTS = "Exists"
-    NOT_EXISTS = "NotExists"
-    IN = "In"
-
-
-@attr.s
-class Filter(Config):
-    """Filter for FeatureStore search.
-
-    Attributes:
-        name (str): A resource property name.
-        value (str): A value used with ``Name`` and ``Operator`` to determine which resources
-            satisfy the filter's condition.
-        operator (FilterOperatorEnum): A Boolean binary operator that is used to evaluate the
-        filter. If specify ``Value`` without ``Operator``, Amazon SageMaker uses ``Equals``
-        (default: None).
-    """
-
-    name: str = attr.ib()
-    value: str = attr.ib()
-    operator: FilterOperatorEnum = attr.ib(default=None)
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Construct a dictionary based on the attributes provided.
-
-        Returns:
-            dict represents the attributes.
-        """
-        return Config.construct_dict(
-            Name=self.name,
-            Value=self.value,
-            Operator=None if not self.operator else str(self.operator),
-        )
-
-
-@attr.s
-class Identifier(Config):
-    """Identifier of batch get record API.
-
-    Attributes:
-        feature_group_name (str): name of a feature group.
-        record_identifiers_value_as_string (List[str]): string value of record identifier.
-        feature_names (List[str]): list of feature names (default: None).
-    """
-
-    feature_group_name: str = attr.ib()
-    record_identifiers_value_as_string: List[str] = attr.ib()
-    feature_names: List[str] = attr.ib(default=None)
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Construct a dictionary based on the attributes provided.
-
-        Returns:
-            dict represents the attributes.
-        """
-
-        return Config.construct_dict(
-            FeatureGroupName=self.feature_group_name,
-            RecordIdentifiersValueAsString=self.record_identifiers_value_as_string,
-            FeatureNames=None if not self.feature_names else self.feature_names,
-        )

src/sagemaker/session.py

@@ -4536,48 +4536,6 @@ def describe_feature_metadata(
             FeatureGroupName=feature_group_name, FeatureName=feature_name
         )
 
-    def search(
-        self,
-        resource: str,
-        search_expression: Dict[str, any] = None,
-        sort_by: str = None,
-        sort_order: str = None,
-        next_token: str = None,
-        max_results: int = None,
-    ) -> Dict[str, Any]:
-        """Search for SageMaker resources satisfying given filters.
-
-        Args:
-            resource (str): The name of the Amazon SageMaker resource to search for.
-            search_expression (Dict[str, any]): A Boolean conditional statement. Resources must
-                satisfy this condition to be included in search results.
-            sort_by (str): The name of the resource property used to sort the ``SearchResults``.
-                The default is ``LastModifiedTime``.
-            sort_order (str): How ``SearchResults`` are ordered.
-                Valid values are ``Ascending`` or ``Descending``. The default is ``Descending``.
-                next_token (str): If more than ``MaxResults`` resources match the specified
-                ``SearchExpression``, the response includes a ``NextToken``. The ``NextToken`` can
-                be passed to the next ``SearchRequest`` to continue retrieving results.
-            max_results (int): The maximum number of results to return.
-
-        Returns:
-            Response dict from service.
-        """
-        search_args = {"Resource": resource}
-
-        if search_expression:
-            search_args["SearchExpression"] = search_expression
-        if sort_by:
-            search_args["SortBy"] = sort_by
-        if sort_order:
-            search_args["SortOrder"] = sort_order
-        if next_token:
-            search_args["NextToken"] = next_token
-        if max_results:
-            search_args["MaxResults"] = max_results
-
-        return self.sagemaker_client.search(**search_args)
-
     def put_record(
         self,
         feature_group_name: str,
@@ -4637,20 +4595,6 @@ def get_record(
 
         return self.sagemaker_featurestore_runtime_client.get_record(**get_record_args)
 
-    def batch_get_record(self, identifiers: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
-        """Gets a batch of record from FeatureStore.
-
-        Args:
-            identifiers (Sequence[Dict[str, Any]]): list of identifiers to uniquely identify records
-                in FeatureStore.
-
-        Returns:
-            Response dict from service.
-        """
-        batch_get_record_args = {"Identifiers": identifiers}
-
-        return self.sagemaker_featurestore_runtime_client.batch_get_record(**batch_get_record_args)
-
     def start_query_execution(
         self,
         catalog: str,

tests/integ/test_feature_store.py

@@ -26,14 +26,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,
-    Filter,
-    ResourceEnum,
-    Identifier,
-)
+from sagemaker.feature_store.inputs import FeatureValue, FeatureParameter, TableFormatEnum
 from sagemaker.session import get_execution_role, Session
 from tests.integ.timeout import timeout
 
@@ -328,7 +321,7 @@ def test_create_feature_group_glue_table_format(
         assert table_format == "Glue"
 
 
-def test_get_and_batch_get_record(
+def test_get_record(
     feature_store_session,
     role,
     feature_group_name,
@@ -374,24 +367,6 @@ def test_get_and_batch_get_record(
         )
         assert retrieved_record is None
 
-        # Retrieve data using batch_get_record
-        feature_store = FeatureStore(sagemaker_session=feature_store_session)
-        records = feature_store.batch_get_record(
-            identifiers=[
-                Identifier(
-                    feature_group_name=feature_group_name,
-                    record_identifiers_value_as_string=[record_identifier_value_as_string],
-                    feature_names=record_names,
-                )
-            ]
-        )["Records"]
-        assert records[0]["FeatureGroupName"] == feature_group_name
-        assert records[0]["RecordIdentifierValueAsString"] == record_identifier_value_as_string
-        assert len(records[0]["Record"]) == len(record_names)
-        for feature in records[0]["Record"]:
-            assert feature["FeatureName"] in record_names
-            assert feature["FeatureName"] is not removed_feature_name
-
 
 def test_delete_record(
     feature_store_session,
@@ -528,28 +503,6 @@ def test_feature_metadata(
         assert 1 == len(describe_feature_metadata.get("Parameters"))
 
 
-def test_search(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.search(
-            resource=ResourceEnum.FEATURE_GROUP,
-            filters=[Filter(name="FeatureGroupName", value=feature_group_name)],
-        )
-    print(output)
-    assert output["Results"][0]["FeatureGroup"]["FeatureGroupName"] == feature_group_name
-
-
 def test_ingest_without_string_feature(
     feature_store_session,
     role,