feat: Add to_csv method in DatasetBuilder (aws#699)

Author: imingtsou

src/sagemaker/feature_store/dataset_builder.py

@@ -22,6 +22,7 @@
 import attr
 import pandas as pd
 
+from sagemaker import Session
 from sagemaker.feature_store.feature_group import FeatureGroup
 
 
@@ -33,6 +34,7 @@ class DatasetBuilder:
     an output path and a KMS key ID.
 
     Attributes:
+        _sagemaker_session (Session): Session instance to perform boto calls.
         _base (Union[FeatureGroup, DataFrame]): A base which can be either a FeatureGroup or a
             pandas.DataFrame and will be used to merge other FeatureGroups and generate a Dataset.
         _output_path (str): An S3 URI which stores the output .csv file.
@@ -59,6 +61,7 @@ class DatasetBuilder:
             dataset will be before it.
     """
 
+    _sagemaker_session: Session = attr.ib()
     _base: Union[FeatureGroup, pd.DataFrame] = attr.ib()
     _output_path: str = attr.ib()
     _record_identifier_feature_name: str = attr.ib(default=None)
@@ -155,3 +158,104 @@ def with_event_time_range(
         self._event_time_starting_timestamp = starting_timestamp
         self._event_time_ending_timestamp = ending_timestamp
         return self
+
+    def to_csv(self):
+        """Get query string and result in .csv format
+
+        Returns:
+            The S3 path of the .csv file.
+            The query string executed.
+        """
+        if isinstance(self._base, FeatureGroup):
+            # TODO: handle pagination and input feature validation
+            base_feature_group = self._base.describe()
+            data_catalog_config = base_feature_group.get("OfflineStoreConfig", None).get(
+                "DataCatalogConfig", None
+            )
+            if not data_catalog_config:
+                raise RuntimeError("No metastore is configured with the base FeatureGroup.")
+            disable_glue = base_feature_group.get("DisableGlueTableCreation", False)
+            self._record_identifier_feature_name = base_feature_group.get(
+                "RecordIdentifierFeatureName", None
+            )
+            self._event_time_identifier_feature_name = base_feature_group.get(
+                "EventTimeFeatureName", None
+            )
+            base_features = [
+                feature.get("FeatureName", None)
+                for feature in base_feature_group.get("FeatureDefinitions", None)
+            ]
+
+            query = self._sagemaker_session.start_query_execution(
+                catalog=data_catalog_config.get("Catalog", None)
+                if disable_glue
+                else "AwsDataCatalog",
+                database=data_catalog_config.get("Database", None),
+                query_string=self._construct_query_string(
+                    data_catalog_config.get("TableName", None),
+                    data_catalog_config.get("Database", None),
+                    base_features,
+                ),
+                output_location=self._output_path,
+                kms_key=self._kms_key_id,
+            )
+            query_id = query.get("QueryExecutionId", None)
+            self._sagemaker_session.wait_for_athena_query(
+                query_execution_id=query_id,
+            )
+            query_state = (
+                self._sagemaker_session.get_query_execution(
+                    query_execution_id=query_id,
+                )
+                .get("QueryExecution", None)
+                .get("Status", None)
+                .get("State", None)
+            )
+            if query_state != "SUCCEEDED":
+                raise RuntimeError(f"Failed to execute query {query_id}.")
+
+            return query_state.get("QueryExecution", None).get("ResultConfiguration", None).get(
+                "OutputLocation", None
+            ), query_state.get("QueryExecution", None).get("Query", None)
+        raise ValueError("Base must be either a FeatureGroup or a DataFrame.")
+
+    def _construct_query_string(
+        self, base_table_name: str, database: str, base_features: list
+    ) -> str:
+        """Internal method for constructing SQL query string by parameters.
+
+        Args:
+            base_table_name (str): The Athena table name of base FeatureGroup or pandas.DataFrame.
+            database (str): The Athena database of the base table.
+            base_features (list): The list of features of the base table.
+        Returns:
+            The query string.
+        """
+        included_features = ", ".join(
+            [
+                f'base."{include_feature_name}"'
+                for include_feature_name in self._included_feature_names
+            ]
+        )
+        query_string = f"SELECT {included_features}\n"
+        if self._include_duplicated_records:
+            query_string += f'FROM "{database}"."{base_table_name}" base\n'
+            if not self._include_deleted_records:
+                query_string += "WHERE NOT is_deleted\n"
+        else:
+            base_features.remove(self._event_time_identifier_feature_name)
+            dedup_features = ", ".join([f'dedup_base."{feature}"' for feature in base_features])
+            query_string += (
+                "FROM (\n"
+                + "SELECT *, row_number() OVER (\n"
+                + f"PARTITION BY {dedup_features}\n"
+                + f'ORDER BY dedup_base."{self._event_time_identifier_feature_name}" '
+                + 'DESC, dedup_base."api_invocation_time" DESC, dedup_base."write_time" DESC\n'
+                + ") AS row_base\n"
+                + f'FROM "{database}"."{base_table_name}" dedup_base\n'
+                + ") AS base\n"
+                + "WHERE row_base = 1\n"
+            )
+            if not self._include_deleted_records:
+                query_string += "AND NOT is_deleted\n"
+        return query_string

src/sagemaker/feature_store/feature_store.py

@@ -40,8 +40,8 @@ class FeatureStore:
 
     sagemaker_session: Session = attr.ib(default=Session)
 
-    @staticmethod
     def create_dataset(
+        self,
         base: Union[FeatureGroup, pd.DataFrame],
         output_path: str,
         record_identifier_feature_name: str = None,
@@ -76,6 +76,7 @@ def create_dataset(
                     + "identifier feature name if specify DataFrame as base."
                 )
         return DatasetBuilder(
+            self.sagemaker_session,
             base,
             output_path,
             record_identifier_feature_name,

tests/unit/sagemaker/feature_store/test_dataset_builder.py

@@ -20,50 +20,85 @@
 from sagemaker.feature_store.dataset_builder import DatasetBuilder
 
 
+@pytest.fixture
+def sagemaker_session_mock():
+    return Mock()
+
+
 @pytest.fixture
 def feature_group_mock():
     return Mock()
 
 
-def test_point_in_time_accurate_join(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_point_in_time_accurate_join(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     dataset_builder.point_in_time_accurate_join()
     assert dataset_builder._point_in_time_accurate_join
 
 
-def test_include_duplicated_records(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_include_duplicated_records(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     dataset_builder.include_duplicated_records()
     assert dataset_builder._include_duplicated_records
 
 
-def test_include_deleted_records(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_include_deleted_records(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     dataset_builder.include_deleted_records()
     assert dataset_builder._include_deleted_records
 
 
-def test_with_number_of_recent_records_by_record_identifier(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_with_number_of_recent_records_by_record_identifier(
+    sagemaker_session_mock, feature_group_mock
+):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     dataset_builder.with_number_of_recent_records_by_record_identifier(5)
     assert dataset_builder._number_of_recent_records == 5
 
 
-def test_with_number_of_records_from_query_results(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_with_number_of_records_from_query_results(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     dataset_builder.with_number_of_records_from_query_results(100)
     assert dataset_builder._number_of_records == 100
 
 
-def test_as_of(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_as_of(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     time = datetime.datetime.now()
     dataset_builder.as_of(time)
     assert dataset_builder._write_time_ending_timestamp == time
 
 
-def test_with_event_time_range(feature_group_mock):
-    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+def test_with_event_time_range(sagemaker_session_mock, feature_group_mock):
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
     start = datetime.datetime.now()
     end = start + datetime.timedelta(minutes=1)
     dataset_builder.with_event_time_range(start, end)

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -500,13 +500,16 @@ def query(sagemaker_session_mock):
 
 def test_athena_query_run(sagemaker_session_mock, query):
     sagemaker_session_mock.start_query_execution.return_value = {"QueryExecutionId": "query_id"}
-    query.run(query_string="query", output_location="s3://some-bucket/some-path")
+    query.run(
+        query_string="query", output_location="s3://some-bucket/some-path", workgroup="workgroup"
+    )
     sagemaker_session_mock.start_query_execution.assert_called_with(
         catalog="catalog",
         database="database",
         query_string="query",
         output_location="s3://some-bucket/some-path",
         kms_key=None,
+        workgroup="workgroup",
     )
     assert "some-bucket" == query._result_bucket
     assert "some-path" == query._result_file_prefix

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -39,6 +39,7 @@ def test_minimal_create_dataset(sagemaker_session_mock, feature_group_mock):
         base=feature_group_mock,
         output_path="file/to/path",
     )
+    assert dataset_builder._sagemaker_session == sagemaker_session_mock
     assert dataset_builder._base == feature_group_mock
     assert dataset_builder._output_path == "file/to/path"
 
@@ -51,6 +52,7 @@ def test_complete_create_dataset(sagemaker_session_mock, feature_group_mock):
         output_path="file/to/path",
         kms_key_id="kms-key-id",
     )
+    assert dataset_builder._sagemaker_session == sagemaker_session_mock
     assert dataset_builder._base == feature_group_mock
     assert dataset_builder._included_feature_names == ["feature_1", "feature_2"]
     assert dataset_builder._output_path == "file/to/path"
@@ -67,6 +69,7 @@ def test_create_dataset_with_dataframe(sagemaker_session_mock):
         output_path="file/to/path",
         kms_key_id="kms-key-id",
     )
+    assert dataset_builder._sagemaker_session == sagemaker_session_mock
     assert dataset_builder._base.equals(DATAFRAME)
     assert dataset_builder._record_identifier_feature_name == "feature_1"
     assert dataset_builder._event_time_identifier_feature_name == "feature_2"