Add integration tests for create_dataset (aws#743)

Author: imingtsou

src/sagemaker/feature_store/dataset_builder.py

@@ -16,14 +16,14 @@
 """
 from __future__ import absolute_import
 
+import copy
 import datetime
-import os
 from typing import Any, Dict, List, Sequence, Tuple, Union
 
 import attr
 import pandas as pd
 
-from sagemaker import Session, s3, utils
+from sagemaker import Session, utils
 from sagemaker.feature_store.feature_group import FeatureGroup
 
 
@@ -166,7 +166,7 @@ class DatasetBuilder:
     _write_time_ending_timestamp: datetime.datetime = attr.ib(init=False, default=None)
     _event_time_starting_timestamp: datetime.datetime = attr.ib(init=False, default=None)
     _event_time_ending_timestamp: datetime.datetime = attr.ib(init=False, default=None)
-    _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = attr.ib(init=False, default=[])
+    _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = attr.ib(init=False, factory=list)
 
     _DATAFRAME_TYPE_TO_COLUMN_TYPE_MAP = {
         "object": "STRING",
@@ -193,8 +193,6 @@ def with_feature_group(
         Returns:
             This DatasetBuilder object.
         """
-        if not target_feature_name_in_base:
-            target_feature_name_in_base = self._record_identifier_feature_name
         self._feature_groups_to_be_merged.append(
             construct_feature_group_to_be_merged(
                 feature_group, included_feature_names, target_feature_name_in_base
@@ -292,18 +290,11 @@ def to_csv(self) -> Tuple[str, str]:
         """
         if isinstance(self._base, pd.DataFrame):
             temp_id = utils.unique_name_from_base("dataframe-base")
-            local_file_name = f"{temp_id}.csv"
-            desired_s3_folder = f"{self._output_path}/{temp_id}"
-            self._base.to_csv(local_file_name, index=False, header=False)
-            s3.S3Uploader.upload(
-                local_path=local_file_name,
-                desired_s3_uri=desired_s3_folder,
-                sagemaker_session=self._sagemaker_session,
-                kms_key=self._kms_key_id,
-            )
-            os.remove(local_file_name)
-            temp_table_name = f"dataframe_{temp_id}"
-            self._create_temp_table(temp_table_name, desired_s3_folder)
+            s3_file_name = f"{temp_id}.csv"
+            s3_folder = f"{self._output_path}/{temp_id}"
+            self._base.to_csv(f"{s3_folder}/{s3_file_name}", index=False, header=False)
+            temp_table_name = f'dataframe_{temp_id.replace("-", "_")}'
+            self._create_temp_table(temp_table_name, s3_folder)
             base_features = list(self._base.columns)
             query_string = self._construct_query_string(
                 FeatureGroupToBeMerged(
@@ -325,6 +316,10 @@ def to_csv(self) -> Tuple[str, str]:
             base_feature_group = construct_feature_group_to_be_merged(
                 self._base, self._included_feature_names
             )
+            self._record_identifier_feature_name = base_feature_group.record_identifier_feature_name
+            self._event_time_identifier_feature_name = (
+                base_feature_group.event_time_identifier_feature_name
+            )
             query_string = self._construct_query_string(base_feature_group)
             query_result = self._run_query(
                 query_string,
@@ -344,15 +339,7 @@ def to_dataframe(self) -> Tuple[pd.DataFrame, str]:
             The query string executed.
         """
         csv_file, query_string = self.to_csv()
-        s3.S3Downloader.download(
-            s3_uri=csv_file,
-            local_path="./",
-            kms_key=self._kms_key_id,
-            sagemaker_session=self._sagemaker_session,
-        )
-        local_file_name = csv_file.split("/")[-1]
-        df = pd.read_csv(local_file_name)
-        os.remove(local_file_name)
+        df = pd.read_csv(csv_file)
         return df, query_string
 
     def _construct_where_query_string(
@@ -368,21 +355,25 @@ def _construct_where_query_string(
             The WHERE query string.
         """
         where_conditions = []
-        if not self._include_deleted_records:
+        if not self._include_duplicated_records:
+            where_conditions.append(f"row_{suffix} = 1")
+        if not self._include_deleted_records and not isinstance(self._base, pd.DataFrame):
             where_conditions.append("NOT is_deleted")
         if self._write_time_ending_timestamp:
             where_conditions.append(
-                f'table_{suffix}."write_time" <= {self._write_time_ending_timestamp}'
+                f'table_{suffix}."write_time" <= '
+                f"to_timestamp('{self._write_time_ending_timestamp.replace(microsecond=0)}', "
+                f"'yyyy-mm-dd hh24:mi:ss')"
             )
         if self._event_time_starting_timestamp:
             where_conditions.append(
                 f'table_{suffix}."{event_time_identifier_feature_name}" >= '
-                + str(self._event_time_starting_timestamp)
+                + str(self._event_time_starting_timestamp.timestamp())
             )
         if self._event_time_ending_timestamp:
             where_conditions.append(
                 f'table_{suffix}."{event_time_identifier_feature_name}" <= '
-                + str(self._event_time_ending_timestamp)
+                + str(self._event_time_ending_timestamp.timestamp())
             )
         if len(where_conditions) == 0:
             return ""
@@ -410,20 +401,27 @@ def _construct_table_query(self, feature_group: FeatureGroupToBeMerged, suffix:
                 f'FROM "{feature_group.database}"."{feature_group.table_name}" table_{suffix}\n'
             )
         else:
-            features = feature_group.features
+            features = copy.deepcopy(feature_group.features)
             features.remove(feature_group.event_time_identifier_feature_name)
             dedup_features = ", ".join([f'dedup_{suffix}."{feature}"' for feature in features])
+            rank_query_string = (
+                f'ORDER BY dedup_{suffix}."{feature_group.event_time_identifier_feature_name}" '
+                + f'DESC, dedup_{suffix}."api_invocation_time" DESC, dedup_{suffix}."write_time" '
+                + "DESC\n"
+            )
+            if isinstance(self._base, pd.DataFrame):
+                rank_query_string = (
+                    f'ORDER BY dedup_{suffix}."{feature_group.event_time_identifier_feature_name}" '
+                    + "DESC\n"
+                )
             query_string += (
                 "FROM (\n"
                 + "SELECT *, row_number() OVER (\n"
                 + f"PARTITION BY {dedup_features}\n"
-                + f'ORDER BY dedup_{suffix}."{feature_group.event_time_identifier_feature_name}" '
-                + f'DESC, dedup_{suffix}."api_invocation_time" DESC, '
-                + f'dedup_{suffix}."write_time" DESC\n'
+                + rank_query_string
                 + f") AS row_{suffix}\n"
                 + f'FROM "{feature_group.database}"."{feature_group.table_name}" dedup_{suffix}\n'
                 + f") AS table_{suffix}\n"
-                + f"WHERE row_{suffix} = 1\n"
             )
         return query_string + self._construct_where_query_string(
             suffix, feature_group.event_time_identifier_feature_name
@@ -446,8 +444,8 @@ def _construct_query_string(self, base: FeatureGroupToBeMerged) -> str:
                     for i, feature_group in enumerate(self._feature_groups_to_be_merged)
                 ]
             )
-            query_string += f"{with_subquery_string}\n"
-        query_string += "SELECT *\nFROM fg_base"
+            query_string += with_subquery_string
+        query_string += "\nSELECT *\nFROM fg_base"
         if len(self._feature_groups_to_be_merged) > 0:
             join_subquery_string = "".join(
                 [
@@ -470,6 +468,8 @@ def _construct_join_condition(self, feature_group: FeatureGroupToBeMerged, suffi
         Returns:
             The JOIN query string.
         """
+        if not feature_group.target_feature_name_in_base:
+            feature_group.target_feature_name_in_base = self._record_identifier_feature_name
         join_condition_string = (
             f"\nJOIN fg_{suffix}\n"
             + f'ON fg_base."{feature_group.target_feature_name_in_base}" = '

tests/integ/test_feature_store.py

@@ -82,6 +82,11 @@ def feature_group_name():
     return f"my-feature-group-{int(time.time() * 10**7)}"
 
 
+@pytest.fixture
+def base_name():
+    return f"my-base-{int(time.time() * 10**7)}"
+
+
 @pytest.fixture
 def offline_store_s3_uri(feature_store_session, region_name):
     bucket = f"sagemaker-test-featurestore-{region_name}-{feature_store_session.account_id()}"
@@ -109,6 +114,32 @@ def pandas_data_frame():
     return df
 
 
+@pytest.fixture
+def base_dataframe():
+    base_data = [
+        [1, 187512346.0, 123, 128],
+        [2, 187512347.0, 168, 258],
+        [3, 187512348.0, 125, 184],
+        [1, 187512349.0, 195, 206],
+    ]
+    return pd.DataFrame(
+        base_data, columns=["base_id", "base_time", "base_feature_1", "base_feature_2"]
+    )
+
+
+@pytest.fixture
+def feature_group_dataframe():
+    feature_group_data = [
+        [1, 187512246.0, 456, 325],
+        [2, 187512247.0, 729, 693],
+        [3, 187512348.0, 129, 901],
+        [1, 187512449.0, 289, 286],
+    ]
+    return pd.DataFrame(
+        feature_group_data, columns=["fg_id", "fg_time", "fg_feature_1", "fg_feature_2"]
+    )
+
+
 @pytest.fixture
 def pandas_data_frame_without_string():
     df = pd.DataFrame(
@@ -527,6 +558,135 @@ def test_ingest_multi_process(
     assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
 
 
+def test_create_dataset_with_feature_group_base(
+    feature_store_session,
+    region_name,
+    role,
+    base_name,
+    feature_group_name,
+    offline_store_s3_uri,
+    base_dataframe,
+    feature_group_dataframe,
+):
+    base = FeatureGroup(name=base_name, sagemaker_session=feature_store_session)
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    with cleanup_feature_group(base) and cleanup_feature_group(feature_group):
+        _create_feature_group_and_ingest_data(
+            base, base_dataframe, offline_store_s3_uri, "base_id", "base_time", role
+        )
+        _create_feature_group_and_ingest_data(
+            feature_group, feature_group_dataframe, offline_store_s3_uri, "fg_id", "fg_time", role
+        )
+        base_table_name = _get_athena_table_name_after_data_replication(
+            feature_store_session, base, offline_store_s3_uri
+        )
+        feature_group_table_name = _get_athena_table_name_after_data_replication(
+            feature_store_session, feature_group, offline_store_s3_uri
+        )
+
+        with timeout(minutes=10) and cleanup_offline_store(
+            base_table_name, feature_store_session
+        ) and cleanup_offline_store(feature_group_table_name, feature_store_session):
+            feature_store = FeatureStore(sagemaker_session=feature_store_session)
+            df, query_string = (
+                feature_store.create_dataset(base=base, output_path=offline_store_s3_uri)
+                .with_feature_group(feature_group)
+                .to_dataframe()
+            )
+            sorted_df = df.sort_values(by=list(df.columns)).reset_index(drop=True)
+            merged_df = base_dataframe.merge(
+                feature_group_dataframe, left_on="base_id", right_on="fg_id"
+            )
+            expect_df = merged_df.sort_values(by=list(merged_df.columns)).reset_index(drop=True)
+            assert sorted_df.equals(expect_df)
+            assert (
+                query_string
+                == 'WITH fg_base AS (SELECT table_base."base_id", table_base."base_time", '
+                + 'table_base."base_feature_1", table_base."base_feature_2"\n'
+                + "FROM (\n"
+                + "SELECT *, row_number() OVER (\n"
+                + 'PARTITION BY dedup_base."base_id", dedup_base."base_feature_1", '
+                + 'dedup_base."base_feature_2"\n'
+                + 'ORDER BY dedup_base."base_time" DESC, dedup_base."api_invocation_time" DESC, '
+                + 'dedup_base."write_time" DESC\n'
+                + ") AS row_base\n"
+                + f'FROM "sagemaker_featurestore"."{base_table_name}" dedup_base\n'
+                + ") AS table_base\n"
+                + "WHERE row_base = 1\n"
+                + "AND NOT is_deleted),\n"
+                + 'fg_0 AS (SELECT table_0."fg_id", table_0."fg_time", table_0."fg_feature_1", '
+                + 'table_0."fg_feature_2"\n'
+                + "FROM (\n"
+                + "SELECT *, row_number() OVER (\n"
+                + 'PARTITION BY dedup_0."fg_id", dedup_0."fg_feature_1", dedup_0."fg_feature_2"\n'
+                + 'ORDER BY dedup_0."fg_time" DESC, dedup_0."api_invocation_time" DESC, '
+                + 'dedup_0."write_time" DESC\n'
+                + ") AS row_0\n"
+                + f'FROM "sagemaker_featurestore"."{feature_group_table_name}" dedup_0\n'
+                + ") AS table_0\n"
+                + "WHERE row_0 = 1\n"
+                + "AND NOT is_deleted)\n"
+                + "SELECT *\n"
+                + "FROM fg_base\n"
+                + "JOIN fg_0\n"
+                + 'ON fg_base."base_id" = fg_0."fg_id"'
+            )
+
+
+def _create_feature_group_and_ingest_data(
+    feature_group: FeatureGroup,
+    dataframe: DataFrame,
+    offline_store_s3_uri: str,
+    record_identifier_name: str,
+    event_time_name: str,
+    role: str,
+):
+    feature_group.load_feature_definitions(data_frame=dataframe)
+    feature_group.create(
+        s3_uri=offline_store_s3_uri,
+        record_identifier_name=record_identifier_name,
+        event_time_feature_name=event_time_name,
+        role_arn=role,
+        enable_online_store=True,
+    )
+    _wait_for_feature_group_create(feature_group)
+
+    ingestion_manager = feature_group.ingest(data_frame=dataframe, max_workers=3, wait=False)
+    ingestion_manager.wait()
+    assert 0 == len(ingestion_manager.failed_rows)
+
+
+def _get_athena_table_name_after_data_replication(
+    feature_store_session, feature_group: FeatureGroup, offline_store_s3_uri
+):
+    feature_group_metadata = feature_group.describe()
+    resolved_output_s3_uri = (
+        feature_group_metadata.get("OfflineStoreConfig", None)
+        .get("S3StorageConfig", None)
+        .get("ResolvedOutputS3Uri", None)
+    )
+    s3_prefix = resolved_output_s3_uri.replace(f"{offline_store_s3_uri}/", "")
+    region_name = feature_store_session.boto_session.region_name
+    s3_client = feature_store_session.boto_session.client(
+        service_name="s3", region_name=region_name
+    )
+    while True:
+        objects_in_bucket = s3_client.list_objects(
+            Bucket=offline_store_s3_uri.replace("s3://", ""), Prefix=s3_prefix
+        )
+        if "Contents" in objects_in_bucket and len(objects_in_bucket["Contents"]) > 1:
+            break
+        else:
+            print(f"Waiting for {feature_group.name} data in offline store...")
+            time.sleep(60)
+    print(f"{feature_group.name} data available.")
+    return (
+        feature_group_metadata.get("OfflineStoreConfig", None)
+        .get("DataCatalogConfig", None)
+        .get("TableName", None)
+    )
+
+
 def _wait_for_feature_group_create(feature_group: FeatureGroup):
     status = feature_group.describe().get("FeatureGroupStatus")
     while status == "Creating":
@@ -560,3 +720,28 @@ def cleanup_feature_group(feature_group: FeatureGroup):
             feature_group.delete()
         except Exception:
             raise RuntimeError(f"Failed to delete feature group with name {feature_group.name}")
+
+
+@contextmanager
+def cleanup_offline_store(table_name: str, feature_store_session: Session):
+    try:
+        yield
+    finally:
+        try:
+            region_name = feature_store_session.boto_session.region_name
+            s3_client = feature_store_session.boto_session.client(
+                service_name="s3", region_name=region_name
+            )
+            account_id = feature_store_session.account_id()
+            bucket_name = f"sagemaker-test-featurestore-{region_name}-{account_id}"
+            response = s3_client.list_objects_v2(
+                Bucket=bucket_name,
+                Prefix=f"{account_id}/sagemaker/{region_name}/offline-store/{table_name}/",
+            )
+            files_in_folder = response["Contents"]
+            files_to_delete = []
+            for f in files_in_folder:
+                files_to_delete.append({"Key": f["Key"]})
+            s3_client.delete_objects(Bucket=bucket_name, Delete={"Objects": files_to_delete})
+        except Exception:
+            raise RuntimeError(f"Failed to delete data under {table_name}")