format

aws · Dec 17, 2024 · 7ccab0e · 7ccab0e
1 parent 433783b
commit 7ccab0e
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/sagemaker/feature_store/dataset_builder.py b/src/sagemaker/feature_store/dataset_builder.py
@@ -438,17 +438,13 @@ def to_csv_file(self) -> Tuple[str, str]:
             os.remove(local_file_name)
             temp_table_name = f'dataframe_{temp_id.replace("-", "_")}'
             self._create_temp_table(temp_table_name, desired_s3_folder)
-            query_result = self._run_query(
-                **self._to_athena_query(temp_table_name=temp_table_name)
-            )
+            query_result = self._run_query(**self._to_athena_query(temp_table_name=temp_table_name))
             # TODO: cleanup temp table, need more clarification, keep it for now
             return query_result.get("QueryExecution", {}).get("ResultConfiguration", {}).get(
                 "OutputLocation", None
             ), query_result.get("QueryExecution", {}).get("Query", None)
         if isinstance(self._base, FeatureGroup):
-            query_result = self._run_query(
-                **self._to_athena_query()
-            )
+            query_result = self._run_query(**self._to_athena_query())
             return query_result.get("QueryExecution", {}).get("ResultConfiguration", {}).get(
                 "OutputLocation", None
             ), query_result.get("QueryExecution", {}).get("Query", None)