fix: Clean up temp CSV in AthenaQuery.as_dataframe

[lucasjia-aws]

Issue

#5100

Problem

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

sagemaker-mlops/src/sagemaker/mlops/feature_store/athena_query.py returns pd.read_csv(output_file, ...) directly with no cleanup. The other download paths in the same module (feature_utils.download_csv_from_s3) already remove their local file after reading; as_dataframe was the only one that 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 sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_athena_query.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 sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store (excluding feature_processor, which needs pyspark): 389 passed.
• black --check and flake8 pass on the changed files.

V2 counterpart: #6345

[github-actions]

Claude finished @lucasjia-aws's task in 52s —— View job

---

Code Review — Clean up temp CSV in AthenaQuery.as_dataframe

• Read PR diff and changed files
• Review correctness, tests, and conventions

---

This is a clean, well-scoped fix. The try/finally with a best-effort _remove_temp_file helper correctly removes the temp CSV on the success path, on download failure, and on read_csv parse failure. The tests are thorough and actually assert on filesystem state (using a real tmp_path rather than mocking os.remove away), and cover all three paths. Nice.

A few observations, none blocking:

1. Consistency — actually an improvement. feature_utils.download_csv_from_s3 (feature_utils.py:328) calls a bare os.remove(local_file) with no try/finally, so it leaks the file if read_csv raises and isn't best-effort. This PR's approach is more robust than the existing pattern it's modeled on — worth considering whether download_csv_from_s3 should adopt the same helper in a follow-up (out of scope here).

2. Lazy readers (chunksize / iterator=True) — edge case, already documented. Since as_dataframe(**kwargs) forwards kwargs to pd.read_csv, a caller passing chunksize/iterator=True gets back a lazy TextFileReader, and finally removes the file before the caller iterates:

• On Linux/macOS: os.remove unlinks the file but the reader's open handle stays valid, so iteration still returns correct data. ✅
• On Windows: os.remove on the open file raises PermissionError, which _remove_temp_file swallows — iteration works, but the file is not cleaned up (the leak this PR fixes persists for this narrow case).

No data loss in either case, and the PR body already calls this out. Reasonable to leave as best-effort. If you wanted to be strict, you could detect chunksize/iterator and materialize before cleanup, but that would change return semantics — not worth it.

3. _remove_temp_file scope. Catching OSError (superclass of PermissionError, FileNotFoundError, etc.) and logging at warning is the right call — cleanup should never mask the real result. The os.path.exists guard also avoids a spurious warning when the download fails before creating the file. 👍

Correctness / security / API: No breaking changes to the public signature; behavior is unchanged except for the added cleanup. No credential or input-validation concerns. No new lint issues apparent.

Verdict: Looks good to merge. The only real caveat (Windows + lazy reader) is documented and non-data-affecting.