feat: Add to_dataframe method in DatasetBuilder (aws#729)

Author: imingtsou

src/sagemaker/feature_store/dataset_builder.py

@@ -17,7 +17,7 @@
 from __future__ import absolute_import
 
 import datetime
-from typing import Any, Dict, List, Sequence, Union
+from typing import Any, Dict, List, Sequence, Tuple, Union
 
 import attr
 import pandas as pd
@@ -248,7 +248,7 @@ def with_event_time_range(
         self._event_time_ending_timestamp = ending_timestamp
         return self
 
-    def to_csv(self):
+    def to_csv(self) -> Tuple[str, str]:
         """Get query string and result in .csv format
 
         Returns:
@@ -314,6 +314,23 @@ def to_csv(self):
             ), query_result.get("QueryExecution", None).get("Query", None)
         raise ValueError("Base must be either a FeatureGroup or a DataFrame.")
 
+    def to_dataframe(self) -> Tuple[str, pd.DataFrame]:
+        """Get query string and result in pandas.Dataframe
+
+        Returns:
+            The pandas.DataFrame object.
+            The query string executed.
+        """
+        query_string, csv_file = self.to_csv()
+        s3.S3Downloader.download(
+            s3_uri=csv_file,
+            local_path="./",
+            kms_key=self._kms_key_id,
+            sagemaker_session=self._sagemaker_session,
+        )
+        # TODO: do we need to clean up local file?
+        return query_string, pd.read_csv(csv_file.split("/")[-1])
+
     def _construct_where_query_string(self, suffix: str, event_time_identifier_feature_name: str):
         """Internal method for constructing SQL WHERE query string by parameters.
 
@@ -425,6 +442,8 @@ def _construct_query_string(
                 ]
             )
             query_string += join_subquery_string
+        if self._number_of_records:
+            query_string += f"\nLIMIT {self._number_of_records}"
         return query_string
 
     def _construct_join_condition(self, feature_group: FeatureGroupToBeMerged, suffix: str):

tests/unit/sagemaker/feature_store/test_dataset_builder.py

@@ -14,10 +14,12 @@
 
 import datetime
 
+import pandas as pd
 import pytest
 from mock import Mock
 
 from sagemaker.feature_store.dataset_builder import DatasetBuilder
+from sagemaker.feature_store.feature_group import FeatureGroup
 
 
 @pytest.fixture
@@ -30,6 +32,58 @@ def feature_group_mock():
     return Mock()
 
 
+def test_with_feature_group_throw_runtime_error(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group,
+        output_path="file/to/path",
+    )
+    sagemaker_session_mock.describe_feature_group.return_value = {"OfflineStoreConfig": {}}
+    with pytest.raises(RuntimeError) as error:
+        dataset_builder.with_feature_group(
+            feature_group, "target-feature", ["feature-1", "feature-2"]
+        )
+    assert "No metastore is configured with FeatureGroup MyFeatureGroup." in str(error)
+
+
+def test_with_feature_group(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    dataframe = pd.DataFrame({"feature-1": [420, 380, 390], "feature-2": [50, 40, 45]})
+    feature_group.load_feature_definitions(dataframe)
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group,
+        output_path="file/to/path",
+    )
+    sagemaker_session_mock.describe_feature_group.return_value = {
+        "OfflineStoreConfig": {"DataCatalogConfig": {"TableName": "table", "Database": "database"}},
+        "RecordIdentifierFeatureName": "feature-1",
+        "EventTimeFeatureName": "feature-2",
+    }
+    dataset_builder.with_feature_group(feature_group, "target-feature", ["feature-1", "feature-2"])
+    assert len(dataset_builder._feature_groups_to_be_merged) == 1
+    assert dataset_builder._feature_groups_to_be_merged[0].features == ["feature-1", "feature-2"]
+    assert dataset_builder._feature_groups_to_be_merged[0].included_feature_names == [
+        "feature-1",
+        "feature-2",
+    ]
+    assert dataset_builder._feature_groups_to_be_merged[0].database == "database"
+    assert dataset_builder._feature_groups_to_be_merged[0].table_name == "table"
+    assert (
+        dataset_builder._feature_groups_to_be_merged[0].record_identifier_feature_name
+        == "feature-1"
+    )
+    assert (
+        dataset_builder._feature_groups_to_be_merged[0].event_time_identifier_feature_name
+        == "feature-2"
+    )
+    assert (
+        dataset_builder._feature_groups_to_be_merged[0].target_feature_name_in_base
+        == "target-feature"
+    )
+
+
 def test_point_in_time_accurate_join(sagemaker_session_mock, feature_group_mock):
     dataset_builder = DatasetBuilder(
         sagemaker_session=sagemaker_session_mock,