black autoreformat

Author: verayu43

src/sagemaker/feature_store/feature_group.py

@@ -129,13 +129,17 @@ def as_dataframe(self) -> DataFrame:
         Returns:
             A pandas DataFrame contains the query result.
         """
-        query_state = self.get_query_execution().get("QueryExecution").get("Status").get("State")
+        query_state = (
+            self.get_query_execution().get("QueryExecution").get("Status").get("State")
+        )
         if query_state != "SUCCEEDED":
             if query_state in ("QUEUED", "RUNNING"):
                 raise RuntimeError(
                     f"Current query {self._current_query_execution_id} is still being executed."
                 )
-            raise RuntimeError(f"Failed to execute query {self._current_query_execution_id}")
+            raise RuntimeError(
+                f"Failed to execute query {self._current_query_execution_id}"
+            )
 
         output_filename = os.path.join(
             tempfile.gettempdir(), f"{self._current_query_execution_id}.csv"
@@ -195,7 +199,10 @@ def _ingest_single_batch(
             List of row indices that failed to be ingested.
         """
         retry_config = client_config.retries
-        if "max_attempts" not in retry_config and "total_max_attempts" not in retry_config:
+        if (
+            "max_attempts" not in retry_config
+            and "total_max_attempts" not in retry_config
+        ):
             client_config = copy.deepcopy(client_config)
             client_config.retries = {"max_attempts": 10, "mode": "standard"}
         sagemaker_featurestore_runtime_client = boto3.Session().client(
@@ -207,7 +214,8 @@ def _ingest_single_batch(
         for row in data_frame[start_index:end_index].itertuples():
             record = [
                 FeatureValue(
-                    feature_name=data_frame.columns[index - 1], value_as_string=str(row[index])
+                    feature_name=data_frame.columns[index - 1],
+                    value_as_string=str(row[index]),
                 )
                 for index in range(1, len(row))
                 if pd.notna(row[index])
@@ -252,7 +260,9 @@ def wait(self, timeout=None):
             self._processing_pool.clear()
 
         self._failed_indices = [
-            failed_index for failed_indices in results for failed_index in failed_indices
+            failed_index
+            for failed_indices in results
+            for failed_index in failed_indices
         ]
 
         if len(self._failed_indices) > 0:
@@ -469,7 +479,8 @@ def create(
             record_identifier_name=record_identifier_name,
             event_time_feature_name=event_time_feature_name,
             feature_definitions=[
-                feature_definition.to_dict() for feature_definition in self.feature_definitions
+                feature_definition.to_dict()
+                for feature_definition in self.feature_definitions
             ],
             role_arn=role_arn,
             description=description,
@@ -478,12 +489,16 @@ def create(
 
         # online store configuration
         if enable_online_store:
-            online_store_config = OnlineStoreConfig(enable_online_store=enable_online_store)
+            online_store_config = OnlineStoreConfig(
+                enable_online_store=enable_online_store
+            )
             if online_store_kms_key_id is not None:
-                online_store_config.online_store_security_config = OnlineStoreSecurityConfig(
-                    kms_key_id=online_store_kms_key_id
+                online_store_config.online_store_security_config = (
+                    OnlineStoreSecurityConfig(kms_key_id=online_store_kms_key_id)
                 )
-            create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
+            create_feature_store_args.update(
+                {"online_store_config": online_store_config.to_dict()}
+            )
 
         # offline store configuration
         if s3_uri:
@@ -627,19 +642,25 @@ def athena_query(self) -> AthenaQuery:
             An instance of AthenaQuery initialized with data catalog configurations.
         """
         response = self.describe()
-        data_catalog_config = response.get("OfflineStoreConfig").get("DataCatalogConfig", None)
+        data_catalog_config = response.get("OfflineStoreConfig").get(
+            "DataCatalogConfig", None
+        )
         disable_glue = data_catalog_config.get("DisableGlueTableCreation", False)
         if data_catalog_config:
             query = AthenaQuery(
-                catalog=data_catalog_config["Catalog"] if disable_glue else "AwsDataCatalog",
+                catalog=data_catalog_config["Catalog"]
+                if disable_glue
+                else "AwsDataCatalog",
                 database=data_catalog_config["Database"],
                 table_name=data_catalog_config["TableName"],
                 sagemaker_session=self.sagemaker_session,
             )
             return query
         raise RuntimeError("No metastore is configured with this feature group.")
 
-    def as_hive_ddl(self, database: str = "sagemaker_featurestore", table_name: str = None) -> str:
+    def as_hive_ddl(
+        self, database: str = "sagemaker_featurestore", table_name: str = None
+    ) -> str:
         """Generate Hive DDL commands to define or change structure of tables or databases in Hive.
 
         Schema of the table is generated based on the feature definitions. Columns are named