Skip to content

fix: Clean up temp CSV in AthenaQuery.as_dataframe (v2) - #6345

Open
lucasjia-aws wants to merge 1 commit into
aws:master-v2from
lucasjia-aws:fix/5100-athena-query-temp-cleanup-v2
Open

lucasjia-aws wants to merge 1 commit into
aws:master-v2from
lucasjia-aws:fix/5100-athena-query-temp-cleanup-v2

Conversation

@lucasjia-aws

Copy link
Copy Markdown
Collaborator

Issue

#5100

Problem

sagemaker.feature_store.feature_group.AthenaQuery.as_dataframe() downloads the Athena query result to <tempdir>/<execution_id>.csv, loads it with pd.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_dataframe returns pd.read_csv(filepath_or_buffer=output_filename, ...) directly with no cleanup. DatasetBuilder.to_dataframe in the same package already removes its downloaded file after reading; as_dataframe did not.

Fix

  • Wrap the download and read_csv in try/finally and remove the temporary file in finally, so it is deleted after a successful read and also when the download or parsing raises.
  • Cleanup is best-effort via a small _remove_temp_file helper: it only removes the file if it exists, and an OSError (for example on Windows when a chunksize/iterator reader still holds the file open) is logged as a warning instead of failing the call.
  • Docstring notes that the temporary file is removed.

Testing

  • Added unit tests in 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 when read_csv raises, and a failing os.remove does not break the call.
  • The two new cleanup tests fail without the fix and pass with it.
  • pytest tests/unit/sagemaker/feature_store/test_feature_group.py: 56 passed.
  • black --check and flake8 pass on the changed files.

V3 counterpart: #6344

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

1 active deployment
auto-approve — 2a8ebec8 Deployed Sep 28, 2026 by lucasjia-aws via wait-for-approval #239
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant