change: use ResolvedOutputS3Uir for Hive DDL LOCATION (#2213)

* Update feature_group.py

* Update test_feature_store.py

* Update feature_group.py

Co-authored-by: Dan <[email protected]>
Co-authored-by: icywang86rui <[email protected]>

Author: 3 people

src/sagemaker/feature_store/feature_group.py

@@ -512,14 +512,11 @@ def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str
         if not table_name:
             table_name = self.name
 
-        s3_uri = self.describe().get("OfflineStoreConfig").get("S3StorageConfig").get("S3Uri")
-        offline_store_s3_uri = os.path.join(
-            s3_uri,
-            self.sagemaker_session.account_id(),
-            "sagemaker",
-            self.sagemaker_session.boto_session.region_name,
-            "offline-store",
-            self.name,
+        resolved_output_s3_uri = (
+            self.describe()
+            .get("OfflineStoreConfig")
+            .get("S3StorageConfig")
+            .get("ResolvedOutputS3Uri")
         )
 
         ddl = f"CREATE EXTERNAL TABLE IF NOT EXISTS {database}.{table_name} (\n"
@@ -537,6 +534,6 @@ def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str
             "  STORED AS\n"
             "  INPUTFORMAT 'parquet.hive.DeprecatedParquetInputFormat'\n"
             "  OUTPUTFORMAT 'parquet.hive.DeprecatedParquetOutputFormat'\n"
-            f"LOCATION '{offline_store_s3_uri}'"
+            f"LOCATION '{resolved_output_s3_uri}'"
         )
         return ddl

tests/integ/test_feature_store.py

@@ -142,8 +142,7 @@ def create_table_ddl():
         "  STORED AS\n"
         "  INPUTFORMAT 'parquet.hive.DeprecatedParquetInputFormat'\n"
         "  OUTPUTFORMAT 'parquet.hive.DeprecatedParquetOutputFormat'\n"
-        "LOCATION 's3://sagemaker-test-featurestore-{region}-{account}"
-        "/{account}/sagemaker/us-east-2/offline-store/{feature_group_name}'"
+        "LOCATION '{resolved_output_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,6 +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,
             )
             == 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"