Skip to content

Commit badb9bd

Browse files
committed
fix: docstring style and linting
1 parent 34fb6c6 commit badb9bd

File tree

2 files changed

+43
-27
lines changed

2 files changed

+43
-27
lines changed
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
FeatureGroup Utilities
22
----------------------
33

4-
.. automodule:: sagemaker.feature_group_utils
4+
.. automodule:: sagemaker.feature_store.feature_utils
55
:members:
66
:undoc-members:
77
:show-inheritance:

src/sagemaker/feature_store/feature_utils.py

Lines changed: 42 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,15 @@
3131

3232
def get_session_from_role(region: str, assume_role: str = None) -> Session:
3333
"""Method use to get the :class:`sagemaker.session.Session` from a role and a region.
34-
Helpful in case it's invoke from a session with a role without permission it can assume
35-
another role temporarily to perform certain tasks.
34+
35+
Description:
36+
Helpful in case it's invoke from a session with a role without permission it can assume
37+
another role temporarily to perform certain tasks.
38+
3639
Args:
3740
assume_role: role name
3841
region: region name
42+
3943
Returns:
4044
"""
4145
boto_session = boto3.Session(region_name=region)
@@ -74,24 +78,26 @@ def get_session_from_role(region: str, assume_role: str = None) -> Session:
7478

7579

7680
def get_feature_group_as_dataframe(
77-
feature_group_name: str,
78-
athena_bucket: str,
79-
query: str = """SELECT * FROM "sagemaker_featurestore"."#{table}"
80-
WHERE is_deleted=False """,
81-
role: str = None,
82-
region: str = None,
83-
session=None,
84-
event_time_feature_name: str = None,
85-
latest_ingestion: bool = True,
86-
verbose: bool = True,
87-
**pandas_read_csv_kwargs,
81+
feature_group_name: str,
82+
athena_bucket: str,
83+
query: str = """SELECT * FROM "sagemaker_featurestore"."#{table}"
84+
WHERE is_deleted=False """,
85+
role: str = None,
86+
region: str = None,
87+
session=None,
88+
event_time_feature_name: str = None,
89+
latest_ingestion: bool = True,
90+
verbose: bool = True,
91+
**pandas_read_csv_kwargs,
8892
) -> DataFrame:
8993
"""Get a :class:`sagemaker.feature_store.feature_group.FeatureGroup` as a pandas.DataFrame
94+
9095
Description:
9196
Method to run an athena query over a Feature Group in a Feature Store
9297
to retrieve its data.It needs the sagemaker.Session linked to a role
9398
or the role and region used to work Feature Stores.Returns a dataframe
9499
with the data.
100+
95101
Args:
96102
region (str): region of the target Feature Store
97103
feature_group_name (str): feature store name
@@ -110,6 +116,7 @@ def get_feature_group_as_dataframe(
110116
If False it will take whatever is specified in the query, or
111117
if not specify it, it will get all the data that wasn't deleted.
112118
verbose (bool): if True show messages, if False is silent.
119+
113120
Returns:
114121
dataset (pandas.DataFrame): dataset with the data retrieved from feature group
115122
"""
@@ -121,8 +128,8 @@ def get_feature_group_as_dataframe(
121128
if latest_ingestion:
122129
if event_time_feature_name is not None:
123130
query += str(
124-
f"AND {event_time_feature_name}=(SELECT " +
125-
f"MAX({event_time_feature_name}) FROM " +
131+
f"AND {event_time_feature_name}=(SELECT "
132+
f"MAX({event_time_feature_name}) FROM "
126133
'"sagemaker_featurestore"."#{table}")'
127134
)
128135
else:
@@ -169,11 +176,14 @@ def get_feature_group_as_dataframe(
169176

170177
def _format_column_names(data: pandas.DataFrame) -> pandas.DataFrame:
171178
"""Formats the column names for :class:`sagemaker.feature_store.feature_group.FeatureGroup`
179+
172180
Description:
173181
Module to format correctly the name of the columns of a DataFrame
174182
to later generate the features names of a Feature Group
183+
175184
Args:
176185
data (pandas.DataFrame): dataframe used
186+
177187
Returns:
178188
pandas.DataFrame
179189
"""
@@ -183,8 +193,11 @@ def _format_column_names(data: pandas.DataFrame) -> pandas.DataFrame:
183193

184194
def _cast_object_to_string(data_frame: pandas.DataFrame) -> pandas.DataFrame:
185195
"""Cast properly pandas object types to strings
186-
Method to convert 'object' and 'O' column dtypes of a pandas.DataFrame to
187-
a valid string type recognized by Feature Groups.
196+
197+
Description:
198+
Method to convert 'object' and 'O' column dtypes of a pandas.DataFrame to
199+
a valid string type recognized by Feature Groups.
200+
188201
Args:
189202
data_frame: dataframe used
190203
Returns:
@@ -196,23 +209,25 @@ def _cast_object_to_string(data_frame: pandas.DataFrame) -> pandas.DataFrame:
196209

197210

198211
def prepare_fg_from_dataframe_or_file(
199-
dataframe_or_path: Union[str, Path, pandas.DataFrame],
200-
feature_group_name: str,
201-
role: str = None,
202-
region: str = None,
203-
session=None,
204-
record_id: str = "record_id",
205-
event_id: str = "data_as_of_date",
206-
verbose: bool = False,
207-
**pandas_read_csv_kwargs
212+
dataframe_or_path: Union[str, Path, pandas.DataFrame],
213+
feature_group_name: str,
214+
role: str = None,
215+
region: str = None,
216+
session=None,
217+
record_id: str = "record_id",
218+
event_id: str = "data_as_of_date",
219+
verbose: bool = False,
220+
**pandas_read_csv_kwargs,
208221
) -> FeatureGroup:
209222
"""Prepares a dataframe to create a :class:`sagemaker.feature_store.feature_group.FeatureGroup`
223+
210224
Description:
211225
Function to prepare a dataframe for creating a Feature Group from a pandas.DataFrame
212226
or a path to a file with proper dtypes, feature names and mandatory features (record_id,
213227
event_id). It needs the sagemaker.Session linked to a role or the role and region used
214228
to work Feature Stores. If record_id or event_id are not specified it will create ones
215229
by default with the names 'record_id' and 'data_as_of_date'.
230+
216231
Args:
217232
**pandas_read_csv_kwargs (object):
218233
feature_group_name (str): feature group name
@@ -228,6 +243,7 @@ def prepare_fg_from_dataframe_or_file(
228243
role (str) : role used to get the session.
229244
region (str) : region used to get the session.
230245
session (str): session of SageMaker used to work with the feature store
246+
231247
Returns:
232248
:class:`sagemaker.feature_store.feature_group.FeatureGroup`: FG prepared with all
233249
the methods and definitions properly defined

0 commit comments

Comments
 (0)