fix: Clean up temp CSV in AthenaQuery.as_dataframe (v2) - #6345
Open
lucasjia-aws wants to merge 1 commit into
Open
lucasjia-aws wants to merge 1 commit into
lucasjia-aws wants to merge 1 commit into
Conversation
AthenaQuery.as_dataframe() downloaded the Athena query result to <tempdir>/<execution_id>.csv and never removed it, so every query left a file behind and repeated queries could fill the local disk. Wrap the download and read_csv in try/finally and remove the temporary file afterwards, including when the download or parsing fails. Cleanup is best-effort: an OSError is logged as a warning and does not affect the returned DataFrame. Fixes aws#5100
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue
#5100
Problem
sagemaker.feature_store.feature_group.AthenaQuery.as_dataframe()downloads the Athena query result to<tempdir>/<execution_id>.csv, loads it withpd.read_csv, and never deletes the file. Each query leaves a CSV behind, so running many queries (or queries with large results) can fill up local disk space.Root cause
as_dataframereturnspd.read_csv(filepath_or_buffer=output_filename, ...)directly with no cleanup.DatasetBuilder.to_dataframein the same package already removes its downloaded file after reading;as_dataframedid not.Fix
read_csvintry/finallyand remove the temporary file infinally, so it is deleted after a successful read and also when the download or parsing raises._remove_temp_filehelper: it only removes the file if it exists, and anOSError(for example on Windows when achunksize/iteratorreader still holds the file open) is logged as a warning instead of failing the call.Testing
tests/unit/sagemaker/feature_store/test_feature_group.py: the temp file is removed after a successful load (using a real temp directory), it is removed whenread_csvraises, and a failingos.removedoes not break the call.pytest tests/unit/sagemaker/feature_store/test_feature_group.py: 56 passed.black --checkandflake8pass on the changed files.V3 counterpart: #6344