feature: Add support of collection types in feature store. (#1106)

Author: AWS-pratab

doc/api/prep_data/feature_store.rst

@@ -41,6 +41,25 @@ Feature Definition
     :members:
     :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.feature_definition.CollectionTypeEnum
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.feature_definition.CollectionType
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.feature_definition.ListCollectionType
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.feature_definition.SetCollectionType
+    :members:
+    :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.feature_definition.VectorCollectionType
+    :members:
+    :show-inheritance:
 
 Inputs
 ******

src/sagemaker/feature_store/feature_definition.py

@@ -39,6 +39,87 @@ class FeatureTypeEnum(Enum):
     STRING = "String"
 
 
+class CollectionTypeEnum(Enum):
+    """Enum of collection types.
+
+    The collection type of a feature can be List, Set or Vector.
+    """
+
+    LIST = "List"
+    SET = "Set"
+    VECTOR = "Vector"
+
+
+@attr.s
+class CollectionType(Config):
+    """Collection type and its configuration.
+
+    This initiates a collectiontype object where CollectionType is a subclass of Config.
+
+    Attributes:
+        collection_type (CollectionTypeEnum): The type of the collection
+        collection_config (Dict[str, Any]): The configuration for the collection.
+    """
+
+    collection_type: CollectionTypeEnum = attr.ib()
+    collection_config: Dict[str, Any] = attr.ib()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Construct a dictionary based on each attribute."""
+        return Config.construct_dict(
+            CollectionType=self.collection_type.value, CollectionConfig=self.collection_config
+        )
+
+
+class ListCollectionType(CollectionType):
+    """List collection type
+
+    This class instantiates a ListCollectionType object, as subclass of CollectionType
+    where the collection type is defined as List.
+
+    """
+
+    def __init__(self):
+        """Construct an instance of ListCollectionType."""
+        super(ListCollectionType, self).__init__(CollectionTypeEnum.LIST, None)
+
+
+class SetCollectionType(CollectionType):
+    """Set collection type
+
+    This class instantiates a SetCollectionType object, as subclass of CollectionType
+    where the collection type is defined as Set.
+
+    """
+
+    def __init__(self):
+        """Construct an instance of SetCollectionType."""
+        super(SetCollectionType, self).__init__(CollectionTypeEnum.SET, None)
+
+
+class VectorCollectionType(CollectionType):
+    """Vector collection type
+
+    This class instantiates a VectorCollectionType object, as subclass of CollectionType
+    where the collection type is defined as Vector.
+
+    Attributes:
+        dimension (int): The dimension size for the Vector.
+    """
+
+    def __init__(self, dimension: int):
+        """Construct an instance of VectorCollectionType.
+
+        Attributes:
+        dimension (int): The dimension size for the Vector.
+        """
+        collection_config: Dict[str, Any] = {}
+        vector_config: Dict[str, Any] = {}
+        vector_config["Dimension"] = dimension
+        collection_config["VectorConfig"] = vector_config
+        super(VectorCollectionType, self).__init__(CollectionTypeEnum.VECTOR, collection_config)
+
+
 @attr.s
 class FeatureDefinition(Config):
     """Feature definition.
@@ -48,15 +129,25 @@ class FeatureDefinition(Config):
     Attributes:
         feature_name (str): The name of the feature
         feature_type (FeatureTypeEnum): The type of the feature
+        collection_type (CollectionType): The type of collection for the feature
     """
 
     feature_name: str = attr.ib()
     feature_type: FeatureTypeEnum = attr.ib()
+    collection_type: CollectionType = attr.ib(default=None)
 
     def to_dict(self) -> Dict[str, Any]:
         """Construct a dictionary based on each attribute."""
+
         return Config.construct_dict(
-            FeatureName=self.feature_name, FeatureType=self.feature_type.value
+            FeatureName=self.feature_name,
+            FeatureType=self.feature_type.value,
+            CollectionType=(
+                self.collection_type.collection_type.value if self.collection_type else None
+            ),
+            CollectionConfig=(
+                self.collection_type.collection_config if self.collection_type else None
+            ),
         )
 
 
@@ -69,15 +160,18 @@ class FractionalFeatureDefinition(FeatureDefinition):
     Attributes:
         feature_name (str): The name of the feature
         feature_type (FeatureTypeEnum): A `FeatureTypeEnum.FRACTIONAL` type
+        collection_type (CollectionType): The type of collection for the feature
     """
 
-    def __init__(self, feature_name: str):
+    def __init__(self, feature_name: str, collection_type: CollectionType = None):
         """Construct an instance of FractionalFeatureDefinition.
 
         Args:
             feature_name (str): the name of the feature.
         """
-        super(FractionalFeatureDefinition, self).__init__(feature_name, FeatureTypeEnum.FRACTIONAL)
+        super(FractionalFeatureDefinition, self).__init__(
+            feature_name, FeatureTypeEnum.FRACTIONAL, collection_type
+        )
 
 
 class IntegralFeatureDefinition(FeatureDefinition):
@@ -89,15 +183,18 @@ class IntegralFeatureDefinition(FeatureDefinition):
     Attributes:
         feature_name (str): the name of the feature.
         feature_type (FeatureTypeEnum): a `FeatureTypeEnum.INTEGRAL` type.
+        collection_type (CollectionType): The type of collection for the feature.
     """
 
-    def __init__(self, feature_name: str):
+    def __init__(self, feature_name: str, collection_type: CollectionType = None):
         """Construct an instance of IntegralFeatureDefinition.
 
         Args:
             feature_name (str): the name of the feature.
         """
-        super(IntegralFeatureDefinition, self).__init__(feature_name, FeatureTypeEnum.INTEGRAL)
+        super(IntegralFeatureDefinition, self).__init__(
+            feature_name, FeatureTypeEnum.INTEGRAL, collection_type
+        )
 
 
 class StringFeatureDefinition(FeatureDefinition):
@@ -109,12 +206,15 @@ class StringFeatureDefinition(FeatureDefinition):
     Attributes:
         feature_name (str): the name of the feature.
         feature_type (FeatureTypeEnum): a `FeatureTypeEnum.STRING` type.
+        collection_type (CollectionType): The type of collection for the feature.
     """
 
-    def __init__(self, feature_name: str):
+    def __init__(self, feature_name: str, collection_type: CollectionType = None):
         """Construct an instance of StringFeatureDefinition.
 
         Args:
             feature_name (str): the name of the feature.
         """
-        super(StringFeatureDefinition, self).__init__(feature_name, FeatureTypeEnum.STRING)
+        super(StringFeatureDefinition, self).__init__(
+            feature_name, FeatureTypeEnum.STRING, collection_type
+        )

src/sagemaker/feature_store/inputs.py

@@ -266,10 +266,13 @@ class FeatureValue(Config):
     Attributes:
         feature_name (str): name of the Feature.
         value_as_string (str): value of the Feature in string form.
+        value_as_string_list (List[str]): value of the Feature in string list
+        form used for collection type.
     """
 
     feature_name: str = attr.ib(default=None)
     value_as_string: str = attr.ib(default=None)
+    value_as_string_list: List[str] = attr.ib(default=None)
 
     def to_dict(self) -> Dict[str, Any]:
         """Construct a dictionary based on the attributes provided.
@@ -280,6 +283,7 @@ def to_dict(self) -> Dict[str, Any]:
         return Config.construct_dict(
             FeatureName=self.feature_name,
             ValueAsString=self.value_as_string,
+            ValueAsStringList=self.value_as_string_list,
         )
 
 

tests/integ/test_feature_store.py

@@ -25,7 +25,11 @@
 from pandas import DataFrame
 
 from sagemaker.feature_store.feature_utils import get_feature_group_as_dataframe
-from sagemaker.feature_store.feature_definition import FractionalFeatureDefinition
+from sagemaker.feature_store.feature_definition import (
+    FractionalFeatureDefinition,
+    StringFeatureDefinition,
+    ListCollectionType,
+)
 from sagemaker.feature_store.feature_group import FeatureGroup
 from sagemaker.feature_store.feature_store import FeatureStore
 from sagemaker.feature_store.inputs import (
@@ -186,6 +190,16 @@ def record():
     ]
 
 
+@pytest.fixture
+def collection_type_record():
+    return [
+        FeatureValue(feature_name="feature1", value_as_string="10.0"),
+        FeatureValue(feature_name="feature2", value_as_string="10"),
+        FeatureValue(feature_name="feature3", value_as_string="2020-10-30T03:43:21Z"),
+        FeatureValue(feature_name="feature4", value_as_string_list=["val1", "val2"]),
+    ]
+
+
 @pytest.fixture
 def create_table_ddl():
     return (
@@ -664,6 +678,60 @@ def test_get_and_batch_get_record(
             assert feature["FeatureName"] is not removed_feature_name
 
 
+def test_put_and_get_collection_type_record(
+    feature_store_session,
+    role,
+    feature_group_name,
+    pandas_data_frame,
+    collection_type_record,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_definition_with_collection = [
+        StringFeatureDefinition(feature_name="feature1"),
+        StringFeatureDefinition(feature_name="feature2"),
+        StringFeatureDefinition(feature_name="feature3"),
+        StringFeatureDefinition(feature_name="feature4", collection_type=ListCollectionType()),
+    ]
+    feature_group.feature_definitions = feature_definition_with_collection
+    record_identifier_value_as_string = collection_type_record[0].value_as_string
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=False,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+            online_store_storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
+        )
+        _wait_for_feature_group_create(feature_group)
+        # Ingest data
+        feature_group.put_record(record=collection_type_record)
+        # Retrieve data
+        retrieved_record = feature_group.get_record(
+            record_identifier_value_as_string=record_identifier_value_as_string,
+        )
+
+        assert retrieved_record is not None
+        record_names = list(map(lambda r: r.feature_name, collection_type_record))
+        assert len(retrieved_record) == len(record_names)
+
+        retrieved_feature_map = {}
+        for feature in retrieved_record:
+            assert feature["FeatureName"] in record_names
+            retrieved_feature_map[feature["FeatureName"]] = (
+                feature.get("ValueAsStringList")
+                if feature.get("ValueAsString") is None
+                else feature.get("ValueAsString")
+            )
+
+        assert collection_type_record[0].value_as_string == retrieved_feature_map.get("feature1")
+        assert collection_type_record[1].value_as_string == retrieved_feature_map.get("feature2")
+        assert collection_type_record[2].value_as_string == retrieved_feature_map.get("feature3")
+        assert collection_type_record[3].value_as_string_list == retrieved_feature_map.get(
+            "feature4"
+        )
+
+
 def test_soft_delete_record(
     feature_store_session,
     role,
@@ -949,7 +1017,12 @@ def test_create_dataset_with_feature_group_base(
             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
+            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
@@ -1131,7 +1204,12 @@ def test_create_dataset_with_feature_group_base_with_additional_params(
             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
+            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
@@ -1157,7 +1235,10 @@ def test_create_dataset_with_feature_group_base_with_additional_params(
             )
             sorted_df = df.sort_values(by=list(df.columns)).reset_index(drop=True)
             merged_df = base_dataframe.merge(
-                feature_group_dataframe, left_on="base_time", right_on="fg_time", how="outer"
+                feature_group_dataframe,
+                left_on="base_time",
+                right_on="fg_time",
+                how="outer",
             )
 
             expect_df = merged_df.sort_values(by=list(merged_df.columns)).reset_index(drop=True)