feat: Add DatasetBuilder class (aws#667)

imingtsou · Eric Zou · mizanfiu · commit 699f9b0f5da6 · 2022-12-13T01:02:12.000-08:00
Co-authored-by: Eric Zou &lt;zoueric@amazon.com&gt;
diff --git a/src/sagemaker/feature_store/dataset_builder.py b/src/sagemaker/feature_store/dataset_builder.py
@@ -0,0 +1,157 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Dataset Builder
+
+A Dataset Builder is a builder class for generating a dataset by providing conditions.
+"""
+from __future__ import absolute_import
+
+import datetime
+from typing import Sequence, Union
+
+import attr
+import pandas as pd
+
+from sagemaker.feature_store.feature_group import FeatureGroup
+
+
+@attr.s
+class DatasetBuilder:
+    """DatasetBuilder definition.
+
+    This class instantiates a DatasetBuilder object that comprises a base, a list of feature names,
+    an output path and a KMS key ID.
+
+    Attributes:
+        _base (Union[FeatureGroup, DataFrame]): A base which can be either a FeatureGroup or a
+            pandas.DataFrame and will be used to merge other FeatureGroups and generate a Dataset.
+        _output_path (str): An S3 URI which stores the output .csv file.
+        _record_identifier_feature_name (str): A string representing the record identifier feature
+            if base is a DataFrame.
+        _event_time_identifier_feature_name (str): A string representing the event time identifier
+            feature if base is a DataFrame.
+        _included_feature_names (List[str]): A list of features to be included in the output.
+        _kms_key_id (str): An KMS key id. If set, will be used to encrypt the result file.
+        _point_in_time_accurate_join (bool): A boolean representing whether using point in time join
+            or not.
+        _include_duplicated_records (bool): A boolean representing whether including duplicated
+            records or not.
+        _include_deleted_records (bool): A boolean representing whether including deleted records or
+            not.
+        _number_of_recent_records (int): An int that how many records will be returned for each
+            record identifier.
+        _number_of_records (int): An int that how many records will be returned.
+        _write_time_ending_timestamp (datetime.datetime): A datetime that all records' write time in
+            dataset will be before it.
+        _event_time_starting_timestamp (datetime.datetime): A datetime that all records' event time
+            in dataset will be after it.
+        _event_time_ending_timestamp (datetime.datetime): A datetime that all records' event time in
+            dataset will be before it.
+    """
+
+    _base: Union[FeatureGroup, pd.DataFrame] = attr.ib()
+    _output_path: str = attr.ib()
+    _record_identifier_feature_name: str = attr.ib(default=None)
+    _event_time_identifier_feature_name: str = attr.ib(default=None)
+    _included_feature_names: Sequence[str] = attr.ib(default=None)
+    _kms_key_id: str = attr.ib(default=None)
+
+    _point_in_time_accurate_join: bool = attr.ib(init=False, default=False)
+    _include_duplicated_records: bool = attr.ib(init=False, default=False)
+    _include_deleted_records: bool = attr.ib(init=False, default=False)
+    _number_of_recent_records: int = attr.ib(init=False, default=1)
+    _number_of_records: int = attr.ib(init=False, default=None)
+    _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)
+
+    def point_in_time_accurate_join(self):
+        """Set join type as point in time accurate join.
+
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._point_in_time_accurate_join = True
+        return self
+
+    def include_duplicated_records(self):
+        """Include duplicated records in dataset.
+
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._include_duplicated_records = True
+        return self
+
+    def include_deleted_records(self):
+        """Include deleted records in dataset.
+
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._include_deleted_records = True
+        return self
+
+    def with_number_of_recent_records_by_record_identifier(self, number_of_recent_records: int):
+        """Set number_of_recent_records field with provided input.
+
+        Args:
+            number_of_recent_records (int): An int that how many recent records will be returned for
+                each record identifier.
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._number_of_recent_records = number_of_recent_records
+        return self
+
+    def with_number_of_records_from_query_results(self, number_of_records: int):
+        """Set number_of_records field with provided input.
+
+        Args:
+            number_of_records (int): An int that how many records will be returned.
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._number_of_records = number_of_records
+        return self
+
+    def as_of(self, timestamp: datetime.datetime):
+        """Set write_time_ending_timestamp field with provided input.
+
+        Args:
+            timestamp (datetime.datetime): A datetime that all records' write time in dataset will
+                be before it.
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._write_time_ending_timestamp = timestamp
+        return self
+
+    def with_event_time_range(
+        self,
+        starting_timestamp: datetime.datetime = None,
+        ending_timestamp: datetime.datetime = None,
+    ):
+        """Set event_time_starting_timestamp and event_time_ending_timestamp with provided inputs.
+
+        Args:
+            starting_timestamp (datetime.datetime): A datetime that all records' event time in
+                dataset will be after it (default: None).
+            ending_timestamp (datetime.datetime): A datetime that all records' event time in dataset
+                will be before it (default: None).
+        Returns:
+            This DatasetBuilder object.
+        """
+        self._event_time_starting_timestamp = starting_timestamp
+        self._event_time_ending_timestamp = ending_timestamp
+        return self
diff --git a/src/sagemaker/feature_store/feature_store.py b/src/sagemaker/feature_store/feature_store.py
@@ -18,11 +18,14 @@
 from __future__ import absolute_import
 
 import datetime
-from typing import Dict, Any
+from typing import Any, Dict, Sequence, Union
 
 import attr
+import pandas as pd
 
 from sagemaker import Session
+from sagemaker.feature_store.dataset_builder import DatasetBuilder
+from sagemaker.feature_store.feature_group import FeatureGroup
 
 
 @attr.s
@@ -37,6 +40,50 @@ class FeatureStore:
 
     sagemaker_session: Session = attr.ib(default=Session)
 
+    @staticmethod
+    def create_dataset(
+        base: Union[FeatureGroup, pd.DataFrame],
+        output_path: str,
+        record_identifier_feature_name: str = None,
+        event_time_identifier_feature_name: str = None,
+        included_feature_names: Sequence[str] = None,
+        kms_key_id: str = None,
+    ) -> DatasetBuilder:
+        """Create a Dataset Builder for generating a Dataset.
+
+        Args:
+            base (Union[FeatureGroup, DataFrame]): A base which can be either a FeatureGroup or a
+                pandas.DataFrame and will be used to merge other FeatureGroups and generate a
+                Dataset.
+            output_path (str): An S3 URI which stores the output .csv file.
+            record_identifier_feature_name (str): A string representing the record identifier
+                feature if base is a DataFrame (default: None).
+            event_time_identifier_feature_name (str): A string representing the event time
+                identifier feature if base is a DataFrame (default: None).
+            included_feature_names (List[str]): A list of features to be included in the output
+                (default: None).
+            kms_key_id (str): An KMS key id. If set, will be used to encrypt the result file
+                (default: None).
+
+        Raises:
+            ValueError: Base is a Pandas DataFrame but no record identifier feature name nor event
+                time identifier feature name is provided.
+        """
+        if isinstance(base, pd.DataFrame):
+            if record_identifier_feature_name is None or event_time_identifier_feature_name is None:
+                raise ValueError(
+                    "You must provide a record identifier feature name and an event time "
+                    + "identifier feature name if specify DataFrame as base."
+                )
+        return DatasetBuilder(
+            base,
+            output_path,
+            record_identifier_feature_name,
+            event_time_identifier_feature_name,
+            included_feature_names,
+            kms_key_id,
+        )
+
     def list_feature_groups(
         self,
         name_contains: str = None,
diff --git a/tests/unit/sagemaker/feature_store/test_dataset_builder.py b/tests/unit/sagemaker/feature_store/test_dataset_builder.py
@@ -0,0 +1,71 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import datetime
+
+import pytest
+from mock import Mock
+
+from sagemaker.feature_store.dataset_builder import DatasetBuilder
+
+
+@pytest.fixture
+def feature_group_mock():
+    return Mock()
+
+
+def test_point_in_time_accurate_join(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    dataset_builder.point_in_time_accurate_join()
+    assert dataset_builder._point_in_time_accurate_join
+
+
+def test_include_duplicated_records(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    dataset_builder.include_duplicated_records()
+    assert dataset_builder._include_duplicated_records
+
+
+def test_include_deleted_records(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    dataset_builder.include_deleted_records()
+    assert dataset_builder._include_deleted_records
+
+
+def test_with_number_of_recent_records_by_record_identifier(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    dataset_builder.with_number_of_recent_records_by_record_identifier(5)
+    assert dataset_builder._number_of_recent_records == 5
+
+
+def test_with_number_of_records_from_query_results(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    dataset_builder.with_number_of_records_from_query_results(100)
+    assert dataset_builder._number_of_records == 100
+
+
+def test_as_of(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    time = datetime.datetime.now()
+    dataset_builder.as_of(time)
+    assert dataset_builder._write_time_ending_timestamp == time
+
+
+def test_with_event_time_range(feature_group_mock):
+    dataset_builder = DatasetBuilder(base=feature_group_mock, output_path="file/to/path")
+    start = datetime.datetime.now()
+    end = start + datetime.timedelta(minutes=1)
+    dataset_builder.with_event_time_range(start, end)
+    assert dataset_builder._event_time_starting_timestamp == start
+    assert dataset_builder._event_time_ending_timestamp == end
diff --git a/tests/unit/sagemaker/feature_store/test_feature_store.py b/tests/unit/sagemaker/feature_store/test_feature_store.py
@@ -14,17 +14,83 @@
 
 import datetime
 
+import pandas as pd
 import pytest
 from mock import Mock
 
 from sagemaker.feature_store.feature_store import FeatureStore
 
+DATAFRAME = pd.DataFrame({"feature_1": [420, 380, 390], "feature_2": [50, 40, 45]})
+
 
 @pytest.fixture
 def sagemaker_session_mock():
     return Mock()
 
 
+@pytest.fixture
+def feature_group_mock():
+    return Mock()
+
+
+def test_minimal_create_dataset(sagemaker_session_mock, feature_group_mock):
+    feature_store = FeatureStore(sagemaker_session=sagemaker_session_mock)
+    dataset_builder = feature_store.create_dataset(
+        base=feature_group_mock,
+        output_path="file/to/path",
+    )
+    assert dataset_builder._base == feature_group_mock
+    assert dataset_builder._output_path == "file/to/path"
+
+
+def test_complete_create_dataset(sagemaker_session_mock, feature_group_mock):
+    feature_store = FeatureStore(sagemaker_session=sagemaker_session_mock)
+    dataset_builder = feature_store.create_dataset(
+        base=feature_group_mock,
+        included_feature_names=["feature_1", "feature_2"],
+        output_path="file/to/path",
+        kms_key_id="kms-key-id",
+    )
+    assert dataset_builder._base == feature_group_mock
+    assert dataset_builder._included_feature_names == ["feature_1", "feature_2"]
+    assert dataset_builder._output_path == "file/to/path"
+    assert dataset_builder._kms_key_id == "kms-key-id"
+
+
+def test_create_dataset_with_dataframe(sagemaker_session_mock):
+    feature_store = FeatureStore(sagemaker_session=sagemaker_session_mock)
+    dataset_builder = feature_store.create_dataset(
+        base=DATAFRAME,
+        record_identifier_feature_name="feature_1",
+        event_time_identifier_feature_name="feature_2",
+        included_feature_names=["feature_1", "feature_2"],
+        output_path="file/to/path",
+        kms_key_id="kms-key-id",
+    )
+    assert dataset_builder._base.equals(DATAFRAME)
+    assert dataset_builder._record_identifier_feature_name == "feature_1"
+    assert dataset_builder._event_time_identifier_feature_name == "feature_2"
+    assert dataset_builder._included_feature_names == ["feature_1", "feature_2"]
+    assert dataset_builder._output_path == "file/to/path"
+    assert dataset_builder._kms_key_id == "kms-key-id"
+
+
+def test_create_dataset_with_dataframe_value_error(sagemaker_session_mock):
+    feature_store = FeatureStore(sagemaker_session=sagemaker_session_mock)
+    with pytest.raises(ValueError) as error:
+        feature_store.create_dataset(
+            base=DATAFRAME,
+            included_feature_names=["feature_1", "feature_2"],
+            output_path="file/to/path",
+            kms_key_id="kms-key-id",
+        )
+    assert (
+        "You must provide a record identifier feature name and an event time identifier feature "
+        + "name if specify DataFrame as base."
+        in str(error)
+    )
+
+
 def test_list_feature_groups_with_no_filter(sagemaker_session_mock):
     feature_store = FeatureStore(sagemaker_session=sagemaker_session_mock)
     feature_store.list_feature_groups()
