Add delete_record API (aws#664)

Author: imingtsou

src/sagemaker/feature_store/feature_group.py

@@ -671,6 +671,25 @@ def put_record(self, record: Sequence[FeatureValue]):
             feature_group_name=self.name, record=[value.to_dict() for value in record]
         )
 
+    def delete_record(
+        self,
+        record_identifier_value_as_string: str,
+        event_time: str,
+    ):
+        """Delete a single record from a FeatureGroup.
+
+        Args:
+            record_identifier_value_as_string (String):
+                a String representing the value of the record identifier.
+            event_time (String):
+                a timestamp format String indicating when the deletion event occurred.
+        """
+        return self.sagemaker_session.delete_record(
+            feature_group_name=self.name,
+            record_identifier_value_as_string=record_identifier_value_as_string,
+            event_time=event_time,
+        )
+
     def ingest(
         self,
         data_frame: DataFrame,

src/sagemaker/session.py

@@ -4454,6 +4454,25 @@ def put_record(
             Record=record,
         )
 
+    def delete_record(
+        self,
+        feature_group_name: str,
+        record_identifier_value_as_string: str,
+        event_time: str,
+    ):
+        """Deletes a single record from the FeatureGroup.
+
+        Args:
+            feature_group_name (str): name of the FeatureGroup.
+            record_identifier_value_as_string (str): name of the record identifier.
+            event_time (str): a timestamp indicating when the deletion event occurred.
+        """
+        return self.sagemaker_featurestore_runtime_client.delete_record(
+            FeatureGroupName=feature_group_name,
+            RecordIdentifierValueAsString=record_identifier_value_as_string,
+            EventTime=event_time,
+        )
+
     def get_record(
         self,
         record_identifier_value_as_string: str,

tests/integ/test_feature_store.py

@@ -14,6 +14,7 @@
 
 import json
 import time
+import datetime
 from contextlib import contextmanager
 
 import boto3
@@ -336,6 +337,45 @@ def test_get_record(
         assert retrieved_record is None
 
 
+def test_delete_record(
+    feature_store_session,
+    role,
+    feature_group_name,
+    pandas_data_frame,
+    record,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    record_identifier_value_as_string = record[0].value_as_string
+    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)
+        # Ingest data
+        feature_group.put_record(record=record)
+        # Retrieve data
+        retrieved_record = feature_group.get_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+        )
+        assert retrieved_record is not None
+        # Delete data
+        feature_group.delete_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+            event_time=datetime.datetime.now().replace(microsecond=0).isoformat() + "Z",
+        )
+        # Retrieve data
+        retrieved_record = feature_group.get_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+        )
+        assert retrieved_record is None
+
+
 def test_update_feature_group(
     feature_store_session,
     role,

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -225,6 +225,21 @@ def test_put_record(sagemaker_session_mock):
     )
 
 
+def test_delete_record(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    record_identifier_value_as_string = "1.0"
+    event_time = "2022-09-14"
+    feature_group.delete_record(
+        record_identifier_value_as_string=record_identifier_value_as_string,
+        event_time=event_time,
+    )
+    sagemaker_session_mock.delete_record.assert_called_with(
+        feature_group_name="MyFeatureGroup",
+        record_identifier_value_as_string=record_identifier_value_as_string,
+        event_time=event_time,
+    )
+
+
 def test_load_feature_definition(sagemaker_session_mock):
     feature_group = FeatureGroup(name="SomeGroup", sagemaker_session=sagemaker_session_mock)
     df = pd.DataFrame(