feature: support table format option for create feature group.

Co-authored-by: Mufaddal Rohawala <[email protected]>

Author: 2 people

src/sagemaker/feature_store/feature_group.py

@@ -54,6 +54,7 @@
     DataCatalogConfig,
     FeatureValue,
     FeatureParameter,
+    TableFormatEnum,
 )
 
 logger = logging.getLogger(__name__)
@@ -461,6 +462,7 @@ def create(
         data_catalog_config: DataCatalogConfig = None,
         description: str = None,
         tags: List[Dict[str, str]] = None,
+        table_format: TableFormatEnum = None,
     ) -> Dict[str, Any]:
         """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -470,20 +472,23 @@ def create(
             record_identifier_name (str): name of the record identifier feature.
             event_time_feature_name (str): name of the event time feature.
             role_arn (str): ARN of the role used to call CreateFeatureGroup.
-            online_store_kms_key_id (str): KMS key id for online store.
-            enable_online_store (bool): whether to enable online store or not.
-            offline_store_kms_key_id (str): KMS key id for offline store.
+            online_store_kms_key_id (str): KMS key id for online store (default: None).
+            enable_online_store (bool): whether to enable online store or not (default: False).
+            offline_store_kms_key_id (str): KMS key id for offline store (default: None).
                 If a KMS encryption key is not specified, SageMaker encrypts all data at
                 rest using the default AWS KMS key. By defining your bucket-level key for
                 SSE, you can reduce the cost of AWS KMS requests.
                 For more information, see
                 `Bucket Key
                 <https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucket-key.html>`_
                 in the Amazon S3 User Guide.
-            disable_glue_table_creation (bool): whether to turn off Glue table creation no not.
-            data_catalog_config (DataCatalogConfig): configuration for Metadata store.
-            description (str): description of the FeatureGroup.
-            tags (List[Dict[str, str]]): list of tags for labeling a FeatureGroup.
+            disable_glue_table_creation (bool): whether to turn off Glue table creation
+                or not (default: False).
+            data_catalog_config (DataCatalogConfig): configuration for
+                Metadata store (default: None).
+            description (str): description of the FeatureGroup (default: None).
+            tags (List[Dict[str, str]]): list of tags for labeling a FeatureGroup (default: None).
+            table_format (TableFormatEnum): format of the offline store table (default: None).
 
         Returns:
             Response dict from service.
@@ -518,6 +523,7 @@ def create(
                 s3_storage_config=s3_storage_config,
                 disable_glue_table_creation=disable_glue_table_creation,
                 data_catalog_config=data_catalog_config,
+                table_format=table_format,
             )
             create_feature_store_args.update(
                 {"offline_store_config": offline_store_config.to_dict()}

src/sagemaker/feature_store/inputs.py

@@ -30,6 +30,7 @@
 
 import abc
 from typing import Dict, Any
+from enum import Enum
 
 import attr
 
@@ -158,6 +159,16 @@ def to_dict(self) -> Dict[str, Any]:
         )
 
 
+class TableFormatEnum(Enum):
+    """Enum of table formats.
+
+    The offline store table formats can be Glue or Iceberg.
+    """
+
+    GLUE = "Glue"
+    ICEBERG = "Iceberg"
+
+
 @attr.s
 class OfflineStoreConfig(Config):
     """OfflineStoreConfig for FeatureStore.
@@ -166,11 +177,13 @@ class OfflineStoreConfig(Config):
         s3_storage_config (S3StorageConfig): configuration of S3 storage.
         disable_glue_table_creation (bool): whether to disable the Glue table creation.
         data_catalog_config (DataCatalogConfig): configuration of the data catalog.
+        table_format (TableFormatEnum): format of the offline store table.
     """
 
     s3_storage_config: S3StorageConfig = attr.ib()
     disable_glue_table_creation: bool = attr.ib(default=False)
     data_catalog_config: DataCatalogConfig = attr.ib(default=None)
+    table_format: TableFormatEnum = attr.ib(default=None)
 
     def to_dict(self) -> Dict[str, Any]:
         """Construct a dictionary based on the attributes.
@@ -182,6 +195,7 @@ def to_dict(self) -> Dict[str, Any]:
             DisableGlueTableCreation=self.disable_glue_table_creation,
             S3StorageConfig=self.s3_storage_config,
             DataCatalogConfig=self.data_catalog_config,
+            TableFormat=self.table_format.value if self.table_format else None,
         )
 
 

tests/integ/test_feature_store.py

@@ -24,7 +24,7 @@
 
 from sagemaker.feature_store.feature_definition import FractionalFeatureDefinition
 from sagemaker.feature_store.feature_group import FeatureGroup
-from sagemaker.feature_store.inputs import FeatureValue, FeatureParameter
+from sagemaker.feature_store.inputs import FeatureValue, FeatureParameter, TableFormatEnum
 from sagemaker.session import get_execution_role, Session
 from tests.integ.timeout import timeout
 
@@ -238,6 +238,56 @@ def test_create_feature_store(
     assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
 
 
+def test_create_feature_group_iceberg_table_format(
+    feature_store_session,
+    role,
+    feature_group_name,
+    offline_store_s3_uri,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=offline_store_s3_uri,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+            table_format=TableFormatEnum.ICEBERG,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        table_format = feature_group.describe().get("OfflineStoreConfig").get("TableFormat")
+        assert table_format == "Iceberg"
+
+
+def test_create_feature_group_glue_table_format(
+    feature_store_session,
+    role,
+    feature_group_name,
+    offline_store_s3_uri,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=offline_store_s3_uri,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+            table_format=TableFormatEnum.GLUE,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        table_format = feature_group.describe().get("OfflineStoreConfig").get("TableFormat")
+        assert table_format == "Glue"
+
+
 def test_update_feature_group(
     feature_store_session,
     role,

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -31,7 +31,10 @@
     AthenaQuery,
     IngestionError,
 )
-from sagemaker.feature_store.inputs import FeatureParameter
+from sagemaker.feature_store.inputs import (
+    FeatureParameter,
+    TableFormatEnum,
+)
 
 
 class PicklableMock(Mock):
@@ -115,6 +118,68 @@ def test_feature_store_create(
     )
 
 
+def test_feature_store_create_iceberg_table_format(
+    sagemaker_session_mock, role_arn, feature_group_dummy_definitions, s3_uri
+):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    feature_group.feature_definitions = feature_group_dummy_definitions
+    feature_group.create(
+        s3_uri=s3_uri,
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        role_arn=role_arn,
+        enable_online_store=True,
+        disable_glue_table_creation=False,
+        table_format=TableFormatEnum.ICEBERG,
+    )
+    sagemaker_session_mock.create_feature_group.assert_called_with(
+        feature_group_name="MyFeatureGroup",
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        feature_definitions=[fd.to_dict() for fd in feature_group_dummy_definitions],
+        role_arn=role_arn,
+        description=None,
+        tags=None,
+        online_store_config={"EnableOnlineStore": True},
+        offline_store_config={
+            "DisableGlueTableCreation": False,
+            "TableFormat": "Iceberg",
+            "S3StorageConfig": {"S3Uri": s3_uri},
+        },
+    )
+
+
+def test_feature_store_create_glue_table_format(
+    sagemaker_session_mock, role_arn, feature_group_dummy_definitions, s3_uri
+):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    feature_group.feature_definitions = feature_group_dummy_definitions
+    feature_group.create(
+        s3_uri=s3_uri,
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        role_arn=role_arn,
+        enable_online_store=True,
+        disable_glue_table_creation=False,
+        table_format=TableFormatEnum.GLUE,
+    )
+    sagemaker_session_mock.create_feature_group.assert_called_with(
+        feature_group_name="MyFeatureGroup",
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        feature_definitions=[fd.to_dict() for fd in feature_group_dummy_definitions],
+        role_arn=role_arn,
+        description=None,
+        tags=None,
+        online_store_config={"EnableOnlineStore": True},
+        offline_store_config={
+            "DisableGlueTableCreation": False,
+            "TableFormat": "Glue",
+            "S3StorageConfig": {"S3Uri": s3_uri},
+        },
+    )
+
+
 def test_feature_store_create_online_only(
     sagemaker_session_mock, role_arn, feature_group_dummy_definitions
 ):

tests/unit/sagemaker/feature_store/test_inputs.py

@@ -20,6 +20,7 @@
     DataCatalogConfig,
     OfflineStoreConfig,
     FeatureParameter,
+    TableFormatEnum,
 )
 
 
@@ -86,6 +87,32 @@ def test_offline_data_store_config():
     )
 
 
+def test_offline_data_store_config_with_glue_table_format():
+    config = OfflineStoreConfig(
+        s3_storage_config=S3StorageConfig(s3_uri="uri"), table_format=TableFormatEnum.GLUE
+    )
+    assert ordered(config.to_dict()) == ordered(
+        {
+            "S3StorageConfig": {"S3Uri": "uri"},
+            "DisableGlueTableCreation": False,
+            "TableFormat": "Glue",
+        }
+    )
+
+
+def test_offline_data_store_config_with_iceberg_table_format():
+    config = OfflineStoreConfig(
+        s3_storage_config=S3StorageConfig(s3_uri="uri"), table_format=TableFormatEnum.ICEBERG
+    )
+    assert ordered(config.to_dict()) == ordered(
+        {
+            "S3StorageConfig": {"S3Uri": "uri"},
+            "DisableGlueTableCreation": False,
+            "TableFormat": "Iceberg",
+        }
+    )
+
+
 def test_feature_metadata():
     config = FeatureParameter(key="key", value="value")
     assert ordered(config.to_dict()) == ordered({"Key": "key", "Value": "value"})