feat: Feature/get record api (aws#650)

Co-authored-by: Eric Zou <[email protected]>

Author: 2 people

src/sagemaker/feature_store/feature_group.py

@@ -644,6 +644,23 @@ def load_feature_definitions(
         self.feature_definitions = feature_definitions
         return self.feature_definitions
 
+    def get_record(
+        self, record_identifier_value_as_string: str, feature_names: Sequence[str] = None
+    ) -> Sequence[Dict[str, str]]:
+        """Get a single record in a FeatureGroup
+
+        Args:
+            record_identifier_value_as_string (String):
+                a String representing the value of the record identifier.
+            feature_names (Sequence[String]):
+                a list of Strings representing feature names.
+        """
+        return self.sagemaker_session.get_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+            feature_group_name=self.name,
+            feature_names=feature_names,
+        ).get("Record")
+
     def put_record(self, record: Sequence[FeatureValue]):
         """Put a single record in the FeatureGroup.
 

src/sagemaker/session.py

@@ -4454,6 +4454,29 @@ def put_record(
             Record=record,
         )
 
+    def get_record(
+        self,
+        record_identifier_value_as_string: str,
+        feature_group_name: str,
+        feature_names: Sequence[str],
+    ) -> Dict[str, Sequence[Dict[str, str]]]:
+        """Gets a single record in the FeatureGroup.
+
+        Args:
+            record_identifier_value_as_string (str): name of the record identifier.
+            feature_group_name (str): name of the FeatureGroup.
+            feature_names (Sequence[str]): list of feature names.
+        """
+        get_record_args = {
+            "FeatureGroupName": feature_group_name,
+            "RecordIdentifierValueAsString": record_identifier_value_as_string,
+        }
+
+        if feature_names:
+            get_record_args["FeatureNames"] = feature_names
+
+        return self.sagemaker_featurestore_runtime_client.get_record(**get_record_args)
+
     def start_query_execution(
         self,
         catalog: str,

tests/integ/test_feature_store.py

@@ -271,7 +271,8 @@ def test_create_feature_group_glue_table_format(
     offline_store_s3_uri,
     pandas_data_frame,
 ):
-    feature_group = FeatureGroup(name=feature_group_name, 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):
@@ -288,6 +289,52 @@ def test_create_feature_group_glue_table_format(
         table_format = feature_group.describe().get("OfflineStoreConfig").get("TableFormat")
         assert table_format == "Glue"
 
+def test_get_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,
+        )
+        record_names = list(map(lambda r: r.feature_name, record))
+        assert len(retrieved_record) == len(record_names)
+        for feature in retrieved_record:
+            assert feature["FeatureName"] in record_names
+        removed_feature_name = record_names.pop()
+        # Retrieve data
+        retrieved_record = feature_group.get_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+            feature_names=record_names,
+        )
+        assert len(retrieved_record) == len(record_names)
+        for feature in retrieved_record:
+            assert feature["FeatureName"] in record_names
+            assert feature["FeatureName"] is not removed_feature_name
+        # Retrieve data
+        retrieved_record = feature_group.get_record(
+            record_identifier_value_as_string="1.0",
+        )
+        assert retrieved_record is None
+
 
 def test_update_feature_group(
     feature_store_session,

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -201,6 +201,22 @@ def test_feature_metadata_describe(sagemaker_session_mock):
     )
 
 
+def test_get_record(sagemaker_session_mock):
+    feature_group_name = "MyFeatureGroup"
+    feature_names = ["MyFeature1", "MyFeature2"]
+    record_identifier_value_as_string = "1.0"
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=sagemaker_session_mock)
+    feature_group.get_record(
+        record_identifier_value_as_string=record_identifier_value_as_string,
+        feature_names=feature_names,
+    )
+    sagemaker_session_mock.get_record.assert_called_with(
+        feature_group_name=feature_group_name,
+        record_identifier_value_as_string=record_identifier_value_as_string,
+        feature_names=feature_names,
+    )
+
+
 def test_put_record(sagemaker_session_mock):
     feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
     feature_group.put_record(record=[])