Skip to content

Commit 0e290db

Browse files
suryans-commitsuryans-commit
committed
Fix Flake8 formatting
1 parent e7c6861 commit 0e290db

File tree

4 files changed

+52
-57
lines changed

4 files changed

+52
-57
lines changed

‎src/sagemaker/feature_store/feature_group.py‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import tempfile
2929
from concurrent.futures import as_completed
3030
from concurrent.futures import ThreadPoolExecutor
31-
from typing import Optional, Sequence, List, Dict, Any, Union
3231
from typing import Optional, Sequence, List, Dict, Any, Union, Iterable
3332
from urllib.parse import urlparse
3433

@@ -331,18 +330,20 @@ def _ingest_row(
331330
"""
332331
try:
333332
record = [
334-
FeatureValue(
335-
feature_name=data_frame.columns[index - 1],
336-
value_as_string_list=IngestionManagerPandas._covert_feature_value_to_string_list(
337-
row[index]
338-
),
339-
)
340-
if IngestionManagerPandas._is_feature_collection_type(
341-
feature_name=data_frame.columns[index - 1],
342-
feature_definitions=feature_definitions,
343-
)
344-
else FeatureValue(
345-
feature_name=data_frame.columns[index - 1], value_as_string=str(row[index])
333+
(
334+
FeatureValue(
335+
feature_name=data_frame.columns[index - 1],
336+
value_as_string_list=IngestionManagerPandas._covert_feature_value_to_string_list(
337+
row[index]
338+
),
339+
)
340+
if IngestionManagerPandas._is_feature_collection_type(
341+
feature_name=data_frame.columns[index - 1],
342+
feature_definitions=feature_definitions,
343+
)
344+
else FeatureValue(
345+
feature_name=data_frame.columns[index - 1], value_as_string=str(row[index])
346+
)
346347
)
347348
for index in range(1, len(row))
348349
if IngestionManagerPandas._feature_value_is_not_none(feature_value=row[index])
@@ -925,21 +926,23 @@ def _determine_collection_list_type(series: Series) -> FeatureTypeEnum | None:
925926
"""
926927

927928
if series.apply(
928-
lambda lst: all(isinstance(x, int) or pd.isna(x) for x in lst)
929-
if is_list_like(lst)
930-
else True
929+
lambda lst: (
930+
all(isinstance(x, int) or pd.isna(x) for x in lst) if is_list_like(lst) else True
931+
)
931932
).all():
932933
return FeatureTypeEnum.INTEGRAL
933934
if series.apply(
934-
lambda lst: all(isinstance(x, (float, int)) or pd.isna(x) for x in lst)
935-
if is_list_like(lst)
936-
else True
935+
lambda lst: (
936+
all(isinstance(x, (float, int)) or pd.isna(x) for x in lst)
937+
if is_list_like(lst)
938+
else True
939+
)
937940
).all():
938941
return FeatureTypeEnum.FRACTIONAL
939942
if series.apply(
940-
lambda lst: all(isinstance(x, str) or pd.isna(x) for x in lst)
941-
if is_list_like(lst)
942-
else True
943+
lambda lst: (
944+
all(isinstance(x, str) or pd.isna(x) for x in lst) if is_list_like(lst) else True
945+
)
943946
).all():
944947
return FeatureTypeEnum.STRING
945948
return None
@@ -1056,9 +1059,9 @@ def put_record(
10561059
return self.sagemaker_session.put_record(
10571060
feature_group_name=self.name,
10581061
record=[value.to_dict() for value in record],
1059-
target_stores=[target_store.value for target_store in target_stores]
1060-
if target_stores
1061-
else None,
1062+
target_stores=(
1063+
[target_store.value for target_store in target_stores] if target_stores else None
1064+
),
10621065
ttl_duration=ttl_duration.to_dict() if ttl_duration is not None else None,
10631066
)
10641067

‎tests/integ/test_feature_store.py‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -751,12 +751,6 @@ def test_create_feature_store_ingest_with_online_target_stores(
751751
)
752752
_wait_for_feature_group_create(feature_group)
753753

754-
resolved_output_s3_uri = (
755-
feature_group.describe()
756-
.get("OfflineStoreConfig")
757-
.get("S3StorageConfig")
758-
.get("ResolvedOutputS3Uri")
759-
)
760754
# Ingest data
761755
ingestion_manager = feature_group.ingest(
762756
data_frame=pandas_data_frame,
@@ -783,7 +777,7 @@ def test_create_feature_store_ingest_with_online_target_stores(
783777
],
784778
target_stores=[TargetStoreEnum.OFFLINE_STORE],
785779
)
786-
assert feature_group.get_record(record_identifier_value_as_string="100.0") == None
780+
assert feature_group.get_record(record_identifier_value_as_string="100.0") is None
787781

788782
# Query the integrated Glue table.
789783
athena_query = feature_group.athena_query()
@@ -823,7 +817,7 @@ def test_put_record_with_target_stores(
823817
feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
824818
feature_group.load_feature_definitions(data_frame=pandas_data_frame)
825819
with cleanup_feature_group(feature_group):
826-
output = feature_group.create(
820+
feature_group.create(
827821
s3_uri=offline_store_s3_uri,
828822
record_identifier_name="feature1",
829823
event_time_feature_name="feature3",
@@ -840,7 +834,7 @@ def test_put_record_with_target_stores(
840834
],
841835
target_stores=[TargetStoreEnum.OFFLINE_STORE],
842836
)
843-
assert feature_group.get_record(record_identifier_value_as_string="100.0") == None
837+
assert feature_group.get_record(record_identifier_value_as_string="100.0") is None
844838

845839
feature_group.put_record(
846840
record=[
@@ -2188,7 +2182,7 @@ def test_ingest_in_memory_multi_process_with_collection_types(
21882182
online_storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
21892183
)
21902184

2191-
with (cleanup_feature_group(feature_group)):
2185+
with cleanup_feature_group(feature_group):
21922186
output = feature_group.create(
21932187
record_identifier_name="feature1",
21942188
event_time_feature_name="feature3",
@@ -2229,7 +2223,7 @@ def test_ingest_in_memory_multi_process_with_collection_types(
22292223
pandas_data_frame_with_collection_type.loc[
22302224
len(pandas_data_frame_with_collection_type)
22312225
] = new_row_data
2232-
with pytest.raises(IngestionError) as error:
2226+
with pytest.raises(IngestionError):
22332227
feature_group.ingest(
22342228
data_frame=pandas_data_frame_with_collection_type,
22352229
max_workers=1,
@@ -2252,7 +2246,7 @@ def test_ingest_in_memory_single_process_with_collection_types(
22522246
online_storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
22532247
)
22542248

2255-
with (cleanup_feature_group(feature_group)):
2249+
with cleanup_feature_group(feature_group):
22562250
output = feature_group.create(
22572251
record_identifier_name="feature1",
22582252
event_time_feature_name="feature3",
@@ -2292,7 +2286,7 @@ def test_ingest_in_memory_single_process_with_collection_types(
22922286
pandas_data_frame_with_collection_type.loc[
22932287
len(pandas_data_frame_with_collection_type)
22942288
] = new_row_data
2295-
with pytest.raises(IngestionError) as error:
2289+
with pytest.raises(IngestionError):
22962290
feature_group.ingest(
22972291
data_frame=pandas_data_frame_with_collection_type,
22982292
max_workers=1,
@@ -2315,7 +2309,7 @@ def test_ingest_standard_multi_process_with_collection_types(
23152309
online_storage_type=OnlineStoreStorageTypeEnum.STANDARD,
23162310
)
23172311

2318-
with (cleanup_feature_group(feature_group)):
2312+
with cleanup_feature_group(feature_group):
23192313
output = feature_group.create(
23202314
record_identifier_name="feature1",
23212315
event_time_feature_name="feature3",

‎tests/unit/sagemaker/feature_store/test_feature_group.py‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -104,15 +104,6 @@ def feature_group_dummy_definition_dict():
104104
}
105105

106106

107-
@pytest.fixture
108-
def feature_group_dummy_definitions():
109-
return [
110-
FractionalFeatureDefinition(feature_name="feature1"),
111-
IntegralFeatureDefinition(feature_name="feature2"),
112-
StringFeatureDefinition(feature_name="feature3"),
113-
]
114-
115-
116107
@pytest.fixture
117108
def data_frame_with_collection_type():
118109
df = pd.DataFrame(
@@ -1097,7 +1088,7 @@ def test_as_hive_ddl(create_table_ddl, feature_group_dummy_definitions, sagemake
10971088
"sagemaker.feature_store.feature_group.IngestionManagerPandas._run_multi_process",
10981089
MagicMock(),
10991090
)
1100-
def test_ingestion_manager_run_success():
1091+
def test_ingestion_manager__run_multi_process_success():
11011092
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
11021093
manager = IngestionManagerPandas(
11031094
feature_group_name="MyGroup",

‎tests/unit/test_session.py‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5275,12 +5275,16 @@ def test_feature_group_put_record(sagemaker_session_with_featurestore_runtime_cl
52755275
feature_group_name="MyFeatureGroup",
52765276
record=[{"FeatureName": "feature1", "ValueAsString": "value1"}],
52775277
)
5278-
assert sagemaker_session_with_featurestore_runtime_client.sagemaker_featurestore_runtime_client.put_record.called_with(
5279-
FeatureGroupName="MyFeatureGroup",
5280-
record=[{"FeatureName": "feature1", "ValueAsString": "value1"}],
5278+
assert (
5279+
sagemaker_session_with_featurestore_runtime_client.sagemaker_featurestore_runtime_client
5280+
.put_record.called_with(
5281+
FeatureGroupName="MyFeatureGroup",
5282+
record=[{"FeatureName": "feature1", "ValueAsString": "value1"}],
5283+
)
52815284
)
52825285

52835286

5287+
52845288
def test_feature_group_put_record_with_ttl_and_target_stores(
52855289
sagemaker_session_with_featurestore_runtime_client,
52865290
):
@@ -5290,11 +5294,14 @@ def test_feature_group_put_record_with_ttl_and_target_stores(
52905294
ttl_duration={"Unit": "Seconds", "Value": 123},
52915295
target_stores=["OnlineStore", "OfflineStore"],
52925296
)
5293-
assert sagemaker_session_with_featurestore_runtime_client.sagemaker_featurestore_runtime_client.put_record.called_with(
5294-
FeatureGroupName="MyFeatureGroup",
5295-
record=[{"FeatureName": "feature1", "ValueAsString": "value1"}],
5296-
target_stores=["OnlineStore", "OfflineStore"],
5297-
ttl_duration={"Unit": "Seconds", "Value": 123},
5297+
assert (
5298+
sagemaker_session_with_featurestore_runtime_client.sagemaker_featurestore_runtime_client
5299+
.put_record.called_with(
5300+
FeatureGroupName="MyFeatureGroup",
5301+
record=[{"FeatureName": "feature1", "ValueAsString": "value1"}],
5302+
target_stores=["OnlineStore", "OfflineStore"],
5303+
ttl_duration={"Unit": "Seconds", "Value": 123},
5304+
)
52985305
)
52995306

53005307

0 commit comments

Comments
 (0)