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

[lucasjia-aws]

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

[jam-jee]

Unit tests failing, can you check.

[lucasjia-aws]

This test relies on a multi-threaded race condition and is a known flaky test. The same commit passes on Python 3.9, 3.10, and 3.11; Recent PRs #6242 and #6245 also encountered this same failure previously.

[mohamedzeidan2021]

Approving — the better-evidenced of the pair. Reproduced the leak and its absence the same way as #6344, with the adversarial cases (exception path, lazy reader) also clean.

This one has real end-to-end proof: the integ job ran the full suite and six tests that actually execute as_dataframe() passed against live Athena/S3 — test_create_feature_store, the three test_create_feature_store_ingest_with_*_target_stores variants, and both test_get_feature_group_with_* via get_feature_group_as_dataframe. test_create_feature_store calls it inside a while df.shape[0] < 11: retry loop, so repeated real invocations after deletion are covered.

The fix also matches an existing v2 convention — dataset_builder.py already does read-then-os.remove — and improves on it by being exception-safe. 2 new tests fail on reverted source; 56 passed in the module.

Three non-blocking notes:

• test_athena_query_as_dataframe_cleanup_failure_does_not_raise passes on unfixed master-v2, same as its v3 sibling.
• The pre-existing unit test patches tempfile.gettempdir to "tmp", so the fix now calls _remove_temp_file("tmp/query_id.csv") — a relative path. It's a no-op in CI, but it would delete a real ./tmp/query_id.csv if a developer happened to have one in the repo root. Purely theoretical.
• @patch("pandas.read_csv", Mock(side_effect=ValueError("bad csv"))) shares one decorator-time mock across invocations; harmless here, but neighbouring tests use patch(..., side_effect=...) for a fresh mock per call.

dataset_builder.to_dataframe() downloads into ./ and os.removes without try/finally — same leak class, worth a follow-up.