black autoreformat take 2

Author: verayu43

src/sagemaker/feature_store/feature_group.py

@@ -129,17 +129,13 @@ 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"
@@ -199,10 +195,7 @@ 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(
@@ -260,9 +253,7 @@ 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:
@@ -479,8 +470,7 @@ 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,
@@ -489,16 +479,12 @@ 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:
@@ -642,25 +628,19 @@ 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