diff --git a/src/sagemaker/feature_store/dataset_builder.py b/src/sagemaker/feature_store/dataset_builder.py index 528d836d0ca..080b672870a 100644 --- 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)