Update feature_group.py

Author: can-sun

src/sagemaker/feature_store/feature_group.py

@@ -511,8 +511,13 @@ def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str
         """
         if not table_name:
             table_name = self.name
-            
-        resolved_output_s3_uri = self.describe().get("OfflineStoreConfig").get("S3StorageConfig").get("ResolvedOutputS3Uri")
+
+        resolved_output_s3_uri = (
+            self.describe()
+            .get("OfflineStoreConfig")
+            .get("S3StorageConfig")
+            .get("ResolvedOutputS3Uri")
+        )
 
         ddl = f"CREATE EXTERNAL TABLE IF NOT EXISTS {database}.{table_name} (\n"
         for definition in self.feature_definitions:

tests/integ/test_feature_store.py

@@ -179,8 +179,7 @@ def test_create_feature_store(
 ):
     feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
     feature_group.load_feature_definitions(data_frame=pandas_data_frame)
-    resolved_output_s3_uri = feature_group.describe().get("OfflineStoreConfig").get("S3StorageConfig").get("ResolvedOutputS3Uri")
-    
+
     with cleanup_feature_group(feature_group):
         output = feature_group.create(
             s3_uri=offline_store_s3_uri,
@@ -191,6 +190,12 @@ def test_create_feature_store(
         )
         _wait_for_feature_group_create(feature_group)
 
+        resolved_output_s3_uri = (
+            feature_group.describe()
+            .get("OfflineStoreConfig")
+            .get("S3StorageConfig")
+            .get("ResolvedOutputS3Uri")
+        )
         # Ingest data
         feature_group.put_record(record=record)
         ingestion_manager = feature_group.ingest(
@@ -225,7 +230,7 @@ def test_create_feature_store(
                 feature_group_name=feature_group_name,
                 region=feature_store_session.boto_session.region_name,
                 account=feature_store_session.account_id(),
-                resolved_output_s3_uri=resolved_output_s3_uri
+                resolved_output_s3_uri=resolved_output_s3_uri,
             )
             == feature_group.as_hive_ddl()
         )

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -66,8 +66,7 @@ def create_table_ddl():
         "  STORED AS\n"
         "  INPUTFORMAT 'parquet.hive.DeprecatedParquetInputFormat'\n"
         "  OUTPUTFORMAT 'parquet.hive.DeprecatedParquetOutputFormat'\n"
-        "LOCATION 's3://some-bucket"
-        "/{account}/sagemaker/{region}/offline-store/{feature_group_name}'"
+        "LOCATION 's3://resolved_output_s3_uri'"
     )
 
 
@@ -217,7 +216,12 @@ def test_as_hive_ddl_with_default_values(
     create_table_ddl, feature_group_dummy_definitions, sagemaker_session_mock
 ):
     sagemaker_session_mock.describe_feature_group.return_value = {
-        "OfflineStoreConfig": {"S3StorageConfig": {"S3Uri": "s3://some-bucket"}}
+        "OfflineStoreConfig": {
+            "S3StorageConfig": {
+                "S3Uri": "s3://some-bucket",
+                "ResolvedOutputS3Uri": "s3://resolved_output_s3_uri",
+            }
+        }
     }
     sagemaker_session_mock.account_id.return_value = "1234"
     sagemaker_session_mock.boto_session.region_name = "us-west-2"
@@ -238,7 +242,12 @@ def test_as_hive_ddl_with_default_values(
 
 def test_as_hive_ddl(create_table_ddl, feature_group_dummy_definitions, sagemaker_session_mock):
     sagemaker_session_mock.describe_feature_group.return_value = {
-        "OfflineStoreConfig": {"S3StorageConfig": {"S3Uri": "s3://some-bucket"}}
+        "OfflineStoreConfig": {
+            "S3StorageConfig": {
+                "S3Uri": "s3://some-bucket",
+                "ResolvedOutputS3Uri": "s3://resolved_output_s3_uri",
+            }
+        }
     }
     sagemaker_session_mock.account_id.return_value = "1234"
     sagemaker_session_mock.boto_session.region_name = "us-west-2"