feature: utilities to facilitate working with Feature Groups

Author: JoseJuan98

doc/api/utility/featuregroup_utils.rst

@@ -0,0 +1,7 @@
+FeatureGroup Utilities
+----------------------
+
+.. automodule:: sagemaker.feature_group_utils
+    :members:
+    :undoc-members:
+    :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -129,9 +129,15 @@ def get_query_execution(self) -> Dict[str, Any]:
             query_execution_id=self._current_query_execution_id
         )
 
-    def as_dataframe(self) -> DataFrame:
+    def as_dataframe(self, **kwargs) -> DataFrame:
         """Download the result of the current query and load it into a DataFrame.
 
+        Args:
+            kwargs: key arguments used for the method pandas.read_csv to be able to have
+                    a better tuning on data.
+                    For more info read
+                    https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html
+
         Returns:
             A pandas DataFrame contains the query result.
         """
@@ -152,7 +158,9 @@ def as_dataframe(self) -> DataFrame:
             query_execution_id=self._current_query_execution_id,
             filename=output_filename,
         )
-        return pd.read_csv(output_filename, delimiter=",")
+
+        kwargs.pop("delimiter", None)
+        return pd.read_csv(filepath_or_buffer=output_filename, delimiter=",", **kwargs)
 
 
 @attr.s

tests/integ/test_feature_store.py

@@ -12,9 +12,9 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import datetime
 import json
 import time
-import datetime
 from contextlib import contextmanager
 
 import boto3
@@ -23,6 +23,7 @@
 import pytest
 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_group import FeatureGroup
 from sagemaker.feature_store.feature_store import FeatureStore
@@ -1055,6 +1056,82 @@ def _wait_for_feature_group_update(feature_group: FeatureGroup):
     print(f"FeatureGroup {feature_group.name} successfully updated.")
 
 
+def test_get_feature_group_with_role_region(
+    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):
+        output = 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,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        feature_group.ingest(
+            data_frame=pandas_data_frame, max_workers=3, max_processes=2, wait=True
+        )
+
+        dataset = get_feature_group_as_dataframe(
+            feature_group_name=feature_group_name,
+            region=region_name,
+            role=role,
+            event_time_feature_name="feature3",
+            latest_ingestion=True,
+            athena_bucket=f"{offline_store_s3_uri}/query",
+        )
+
+    assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
+    assert not dataset.empty
+    assert isinstance(dataset, DataFrame)
+
+
+def test_get_feature_group_with_session(
+    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):
+        output = 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,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        feature_group.ingest(
+            data_frame=pandas_data_frame, max_workers=3, max_processes=2, wait=True
+        )
+
+        dataset = get_feature_group_as_dataframe(
+            feature_group_name=feature_group_name,
+            session=feature_store_session,
+            event_time_feature_name="feature3",
+            latest_ingestion=True,
+            athena_bucket=f"{offline_store_s3_uri}/query",
+            low_memory=False,
+        )  # Using kwargs to pass a parameter to pandas.read_csv
+
+    assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
+    assert not dataset.empty
+    assert isinstance(dataset, DataFrame)
+
+
 @contextmanager
 def cleanup_feature_group(feature_group: FeatureGroup):
     try:

tests/unit/sagemaker/feature_store/test_feature_utils.py

@@ -0,0 +1,101 @@
+# 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.
+# language governing permissions and limitations under the License.
+"""Test for Feature Group Utils"""
+from __future__ import absolute_import
+
+import pandas as pd
+import pytest
+from mock import Mock
+
+from sagemaker.feature_store.feature_utils import (
+    _cast_object_to_string,
+    prepare_fg_from_dataframe_or_file,
+)
+from sagemaker.feature_store.feature_definition import (
+    FeatureTypeEnum,
+)
+from sagemaker.feature_store.feature_group import (
+    FeatureGroup,
+)
+
+
+class PicklableMock(Mock):
+    """Mock class use for tests"""
+
+    def __reduce__(self):
+        """Method from class Mock"""
+        return (Mock, ())
+
+
+@pytest.fixture
+def sagemaker_session_mock():
+    """Fixture Mock class"""
+    return Mock()
+
+
+def test_convert_unsupported_types_to_supported(sagemaker_session_mock):
+    feature_group = FeatureGroup(name="FailedGroup", sagemaker_session=sagemaker_session_mock)
+    df = pd.DataFrame(
+        {
+            "float": pd.Series([2.0], dtype="float64"),
+            "int": pd.Series([2], dtype="int64"),
+            "object": pd.Series(["f1"], dtype="object"),
+        }
+    )
+    # Converting object or O type to string
+    df = _cast_object_to_string(data_frame=df)
+
+    feature_definitions = feature_group.load_feature_definitions(data_frame=df)
+    types = [fd.feature_type for fd in feature_definitions]
+
+    assert types == [
+        FeatureTypeEnum.FRACTIONAL,
+        FeatureTypeEnum.INTEGRAL,
+        FeatureTypeEnum.STRING,
+    ]
+
+
+def test_prepare_fg_from_dataframe(sagemaker_session_mock):
+    very_long_name = "long" * 20
+    df = pd.DataFrame(
+        {
+            "space feature": pd.Series([2.0], dtype="float64"),
+            "dot.feature": pd.Series([2], dtype="int64"),
+            very_long_name: pd.Series(["f1"], dtype="string"),
+        }
+    )
+
+    feature_group = prepare_fg_from_dataframe_or_file(
+        dataframe_or_path=df,
+        session=sagemaker_session_mock,
+        feature_group_name="testFG",
+    )
+
+    names = [fd.feature_name for fd in feature_group.feature_definitions]
+    types = [fd.feature_type for fd in feature_group.feature_definitions]
+
+    assert names == [
+        "space_feature",
+        "dotfeature",
+        very_long_name[:62],
+        "record_id",
+        "data_as_of_date",
+    ]
+    assert types == [
+        FeatureTypeEnum.FRACTIONAL,
+        FeatureTypeEnum.INTEGRAL,
+        FeatureTypeEnum.STRING,
+        FeatureTypeEnum.INTEGRAL,
+        FeatureTypeEnum.FRACTIONAL,
+    ]