From ff8012be46704399e8cd5a438417463df46bf88e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:03:40 +0900 Subject: [PATCH 1/8] Remove and rewrite redundant Polars cursor tests Remove the as_arrow() wrapper tests and the row-count-only chunk and chunked fetch tests, whose contracts test_whole_result_reused, test_iter_chunks_data_consistency, test_unload_chunks, and test_fetch_across_chunks check with values. The chunk-size bound joins test_iter_chunks_data_consistency. The offline reader tests build the result set through its constructor over local result files and read it through iter_chunks() instead of calling private readers on an instance without __init__. The mocked storage-options tests become a live test whose invalid credentials fail the read, and an offline test that the given options reach every Polars read of an UNLOAD result. The strict and ignore_errors timestamp reads are one test, which replaces the converter's own non-strict test. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/polars/test_async_cursor.py | 14 - tests/pyathena/polars/test_converter.py | 9 - tests/pyathena/polars/test_cursor.py | 208 ++-------- tests/pyathena/polars/test_result_set.py | 455 ++++++++------------- 4 files changed, 201 insertions(+), 485 deletions(-) diff --git a/tests/pyathena/polars/test_async_cursor.py b/tests/pyathena/polars/test_async_cursor.py index 68290a3a1..12673f425 100644 --- a/tests/pyathena/polars/test_async_cursor.py +++ b/tests/pyathena/polars/test_async_cursor.py @@ -243,18 +243,6 @@ def test_many_as_polars(self, async_polars_cursor): assert df.height == 10000 assert df.width == 1 - @pytest.mark.parametrize( - "async_polars_cursor", - [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], - indirect=["async_polars_cursor"], - ) - def test_as_arrow(self, async_polars_cursor): - query_id, future = async_polars_cursor.execute("SELECT * FROM one_row") - assert query_id is not None - table = future.result().as_arrow() - assert table.num_rows == 1 - assert table.num_columns == 1 - def test_cancel(self, async_polars_cursor): query_id, future = async_polars_cursor.execute( """ @@ -348,7 +336,6 @@ def test_read_options(self, execute_kwargs): converter=MagicMock(), formatter=MagicMock(), retry_config=RetryConfig(), - max_workers=1, **cursor_kwargs, ) as cursor, patch.object(AsyncPolarsCursor, "_execute", return_value="query_id"), @@ -361,4 +348,3 @@ def test_read_options(self, execute_kwargs): expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/polars/test_converter.py b/tests/pyathena/polars/test_converter.py index d103d611f..8dbb07648 100644 --- a/tests/pyathena/polars/test_converter.py +++ b/tests/pyathena/polars/test_converter.py @@ -61,12 +61,3 @@ def test_to_datetimes(dtype, microseconds, lazy): None, None, ] - - -def test_to_datetimes_not_strict(): - """Text that is not a timestamp becomes null when the conversion is not strict.""" - df = pl.DataFrame({"t": ["2020-01-02 03:04:05.123456789012", "invalid"]}) - with pytest.raises(pl.exceptions.InvalidOperationError): - _to_datetimes(df, {"t": pl.Datetime}) - result = _to_datetimes(df, {"t": pl.Datetime}, strict=False) - assert result["t"].to_list() == [datetime(2020, 1, 2, 3, 4, 5, 123456), None] diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 164d785af..c9b0f38a6 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -18,7 +18,7 @@ import polars as pl import pytest -from pyathena.error import DatabaseError, DataError, ProgrammingError +from pyathena.error import DatabaseError, DataError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.cursor import PolarsCursor @@ -37,6 +37,11 @@ # Polars 2.0 reads INT96 at the time unit of the read schema; Polars 1.x at nanoseconds. _POLARS_2 = int(pl.__version__.split(".")[0]) >= 2 +_INVALID_OBJECT_STORE_OPTIONS = { + "aws_access_key_id": "invalid", + "aws_secret_access_key": "invalid", + "aws_region": ENV.region_name, +} class TestPolarsCursor: @@ -142,6 +147,35 @@ def test_as_polars_with_read_kwargs(self, polars_cursor): ).as_polars() assert df.to_dicts() == [{"number_of_rows": "1"}] + @pytest.mark.parametrize( + ("polars_cursor", "chunksize", "storage_options"), + [ + # A whole CSV result file is read through fsspec with PyAthena's filesystem. + ( + {"cursor_kwargs": {"unload": False}}, + None, + {"key": "invalid", "secret": "invalid", "skip_instance_cache": True}, + ), + # The other results are read with Polars' native object store. + ({"cursor_kwargs": {"unload": False}}, 1, _INVALID_OBJECT_STORE_OPTIONS), + ({"cursor_kwargs": {"unload": True}}, None, _INVALID_OBJECT_STORE_OPTIONS), + ], + indirect=["polars_cursor"], + ids=["csv", "csv-chunks", "unload"], + ) + def test_storage_options_replace_defaults(self, polars_cursor, chunksize, storage_options): + """storage_options given to execute() replace PyAthena's S3 settings for the read. + + Their invalid credentials fail the read. + """ + with pytest.raises(OperationalError): + # A whole result is read in execute(), and chunks as they are iterated. + list( + polars_cursor.execute( + "SELECT * FROM one_row", chunksize=chunksize, storage_options=storage_options + ).iter_chunks() + ) + @pytest.mark.parametrize( "polars_cursor", [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], @@ -502,16 +536,6 @@ def test_complex_unload_as_polars(self, polars_cursor): Decimal("0.1"), ) - @pytest.mark.parametrize( - "polars_cursor", - [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], - indirect=["polars_cursor"], - ) - def test_as_arrow(self, polars_cursor): - table = polars_cursor.execute("SELECT * FROM one_row").as_arrow() - assert table.num_rows == 1 - assert table.num_columns == 1 - def test_cancel(self, polars_cursor): def cancel(c): time.sleep(random.randint(5, 10)) @@ -633,27 +657,6 @@ def test_callback(query_id: str): assert callback_results[0] == polars_cursor.query_id assert polars_cursor.query_id is not None - def test_iter_chunks(self): - """Test chunked iteration over query results.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=5) - cursor.execute("SELECT * FROM many_rows LIMIT 15") - chunks = list(cursor.iter_chunks()) - assert len(chunks) > 0 - total_rows = sum(chunk.height for chunk in chunks) - assert total_rows == 15 - for chunk in chunks: - assert isinstance(chunk, pl.DataFrame) - - def test_iter_chunks_without_chunksize(self, polars_cursor): - """Test that iter_chunks works without chunksize, yielding entire DataFrame.""" - polars_cursor.execute("SELECT * FROM one_row") - chunks = list(polars_cursor.iter_chunks()) - # Without chunksize, yields entire DataFrame as single chunk - assert len(chunks) == 1 - assert isinstance(chunks[0], pl.DataFrame) - assert chunks[0].height == 1 - def test_whole_result_reused(self, polars_cursor): """Test that as_polars(), as_arrow(), and iter_chunks() do not consume the rows.""" polars_cursor.execute("SELECT number FROM (VALUES (1), (2), (3)) AS t(number)") @@ -681,35 +684,6 @@ def test_close_stops_chunks(self, polars_cursor): result_set.close() assert list(chunks) == [] - def test_iter_chunks_many_rows(self): - """Test chunked iteration with many rows.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=1000) - cursor.execute("SELECT * FROM many_rows") - chunks = list(cursor.iter_chunks()) - total_rows = sum(chunk.height for chunk in chunks) - assert total_rows == 10000 - assert len(chunks) >= 10 # At least 10 chunks with chunksize=1000 - - @pytest.mark.parametrize( - "polars_cursor", - [ - { - "cursor_kwargs": {"unload": True, "chunksize": 5}, - }, - ], - indirect=["polars_cursor"], - ) - def test_iter_chunks_unload(self, polars_cursor): - """Test chunked iteration with UNLOAD (Parquet).""" - polars_cursor.execute("SELECT * FROM many_rows LIMIT 15") - chunks = list(polars_cursor.iter_chunks()) - assert len(chunks) > 0 - total_rows = sum(chunk.height for chunk in chunks) - assert total_rows == 15 - for chunk in chunks: - assert isinstance(chunk, pl.DataFrame) - @pytest.mark.parametrize( ("polars_cursor", "chunksize", "options", "sizes"), [ @@ -786,67 +760,9 @@ def test_iter_chunks_data_consistency(self): # Should have the same data (sort for comparison) assert regular_df.sort("a").equals(combined_df.sort("a")) - # Should have multiple chunks + # Should have multiple chunks, none above the chunk size assert len(chunked_dfs) > 1 - - def test_iter_chunks_chunk_sizes(self): - """Test that chunks have correct sizes.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=10) - cursor.execute("SELECT * FROM many_rows LIMIT 50") - - chunk_sizes = [] - total_rows = 0 - - for chunk in cursor.iter_chunks(): - chunk_size = chunk.height - chunk_sizes.append(chunk_size) - total_rows += chunk_size - - # Each chunk should not exceed chunksize - assert chunk_size <= 10 - - # Should have processed all 50 rows - assert total_rows == 50 - - # Should have multiple chunks - assert len(chunk_sizes) > 1 - - def test_fetchone_with_chunksize(self): - """Test that fetchone works correctly with chunksize enabled.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=5) - cursor.execute("SELECT * FROM many_rows LIMIT 15") - - rows = [] - while True: - row = cursor.fetchone() - if row is None: - break - rows.append(row) - - assert len(rows) == 15 - - def test_fetchmany_with_chunksize(self): - """Test that fetchmany works correctly with chunksize enabled.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=5) - cursor.execute("SELECT * FROM many_rows LIMIT 15") - - batch1 = cursor.fetchmany(10) - batch2 = cursor.fetchmany(10) - - assert len(batch1) == 10 - assert len(batch2) == 5 - - def test_fetchall_with_chunksize(self): - """Test that fetchall works correctly with chunksize enabled.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=5) - cursor.execute("SELECT * FROM many_rows LIMIT 15") - - rows = cursor.fetchall() - assert len(rows) == 15 + assert all(chunk.height <= 25 for chunk in chunked_dfs) @pytest.mark.parametrize( ("polars_cursor", "ordered"), @@ -877,54 +793,6 @@ def test_fetch_across_chunks(self, polars_cursor, ordered): (5, "e"), ] - def test_iterator_with_chunksize(self): - """Test that cursor iteration works correctly with chunksize enabled.""" - with contextlib.closing(connect(schema_name=ENV.schema)) as conn: - cursor = conn.cursor(PolarsCursor, chunksize=5) - cursor.execute("SELECT * FROM many_rows LIMIT 15") - - rows = list(cursor) - assert len(rows) == 15 - - @pytest.mark.parametrize( - "polars_cursor", - [ - { - "cursor_kwargs": {"unload": True, "chunksize": 5}, - }, - ], - indirect=["polars_cursor"], - ) - def test_fetchone_with_chunksize_unload(self, polars_cursor): - """Test that fetchone works correctly with chunksize and unload enabled.""" - polars_cursor.execute("SELECT * FROM many_rows LIMIT 15") - - rows = [] - while True: - row = polars_cursor.fetchone() - if row is None: - break - rows.append(row) - - assert len(rows) == 15 - assert [d[0] for d in polars_cursor.description] == ["a"] - assert all(isinstance(row[0], int) for row in rows) - - @pytest.mark.parametrize( - "polars_cursor", - [ - { - "cursor_kwargs": {"unload": True, "chunksize": 5}, - }, - ], - indirect=["polars_cursor"], - ) - def test_iterator_with_chunksize_unload(self, polars_cursor): - """Test that cursor iteration works with chunksize and unload enabled.""" - polars_cursor.execute("SELECT * FROM many_rows LIMIT 15") - rows = list(polars_cursor) - assert len(rows) == 15 - @pytest.mark.parametrize( ("polars_cursor", "query", "expected"), [ @@ -1059,7 +927,7 @@ def test_ambiguous_time_zone_value(self, polars_cursor): The table keeps the text, as for every TIMESTAMP WITH TIME ZONE value. """ polars_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match="is ambiguous"): + with pytest.raises(DataError): polars_cursor.fetchall() @pytest.mark.skipif(not ENV.managed_work_group, reason="AWS_ATHENA_MANAGED_WORKGROUP not set") diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index e90f1f79f..12723f9e8 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -6,10 +6,11 @@ # SPDX-License-Identifier: MIT import inspect -from collections.abc import Iterator from datetime import UTC, date, datetime, time, timedelta from decimal import Decimal -from unittest.mock import MagicMock, PropertyMock, patch +from io import BytesIO +from typing import Any +from unittest.mock import MagicMock, patch import polars as pl import pytest @@ -32,31 +33,81 @@ ) -def _chunked_result_set() -> AthenaPolarsResultSet: - """Create a chunked result set without calling Athena. +def _local_result_set( + column_types: list[tuple[str, str]], + output_location: str = "s3://bucket/result.csv", + data_manifest: list[str] | None = None, + chunksize: int | None = None, + **kwargs: Any, +) -> AthenaPolarsResultSet: + """Create a result set of a succeeded query whose result files are local files. + + The given columns stand in for the GetQueryResults metadata, and the HEAD request of + the CSV result file finds it not empty. The connection's S3 client returns the data + manifest of an UNLOAD result, and PyAthena's filesystem reads the local files. + + Args: + column_types: The names and Athena types of the result columns. + output_location: The local path of the CSV result file. + data_manifest: The local paths of the Parquet files of an UNLOAD result. + chunksize: The chunk size of the result set. + **kwargs: The read arguments given to ``execute()``. Returns: - An AthenaPolarsResultSet with only the attributes used by the chunk readers. + The AthenaPolarsResultSet. """ - result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ - result_set._chunksize = 10_000 - result_set._cache_type = None - result_set._kwargs = {} - result_set._metadata = None - result_set._connection = MagicMock() - return result_set + unload = data_manifest is not None + query_execution = AthenaQueryExecution( + { + "QueryExecution": { + "QueryExecutionId": "query_id", + "Query": "UNLOAD (SELECT 1) TO 's3://bucket/unload/'" if unload else "SELECT 1", + "Status": {"State": AthenaQueryExecution.STATE_SUCCEEDED}, + "ResultConfiguration": {"OutputLocation": output_location}, + "Statistics": {"DataManifestLocation": "s3://bucket/manifest.csv"}, + } + } + ) + connection = MagicMock(region_name=None) + connection.session.get_credentials.return_value = None + connection.s3_client.get_object.return_value = { + "Body": BytesIO("\n".join(data_manifest or []).encode()) + } + metadata = tuple( + {"Name": name, "Type": type_, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} + for name, type_ in column_types + ) + + def pre_fetch(self: AthenaPolarsResultSet) -> None: + self._metadata = metadata + + with ( + patch.object(AthenaPolarsResultSet, "_pre_fetch", pre_fetch), + patch.object(AthenaPolarsResultSet, "_get_content_length", return_value=1), + patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), + ): + return AthenaPolarsResultSet( + connection=connection, + converter=DefaultPolarsTypeConverter(), + query_execution=query_execution, + arraysize=3, + retry_config=RetryConfig(), + unload=unload, + chunksize=chunksize, + **kwargs, + ) -def _read_all(result: pl.DataFrame | Iterator[pl.DataFrame]) -> pl.DataFrame: - """Read the result of a whole read or a chunk reader into one DataFrame. +def _read_all(result_set: AthenaPolarsResultSet) -> pl.DataFrame: + """Read the rows of a result set into one DataFrame. Args: - result: A DataFrame, or an iterator of DataFrame chunks. + result_set: A result set read whole or in chunks. Returns: - The DataFrame, or the chunks concatenated. + The DataFrame of the result, or its chunks concatenated. """ - return result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) + return pl.concat(list(result_set.iter_chunks())) def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPolarsResultSet: @@ -199,13 +250,6 @@ def test_fetchmany_size_above_maxsize(self): assert result_set.fetchmany(10**100) == [(1,), (2,)] assert result_set.rownumber == 2 - def test_fetchmany_size_not_integer(self): - """A size that is not an integer raises TypeError without consuming rows.""" - result_set = _fetching_result_set([pl.DataFrame({"n": [1, 2]})], ["n"], arraysize=3) - with pytest.raises(TypeError): - result_set.fetchmany(1.5) - assert result_set.fetchall() == [(1,), (2,)] - def test_fetch_after_close(self): """The fetch methods return no rows after close().""" result_set = _fetching_result_set([pl.DataFrame({"n": [1, 2]})], ["n"], arraysize=3) @@ -215,42 +259,27 @@ def test_fetch_after_close(self): assert result_set.fetchall() == [] assert result_set.rownumber is None - def test_iter_csv_chunks_raises_when_read_fails_partway(self, tmp_path): - """A CSV read that fails partway through the data raises instead of ending early.""" + def test_csv_chunks_raise_when_read_fails_partway(self, tmp_path): + """A chunked CSV read that fails partway through the data raises instead of ending.""" path = tmp_path / "result.csv" path.write_text( "a\n" + "".join(f"{i}\n" for i in range(_ROWS_BEFORE_FAILURE)) + "not-a-number\n" ) - result_set = _chunked_result_set() - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_dtypes", - new_callable=PropertyMock, - return_value={"a": pl.Int64}, - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - pytest.raises(OperationalError, match="not-a-number"), - ): - list(result_set._iter_csv_chunks()) + result_set = _local_result_set([("a", "integer")], str(path), chunksize=10_000) + with pytest.raises(OperationalError): + list(result_set.iter_chunks()) @pytest.mark.parametrize( "read_options", [{}, {"low_memory": True}], ids=["parquet-chunk-reader", "scan-parquet"], ) - def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path, read_options): - """A Parquet read that fails partway through the data raises instead of ending early. + def test_parquet_chunks_raise_when_read_fails_partway(self, tmp_path, read_options): + """A chunked UNLOAD read that fails partway through the data raises instead of ending. - Without read options, the files of the data manifest are read with - ParquetChunkReader; with them, Polars scans the unload location. + Without read options, the files of the data manifest are read one at a time, + and the chunks of the valid file come before the error; with them, Polars scans + the unload location. """ pl.DataFrame({"a": range(_ROWS_BEFORE_FAILURE)}).write_parquet( tmp_path / "0.parquet", row_group_size=10_000 @@ -258,202 +287,92 @@ def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path, read valid = (tmp_path / "0.parquet").read_bytes() # Keep the footer magic but corrupt the metadata so the second file fails to read. (tmp_path / "1.parquet").write_bytes(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:]) - result_set = _chunked_result_set() - result_set._kwargs = dict(read_options) - result_set._unload_location = f"{tmp_path}/" - result_set._data_manifest = [str(tmp_path / "0.parquet"), str(tmp_path / "1.parquet")] + result_set = _local_result_set( + [], + data_manifest=[str(tmp_path / "0.parquet"), str(tmp_path / "1.parquet")], + chunksize=10_000, + **read_options, + ) chunks: list[pl.DataFrame] = [] - with ( - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), - patch("polars.scan_parquet", wraps=pl.scan_parquet) as scan_parquet, - pytest.raises(OperationalError) as error, - ): + with pytest.raises(OperationalError): # extend() keeps the chunks read before the failure. - chunks.extend(result_set._iter_parquet_chunks(pl.Schema({"a": pl.Int64}))) - # The error comes from the corrupt file, after the valid one was read. - assert not isinstance(error.value.__cause__, AttributeError) - assert scan_parquet.called is bool(read_options) + chunks.extend(result_set.iter_chunks()) if not read_options: assert sum(chunk.height for chunk in chunks) == _ROWS_BEFORE_FAILURE - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) - def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader): + @pytest.mark.parametrize("chunksize", [None, 1]) + def test_unload_storage_options_replace_defaults(self, tmp_path, chunksize): + """storage_options given to execute() reach every Polars read of an UNLOAD result. + + Polars ignores them for local files, so the reads are observed. + """ + pl.DataFrame({"a": [1, 2]}).write_parquet(tmp_path / "0.parquet") + storage_options = {"aws_region": "us-west-2"} + with ( + patch("polars.read_parquet", wraps=pl.read_parquet) as read_parquet, + patch("polars.scan_parquet", wraps=pl.scan_parquet) as scan_parquet, + ): + result_set = _local_result_set( + [], + data_manifest=[str(tmp_path / "0.parquet")], + chunksize=chunksize, + storage_options=storage_options, + ) + assert result_set.fetchall() == [(1,), (2,)] + calls = read_parquet.call_args_list + scan_parquet.call_args_list + # The schema scan, and the whole read or the scan of the chunks. + assert len(calls) == 2 + assert all(call.kwargs["storage_options"] is storage_options for call in calls) + + @pytest.mark.parametrize("chunksize", [None, 10_000]) + def test_csv_read_kwargs_replace_defaults(self, tmp_path, chunksize): """Read arguments given to execute() replace the ones the result set chooses.""" path = tmp_path / "result.csv" path.write_text("1;x\n2;y\n") # Polars names headerless columns from column_1 before 2.0 and from column_0 since. first, second = pl.read_csv(b"1;x\n", separator=";", has_header=False).columns - result_set = _chunked_result_set() - result_set._kwargs = { - "separator": ";", - "has_header": False, - "schema_overrides": {first: pl.Utf8}, - } - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_dtypes", - new_callable=PropertyMock, - return_value={"1;x": pl.Int64}, - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - ): - df = _read_all(getattr(result_set, reader)()) + result_set = _local_result_set( + [("a", "integer"), ("b", "varchar")], + str(path), + chunksize=chunksize, + separator=";", + has_header=False, + schema_overrides={first: pl.Utf8}, + ) + df = _read_all(result_set) assert df.to_dict(as_series=False) == {first: ["1", "2"], second: ["x", "y"]} - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) - def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): - """schema_overrides given with new_columns reach Polars with them for a .txt file.""" + @pytest.mark.parametrize("chunksize", [None, 10_000]) + def test_txt_new_columns_with_schema_overrides(self, tmp_path, chunksize): + """schema_overrides given with new_columns apply to the new names for a .txt file.""" path = tmp_path / "result.txt" path.write_text("001\tx\n") - result_set = _chunked_result_set() # Polars 2.0 rejects new_columns that do not name every column of the file. - result_set._kwargs = {"new_columns": ["z", "w"], "schema_overrides": {"z": pl.Utf8}} - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object(AthenaPolarsResultSet, "_get_column_names", return_value=["a", "b"]), - patch.object( - AthenaPolarsResultSet, "_csv_dtypes", new_callable=PropertyMock, return_value={} - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - ): - df = _read_all(getattr(result_set, reader)()) + result_set = _local_result_set( + [("a", "integer"), ("b", "varchar")], + str(path), + chunksize=chunksize, + new_columns=["z", "w"], + schema_overrides={"z": pl.Utf8}, + ) + df = _read_all(result_set) assert df.to_dict(as_series=False) == {"z": ["001"], "w": ["x"]} @pytest.mark.parametrize( - ("reader", "function"), - [ - ("_iter_csv_chunks", "scan_csv"), - ("_read_parquet", "read_parquet"), - ("_iter_parquet_chunks", "scan_parquet"), - ("_read_parquet_schema", "scan_parquet"), - ], - ) - def test_storage_options_replace_defaults(self, reader, function): - """storage_options given to execute() replace PyAthena's without computing them.""" - result_set = _chunked_result_set() - result_set._unload_location = "s3://bucket/unload/" - result_set._kwargs = {"storage_options": {"anon": True}} - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value="s3://bucket/result.csv", - ), - patch.object( - AthenaPolarsResultSet, "_csv_dtypes", new_callable=PropertyMock, return_value={} - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - side_effect=AssertionError("replaced storage options were computed"), - ), - patch.object( - AthenaPolarsResultSet, - "_get_parquet_storage_options", - side_effect=AssertionError("replaced storage options were computed"), - ), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - patch.object(AthenaPolarsResultSet, "_prepare_parquet_location", return_value=True), - patch(f"polars.{function}") as read, - patch("pyathena.polars.result_set.to_column_info"), - ): - # The Parquet readers take the schema of the files. - args = (pl.Schema(),) if reader in ("_read_parquet", "_iter_parquet_chunks") else () - result = getattr(result_set, reader)(*args) - if not isinstance(result, (pl.DataFrame, tuple)): - list(result) - assert read.call_args.kwargs["storage_options"] == {"anon": True} - - def test_csv_storage_options_replace_defaults(self): - """storage_options given to execute() open a CSV result through fsspec, not Polars.""" - result_set = _chunked_result_set() - result_set._kwargs = {"storage_options": {"anon": True}} - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value="s3://bucket/result.csv", - ), - patch.object( - AthenaPolarsResultSet, "_csv_dtypes", new_callable=PropertyMock, return_value={} - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - side_effect=AssertionError("replaced storage options were computed"), - ), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - patch("pyathena.polars.result_set.fsspec.open") as open_, - ): - open_.return_value.__enter__.return_value.read.return_value = b"a\n1\n" - df = result_set._read_csv() - open_.assert_called_once_with("s3://bucket/result.csv", "rb", anon=True) - assert df.to_dict(as_series=False) == {"a": [1]} - - @pytest.mark.parametrize( - ("reader", "kwargs", "expected"), + ("chunksize", "kwargs", "expected"), [ - ("_read_csv", {}, {"t": _TRUNCATED_TIMESTAMPS}), - ("_iter_csv_chunks", {}, {"t": _TRUNCATED_TIMESTAMPS}), - ("_read_csv", {"columns": ["v"]}, {}), - ("_read_csv", {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), - ("_iter_csv_chunks", {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), - ("_read_csv", {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), - ( - "_iter_csv_chunks", - {"schema": {"t": pl.String, "v": pl.String}}, - {"t": _TIMESTAMP_TEXTS}, - ), - ( - "_read_csv", - {"with_column_names": lambda names: [n.upper() for n in names]}, - OperationalError, - ), - ( - "_iter_csv_chunks", - {"with_column_names": lambda names: [n.upper() for n in names]}, - OperationalError, - ), + (None, {}, {"t": _TRUNCATED_TIMESTAMPS}), + (10_000, {}, {"t": _TRUNCATED_TIMESTAMPS}), + (None, {"columns": ["v"]}, {}), + (None, {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), + (10_000, {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), + (None, {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), + (10_000, {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), + (None, {"with_column_names": lambda names: [n.upper() for n in names]}, None), + (10_000, {"with_column_names": lambda names: [n.upper() for n in names]}, None), ], ) - def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected): + def test_read_csv_truncates_timestamps(self, tmp_path, chunksize, kwargs, expected): """Timestamps are read as text and truncated to the time unit. A whole read does so when parsing them fails; a chunked read always does. @@ -464,74 +383,35 @@ def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected) """ path = tmp_path / "result.csv" path.write_bytes(b'"t","v"\n"2020-01-02 03:04:05.123456789012","x"\n,"y"\n') - result_set = _chunked_result_set() - result_set._query_execution = None - result_set._converter = DefaultPolarsTypeConverter() - result_set._kwargs = kwargs - result_set._metadata = tuple( - {"Name": n, "Type": t, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} - for n, t in (("t", "timestamp"), ("v", "varchar")) - ) - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - ): - if expected is OperationalError: - with pytest.raises(OperationalError): - _read_all(getattr(result_set, reader)()) - return - df = _read_all(getattr(result_set, reader)()) + columns = [("t", "timestamp"), ("v", "varchar")] + if expected is None: + with pytest.raises(OperationalError): + _read_all(_local_result_set(columns, str(path), chunksize=chunksize, **kwargs)) + return + df = _read_all(_local_result_set(columns, str(path), chunksize=chunksize, **kwargs)) for name, values in expected.items(): assert df[name].equals(values, check_dtypes=True) if not expected: assert df.columns == ["v"] - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) - def test_read_csv_ignore_errors(self, tmp_path, reader): - """With ignore_errors, timestamp text that is not a timestamp becomes null. + @pytest.mark.parametrize("chunksize", [None, 10_000]) + @pytest.mark.parametrize("ignore_errors", [True, False]) + def test_read_csv_invalid_timestamps(self, tmp_path, chunksize, ignore_errors): + """Timestamp text that is not a timestamp fails the read, or with ignore_errors is null. Timestamps with more than nine fractional digits, which Polars itself reads as null with ignore_errors, are still truncated. """ path = tmp_path / "result.csv" path.write_bytes(b'"t"\n"2020-01-02 03:04:05.123456789012"\n"invalid"\n') - result_set = _chunked_result_set() - result_set._query_execution = None - result_set._converter = DefaultPolarsTypeConverter() - result_set._kwargs = {"ignore_errors": True} - result_set._metadata = ( - {"Name": "t", "Type": "timestamp", "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"}, + if not ignore_errors: + with pytest.raises(OperationalError): + _read_all(_local_result_set([("t", "timestamp")], str(path), chunksize=chunksize)) + return + result_set = _local_result_set( + [("t", "timestamp")], str(path), chunksize=chunksize, ignore_errors=True ) - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - ): - df = _read_all(getattr(result_set, reader)()) - assert df["t"].equals( + assert _read_all(result_set)["t"].equals( pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), check_dtypes=True, ) @@ -623,10 +503,7 @@ def test_itertuples_value_python_does_not_hold(self): assert list(rows) == [] def test_itertuples_nested_value_python_does_not_hold(self): - """The rows before a nested value that Python does not hold are returned too. - - Polars raises the failure within a struct as a PanicException. - """ + """The rows before a nested value that Python does not hold are returned too.""" dates = pl.Series([0, 0, -719528], dtype=pl.Int32).cast(pl.Date) df = pl.DataFrame({"s": pl.DataFrame({"d": dates}).to_struct(), "n": [1, 2, 3]}) rows = PolarsDataFrameIterator(df, {}, ["s", "n"]).itertuples() @@ -635,7 +512,8 @@ def test_itertuples_nested_value_python_does_not_hold(self): ({"d": date(1970, 1, 1)}, 1), ({"d": date(1970, 1, 1)}, 2), ] - with pytest.raises(pl.exceptions.PanicException): + # Polars raises the failure within a struct as a BaseException of its own. + with pytest.raises(BaseException): # noqa: B017, PT011 next(rows) assert list(rows) == [] @@ -648,13 +526,6 @@ def test_itertuples_large_chunk(self): assert rows == [(i, f"v{i}") for i in range(size)] - def test_itertuples_without_column_names(self): - """Without column names, each row of a chunk is empty.""" - df = pl.DataFrame({"n": [1, 2]}) - - assert list(PolarsDataFrameIterator(df, {}, []).itertuples()) == [(), ()] - assert list(PolarsDataFrameIterator(df, {}, []).iterrows()) == [(0, {}), (1, {})] - def test_iterrows_continuous_row_numbers(self): """Row numbers continue across chunks, with each row keyed by the column names.""" chunks = [pl.DataFrame({"n": [1, 2]}), pl.DataFrame({"n": [3]})] From 9529eb2464d0f33681ac102243cc7eb38df787b7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:03:43 +0900 Subject: [PATCH 2/8] Remove and rewrite redundant Arrow cursor tests Remove the as_polars() tests that check pl.from_arrow() on the tables that the as_arrow tests assert, and the timeout tests of private attributes; test_as_polars keeps one CSV query. The duplicate TIMESTAMP column of the offline result-set test joins the live test_duplicate_column_names, and the fetch-across-batches checks join test_fetch_across_empty_batches. test_fetch_after_close closes a result with rows left, so that it fails when close() keeps them. The Parquet chunk reader tests check that more files were opened than read and that reads continue after the first chunk, instead of exact lookahead counts and byte ranges. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 238 ++---------------------- tests/pyathena/arrow/test_reader.py | 109 ++++------- tests/pyathena/arrow/test_result_set.py | 73 ++------ 3 files changed, 61 insertions(+), 359 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index f50e1e66b..c30c679cf 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -15,7 +15,6 @@ from unittest.mock import MagicMock, patch import pandas as pd -import polars as pl import pyarrow as pa import pytest @@ -566,187 +565,10 @@ def test_complex_unload_as_arrow(self, arrow_cursor): ) ] - @pytest.mark.parametrize( - "arrow_cursor", - [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], - indirect=["arrow_cursor"], - ) def test_as_polars(self, arrow_cursor): df = arrow_cursor.execute("SELECT * FROM one_row").as_polars() - assert df.height == 1 - assert df.width == 1 assert df.to_dicts() == [{"number_of_rows": 1}] - @pytest.mark.parametrize( - "arrow_cursor", - [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], - indirect=["arrow_cursor"], - ) - def test_many_as_polars(self, arrow_cursor): - df = arrow_cursor.execute("SELECT * FROM many_rows").as_polars() - assert df.height == 10000 - assert df.width == 1 - assert df.to_dicts() == [{"a": i} for i in range(10000)] - - def test_complex_as_polars(self, arrow_cursor): - df = arrow_cursor.execute( - """ - SELECT - col_boolean - ,col_tinyint - ,col_smallint - ,col_int - ,col_bigint - ,col_float - ,col_double - ,col_string - ,col_varchar - ,col_timestamp - ,CAST(col_timestamp AS time) AS col_time - ,col_date - ,col_binary - ,col_array - ,CAST(col_array AS json) AS col_array_json - ,col_map - ,CAST(col_map AS json) AS col_map_json - ,col_struct - ,col_decimal - FROM one_row_complex - """ - ).as_polars() - assert df.height == 1 - assert df.width == 19 - dtypes = tuple(df.dtypes) - assert dtypes == ( - pl.Boolean, - pl.Int8, - pl.Int16, - pl.Int32, - pl.Int64, - pl.Float32, - pl.Float64, - pl.String, - pl.String, - pl.Datetime("us"), - pl.String, - pl.Date, - pl.String, - pl.String, - pl.String, - pl.String, - pl.String, - pl.String, - pl.String, - ) - rows = df.to_dicts() - assert rows == [ - { - "col_boolean": True, - "col_tinyint": 127, - "col_smallint": 32767, - "col_int": 2147483647, - "col_bigint": 9223372036854775807, - "col_float": 0.5, - "col_double": 0.25, - "col_string": "a string", - "col_varchar": "varchar", - "col_timestamp": datetime(2017, 1, 1, 0, 0, 0), - "col_time": "00:00:00.000", - "col_date": datetime(2017, 1, 2).date(), - "col_binary": "31 32 33", - "col_array": "[1, 2]", - "col_array_json": "[1,2]", - "col_map": "{1=2, 3=4}", - "col_map_json": '{"1":2,"3":4}', - "col_struct": "{a=1, b=2}", - "col_decimal": "0.1", - } - ] - - @pytest.mark.parametrize( - "arrow_cursor", - [ - { - "cursor_kwargs": {"unload": True}, - }, - ], - indirect=["arrow_cursor"], - ) - def test_complex_unload_as_polars(self, arrow_cursor): - # NOT_SUPPORTED: Unsupported Hive type: time - # NOT_SUPPORTED: Unsupported Hive type: json - df = arrow_cursor.execute( - """ - SELECT - col_boolean - ,col_tinyint - ,col_smallint - ,col_int - ,col_bigint - ,col_float - ,col_double - ,col_string - ,col_varchar - ,col_timestamp - ,col_date - ,col_binary - ,col_array - ,col_map - ,col_struct - ,col_decimal - FROM one_row_complex - """ - ).as_polars() - assert df.height == 1 - assert df.width == 16 - # Polars 2.0 converts an Arrow map to Map; Polars 1.x to a list of key-value structs. - if hasattr(pl, "Map"): - map_dtype = pl.Map(pl.Int32, pl.Int32) - map_value = {1: 2, 3: 4} - else: - map_dtype = pl.List(pl.Struct([pl.Field("key", pl.Int32), pl.Field("value", pl.Int32)])) - map_value = [{"key": 1, "value": 2}, {"key": 3, "value": 4}] - dtypes = tuple(df.dtypes) - assert dtypes == ( - pl.Boolean, - pl.Int8, - pl.Int16, - pl.Int32, - pl.Int64, - pl.Float32, - pl.Float64, - pl.String, - pl.String, - pl.Datetime("us"), - pl.Date, - pl.Binary, - pl.List(pl.Int32), - map_dtype, - pl.Struct([pl.Field("a", pl.Int32), pl.Field("b", pl.Int32)]), - pl.Decimal(precision=10, scale=1), - ) - rows = df.to_dicts() - assert rows == [ - { - "col_boolean": True, - "col_tinyint": 127, - "col_smallint": 32767, - "col_int": 2147483647, - "col_bigint": 9223372036854775807, - "col_float": 0.5, - "col_double": 0.25, - "col_string": "a string", - "col_varchar": "varchar", - "col_timestamp": datetime(2017, 1, 1, 0, 0, 0), - "col_date": datetime(2017, 1, 2).date(), - "col_binary": b"123", - "col_array": [1, 2], - "col_map": map_value, - "col_struct": {"a": 1, "b": 2}, - "col_decimal": Decimal("0.1"), - } - ] - def test_cancel(self, arrow_cursor): def cancel(c): time.sleep(random.randint(5, 10)) @@ -980,49 +802,6 @@ def test_callback(query_id: str): assert callback_results[0] == arrow_cursor.query_id assert arrow_cursor.query_id is not None - @pytest.mark.parametrize( - "arrow_cursor", - [ - { - "cursor_kwargs": { - "connect_timeout": 10, - "request_timeout": 30, - } - } - ], - indirect=["arrow_cursor"], - ) - def test_timeout_parameters(self, arrow_cursor): - """Test that timeout parameters are correctly passed to ArrowCursor and result set.""" - # Verify timeout parameters are set on cursor - assert arrow_cursor._connect_timeout == 10 - assert arrow_cursor._request_timeout == 30 - - # Execute a simple query to create a result set - arrow_cursor.execute("SELECT 1") - - # Verify timeout parameters are passed to result set - assert arrow_cursor.result_set._connect_timeout == 10 - assert arrow_cursor.result_set._request_timeout == 30 - - @pytest.mark.parametrize( - "arrow_cursor", - [{"cursor_kwargs": {"connect_timeout": 5.5, "request_timeout": 15.5}}], - indirect=["arrow_cursor"], - ) - def test_timeout_parameters_float(self, arrow_cursor): - """Test that timeout parameters accept float values.""" - # Verify float timeout parameters are set on cursor - assert arrow_cursor._connect_timeout == 5.5 - assert arrow_cursor._request_timeout == 15.5 - - # Execute a simple query to create a result set - arrow_cursor.execute("SELECT 1") - - # Verify float timeout parameters are passed to result set - assert arrow_cursor.result_set._connect_timeout == 5.5 - assert arrow_cursor.result_set._request_timeout == 15.5 - @pytest.mark.parametrize( "arrow_cursor", [ @@ -1189,10 +968,12 @@ def test_managed_results_match_result_file(self): indirect=["arrow_cursor"], ) def test_duplicate_column_names(self, arrow_cursor): + """Columns with the same name keep their own types; timestamps are truncated.""" arrow_cursor.execute( "SELECT 1 AS x, 2 AS x, CAST('01:02:03' AS TIME) AS x, 'a' AS y, " "json_parse('[1]') AS j, 'b' AS j, CAST('12:34:56' AS TIME) AS t, " - "CAST('01:02:03' AS TIME) AS t" + "CAST('01:02:03' AS TIME) AS t, " + "CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP(9)) AS x" ) assert arrow_cursor.fetchall() == [ ( @@ -1204,9 +985,20 @@ def test_duplicate_column_names(self, arrow_cursor): "b", datetime(2017, 1, 1, 12, 34, 56).time(), datetime(2017, 1, 1, 1, 2, 3).time(), + datetime(2020, 1, 2, 3, 4, 5, 123456), ) ] - assert arrow_cursor.as_arrow().column_names == ["x", "x", "x", "y", "j", "j", "t", "t"] + assert arrow_cursor.as_arrow().column_names == [ + "x", + "x", + "x", + "y", + "j", + "j", + "t", + "t", + "x", + ] def test_duplicate_column_names_with_newline(self, arrow_cursor): """Columns with the same name keep their own types, with a newline in the name.""" diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index 7569a4d2b..4fc676e73 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -42,13 +42,11 @@ def __init__(self, block_size): super().__init__() self.block_size = block_size self.fetched = [] - self.opened = [] def size(self, path): return os.path.getsize(path) def _open(self, path, mode="rb", block_size=None, cache_type="bytes", **kwargs): - self.opened.append((block_size, cache_type)) return _RecordingFile( self, path, block_size=block_size or self.block_size, cache_type=cache_type, **kwargs ) @@ -108,115 +106,69 @@ def test_columns(self, paths, max_prefetch_size): def test_columns_not_top_level(self, paths, max_prefetch_size, columns): with ( ParquetChunkReader(LocalFileSystem(), paths, 1, columns=columns) as reader, - pytest.raises(ValueError, match="No top-level columns"), + pytest.raises(ValueError), # noqa: PT011 ): next(reader) - @pytest.mark.parametrize("corrupt", [False, True], ids=["close", "error"]) - def test_releases_files_opened_ahead(self, paths, max_prefetch_size, monkeypatch, corrupt): - """Files opened ahead are closed when reading stops early or fails.""" - if corrupt: - # The third file fails while the files after it may be opened ahead. - with open(paths[2], "rb") as f: - valid = f.read() - with open(paths[2], "wb") as f: - f.write(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:]) - paths = [*paths, paths[0], paths[0]] + def test_releases_files_opened_ahead(self, paths, monkeypatch): + """Files opened ahead are closed when reading fails.""" + monkeypatch.setattr(ParquetChunkReader, "MAX_PREFETCH_SIZE", -1) + # The third file fails while the files after it may be opened ahead. + with open(paths[2], "rb") as f: + valid = f.read() + with open(paths[2], "wb") as f: + f.write(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:]) + fs = LocalFileSystem() opened = [] - original_open = ParquetChunkReader._open + original_open = fs.open - def open_(self, path): - source = original_open(self, path) - opened.append(source) - return source + def open_(path, mode="rb", **kwargs): + f = original_open(path, mode, **kwargs) + opened.append(f) + return f - monkeypatch.setattr(ParquetChunkReader, "_open", open_) + monkeypatch.setattr(fs, "open", open_) # Open ahead at once, so that a file opened ahead is never cancelled. monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) - reader = ParquetChunkReader(LocalFileSystem(), paths, 2) - if corrupt: - with pytest.raises((OSError, ValueError)): - list(reader) - else: - next(reader) - reader.close() - # The file after the last one received was opened ahead. - assert len(opened) == (4 if corrupt else 2) - assert all(source.closed for source in opened) - - def test_open_failure_opens_no_further_file(self, paths, monkeypatch): - """A file that fails to open raises before the next file is opened.""" - opened = [] - - def open_(self, path): - opened.append(path) - raise OSError(f"cannot open {path}") - - monkeypatch.setattr(ParquetChunkReader, "_open", open_) - # Open at once, so that an open submitted too early is not cancelled. - monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) - with pytest.raises(OSError, match="cannot open"): - next(ParquetChunkReader(LocalFileSystem(), paths, 2)) - assert opened == paths[:1] + with pytest.raises((OSError, ValueError)): + list(ParquetChunkReader(fs, [*paths, paths[0], paths[0]], 2)) + # A file after the failed third one was opened ahead. + assert len(opened) > 3 + assert all(f.closed for f in opened) def test_prefetched_file(self, paths): - """A file read whole is opened without a cache, in the filesystem's block size.""" + """A file read whole is fetched in one read of exactly the file.""" fs = _RecordingFileSystem(block_size=32 * 2**10) with ParquetChunkReader(fs, paths[2:], 10) as reader: assert next(reader).num_rows == 7 - assert fs.opened == [(None, "none")] # One read of exactly the file, which PyAthena's filesystem splits into blocks. assert fs.fetched == [(0, os.path.getsize(paths[2]))] @pytest.mark.parametrize("cache_type", [None, "bytes"], ids=["default", "bytes"]) def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): - """A streamed file is read one pre-buffered row group at a time, when needed. - - With the default cache, each row group is fetched as one range of exactly its - column chunks; a given cache type adds its read-ahead to each range. - """ + """A streamed file is read one row group at a time, when needed.""" monkeypatch.setattr(ParquetChunkReader, "MAX_PREFETCH_SIZE", -1) path = tmp_path / "groups.parquet" columns = {f"c{i}": pa.array([os.urandom(256) for _ in range(300)]) for i in range(3)} pq.write_table(pa.table(columns), path, compression="none", row_group_size=100) - metadata = pq.read_metadata(path) - spans = [] - for index in range(metadata.num_row_groups): - chunks = [metadata.row_group(index).column(c) for c in range(metadata.num_columns)] - starts = [c.dictionary_page_offset or c.data_page_offset for c in chunks] - spans.append( - ( - min(starts), - max(s + c.total_compressed_size for s, c in zip(starts, chunks, strict=True)), - ) - ) - assert len(spans) == 3 fs = _RecordingFileSystem(block_size=32 * 2**10) with ParquetChunkReader(fs, [str(path)], 50, cache_type=cache_type) as reader: assert next(reader).num_rows == 50 - # The first fetch reads the footer; the first chunk needs the first row group. - first = fs.fetched[1:] + fetched = len(fs.fetched) assert sum(table.num_rows for table in reader) == 250 - ranges = fs.fetched[1:] - if cache_type is None: - assert first == spans[:1] - assert ranges == spans - else: - assert [start for start, _ in ranges] == [start for start, _ in spans] - assert all( - end > span_end - for (_, end), (_, span_end) in zip(ranges[:2], spans[:2], strict=True) - ) + # The later row groups are fetched after the first chunk. + assert len(fs.fetched) > fetched def test_no_files(self): assert list(ParquetChunkReader(LocalFileSystem(), [], 4)) == [] @pytest.mark.parametrize("chunksize", [0, -1]) def test_invalid_chunksize(self, paths, chunksize): - with pytest.raises(ValueError, match="chunksize must be positive"): + with pytest.raises(ValueError): # noqa: PT011 ParquetChunkReader(LocalFileSystem(), paths, chunksize) def test_close(self, paths, max_prefetch_size, monkeypatch): + """close() closes the file being read and the files opened ahead.""" fs = LocalFileSystem() opened = [] original_open = fs.open @@ -227,9 +179,12 @@ def open_(path, mode="rb", **kwargs): return f monkeypatch.setattr(fs, "open", open_) + # Open ahead at once, so that a file opened ahead is never cancelled. + monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) reader = ParquetChunkReader(fs, paths, 2) assert next(reader).num_rows == 2 reader.close() - assert opened + # The file after the first one was opened ahead. + assert len(opened) > 1 assert all(f.closed for f in opened) assert list(reader) == [] diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index cd214b6b4..f43c5e1e5 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -6,7 +6,6 @@ # SPDX-License-Identifier: MIT import os import sys -from datetime import datetime from unittest.mock import MagicMock, patch import pyarrow as pa @@ -21,22 +20,7 @@ class TestAthenaArrowResultSet: def test_fetch_after_close(self): - """No AWS calls; the query execution and the filesystem are mocked.""" - with patch.object(AthenaArrowResultSet, "_create_s3_file_system"): - result_set = AthenaArrowResultSet( - connection=MagicMock(), - converter=DefaultArrowTypeConverter(), - query_execution=MagicMock(state=AthenaQueryExecution.STATE_FAILED), - arraysize=1, - retry_config=RetryConfig(), - ) - result_set.close() - assert result_set.fetchone() is None - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - - def test_fetch_across_batches(self): - """The buffer is refilled one record batch at a time. + """close() discards the rows that are not fetched yet. No AWS calls; the first page request and the result table are mocked. """ @@ -57,14 +41,13 @@ def test_fetch_across_batches(self): arraysize=1, retry_config=RetryConfig(), ) - assert result_set.fetchmany(2) == [(1,), (2,)] - assert result_set.fetchone() == (3,) + assert result_set.fetchone() == (1,) + result_set.close() + assert result_set.fetchone() is None assert result_set.fetchmany() == [] assert result_set.fetchall() == [] - assert result_set.fetchone() is None - assert result_set.rownumber == 3 - @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"]) + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter", "mixed"]) @pytest.mark.parametrize( ("chunks", "expected"), [ @@ -79,8 +62,9 @@ def test_fetch_across_empty_batches(self, method, chunks, expected): A table read from result files that include an empty one has such batches; the fetches used to stop at the first one. to_batches() yields - no empty batch after the last rows. No AWS calls; the first page - request and the result table are mocked. + no empty batch after the last rows. The fetch methods can be mixed, and + return nothing after the last row. No AWS calls; the first page request + and the result table are mocked. """ table = pa.concat_tables([pa.table({"a": pa.array(chunk, pa.int64())}) for chunk in chunks]) with ( @@ -105,36 +89,17 @@ def test_fetch_across_empty_batches(self, method, chunks, expected): rows = result_set.fetchmany(10) elif method == "fetchall": rows = result_set.fetchall() - else: + elif method == "iter": rows = list(result_set) + else: + rows = [*result_set.fetchmany(2), *filter(None, [result_set.fetchone()])] + rows += result_set.fetchall() assert rows == [(value,) for value in expected] assert result_set.rownumber == len(expected) + assert result_set.fetchmany() == [] + assert result_set.fetchall() == [] assert result_set.fetchone() is None - def test_read_csv_timestamps_with_the_same_name(self): - """Timestamp columns are converted by position, also with names that repeat. - - No AWS calls; the GetQueryResults rows are mocked. - """ - result_set = AthenaArrowResultSet.__new__(AthenaArrowResultSet) # bypass __init__ - result_set._query_execution = None - result_set._converter = DefaultArrowTypeConverter() - result_set._block_size = AthenaArrowResultSet.DEFAULT_BLOCK_SIZE - result_set._metadata = tuple( - {"Name": n, "Type": t, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} - for n, t in (("x", "varchar"), ("x", "timestamp"), ("y", "timestamp")) - ) - data = b'"x","x","y"\n"a","2020-01-02 03:04:05.123456789","2020-01-02 03:04:05.1"\n"b",,\n' - with patch.object(AthenaArrowResultSet, "_fetch_all_rows_as_csv", return_value=data): - table = result_set._read_csv() - assert table.column_names == ["x", "x", "y"] - assert table.schema.types == [pa.string(), pa.timestamp("us"), pa.timestamp("us")] - rows = list(zip(*(column.to_pylist() for column in table.columns), strict=True)) - assert rows == [ - ("a", datetime(2020, 1, 2, 3, 4, 5, 123456), datetime(2020, 1, 2, 3, 4, 5, 100000)), - ("b", None, None), - ] - @pytest.mark.parametrize( ("connection_kwargs", "expected_access_key"), [ @@ -191,13 +156,3 @@ def test_create_s3_file_system_uses_session_credentials( options = fs.__reduce__()[1][0] assert options["access_key"] == expected_access_key assert options["role_arn"] == "" - - def test_open_csv_closes_file_when_stream_fails(self): - """The input file is closed if the read-ahead stream cannot be created.""" - result_set = AthenaArrowResultSet.__new__(AthenaArrowResultSet) # bypass __init__ - result_set._fs = MagicMock() - file = result_set._fs.open_input_file.return_value - file.size.side_effect = OSError("HEAD failed") - with pytest.raises(OSError, match="HEAD failed"): - result_set._open_csv("bucket/key.csv") - file.close.assert_called_once_with() From d894953386548d1f5a45504929bcc20c8cfa5fe2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:03:45 +0900 Subject: [PATCH 3/8] Remove and rewrite redundant S3FS cursor tests Remove the single-cell NULL and empty-string tests, the basic reader queries, and the async copies of sync queries that read through the same result set; test_null_vs_empty_string checks both readers in one query. The tab, newline, and comma values are columns of one SELECT per reader, and test_executemany runs on the empty_table fixture instead of being skipped. The reader tests of Python's csv module and of the class declaration are removed. test_context_manager leaves the context before the rows run out, so that it fails when close() does not close the file. The CSV open options are checked through a result set built by its constructor. test_on_start_query_execution stays until #1243 makes the connection callback test run the callback for every cursor class. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/s3fs/test_async_cursor.py | 94 --------- tests/pyathena/s3fs/test_cursor.py | 230 +++-------------------- tests/pyathena/s3fs/test_reader.py | 95 +--------- tests/pyathena/s3fs/test_result_set.py | 78 ++++---- 4 files changed, 71 insertions(+), 426 deletions(-) diff --git a/tests/pyathena/s3fs/test_async_cursor.py b/tests/pyathena/s3fs/test_async_cursor.py index 06d082683..e77a46690 100644 --- a/tests/pyathena/s3fs/test_async_cursor.py +++ b/tests/pyathena/s3fs/test_async_cursor.py @@ -8,8 +8,6 @@ import contextlib import random import time -from datetime import datetime -from decimal import Decimal from unittest.mock import MagicMock, patch import pytest @@ -25,20 +23,6 @@ class TestAsyncS3FSCursor: - @pytest.mark.parametrize( - ("async_s3fs_cursor", "expected_empty"), - [ - ({}, ""), - ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), - ], - indirect=["async_s3fs_cursor"], - ) - def test_null_vs_empty_string(self, async_s3fs_cursor, expected_empty): - _, future = async_s3fs_cursor.execute( - "SELECT NULL AS null_col, '' AS empty_col, 'text' AS value_col" - ) - assert future.result().fetchall() == [(None, expected_empty, "text")] - def test_fetchone(self, async_s3fs_cursor): query_id, future = async_s3fs_cursor.execute("SELECT * FROM one_row") result_set = future.result() @@ -77,78 +61,6 @@ def test_invalid_arraysize(self, async_s3fs_cursor): with pytest.raises(ProgrammingError): async_s3fs_cursor.arraysize = -1 - def test_complex(self, async_s3fs_cursor): - query_id, future = async_s3fs_cursor.execute( - """ - SELECT - col_boolean - ,col_tinyint - ,col_smallint - ,col_int - ,col_bigint - ,col_float - ,col_double - ,col_string - ,col_varchar - ,col_timestamp - ,CAST(col_timestamp AS time) AS col_time - ,col_date - ,col_binary - ,col_array - ,CAST(col_array AS json) AS col_array_json - ,col_map - ,CAST(col_map AS json) AS col_map_json - ,col_struct - ,col_decimal - FROM one_row_complex - """ - ) - result_set = future.result() - assert result_set.description == [ - ("col_boolean", "boolean", None, None, 0, 0, "UNKNOWN"), - ("col_tinyint", "tinyint", None, None, 3, 0, "UNKNOWN"), - ("col_smallint", "smallint", None, None, 5, 0, "UNKNOWN"), - ("col_int", "integer", None, None, 10, 0, "UNKNOWN"), - ("col_bigint", "bigint", None, None, 19, 0, "UNKNOWN"), - ("col_float", "float", None, None, 17, 0, "UNKNOWN"), - ("col_double", "double", None, None, 17, 0, "UNKNOWN"), - ("col_string", "varchar", None, None, 2147483647, 0, "UNKNOWN"), - ("col_varchar", "varchar", None, None, 10, 0, "UNKNOWN"), - ("col_timestamp", "timestamp", None, None, 3, 0, "UNKNOWN"), - ("col_time", "time", None, None, 3, 0, "UNKNOWN"), - ("col_date", "date", None, None, 0, 0, "UNKNOWN"), - ("col_binary", "varbinary", None, None, 1073741824, 0, "UNKNOWN"), - ("col_array", "array", None, None, 0, 0, "UNKNOWN"), - ("col_array_json", "json", None, None, 0, 0, "UNKNOWN"), - ("col_map", "map", None, None, 0, 0, "UNKNOWN"), - ("col_map_json", "json", None, None, 0, 0, "UNKNOWN"), - ("col_struct", "row", None, None, 0, 0, "UNKNOWN"), - ("col_decimal", "decimal", None, None, 10, 1, "UNKNOWN"), - ] - assert result_set.fetchall() == [ - ( - True, - 127, - 32767, - 2147483647, - 9223372036854775807, - 0.5, - 0.25, - "a string", - "varchar", - datetime(2017, 1, 1, 0, 0, 0), - datetime(2017, 1, 1, 0, 0, 0).time(), - datetime(2017, 1, 2).date(), - b"123", - [1, 2], - [1, 2], - {"1": "2", "3": "4"}, - {"1": 2, "3": 4}, - {"a": "1", "b": "2"}, - Decimal("0.1"), - ) - ] - def test_cancel(self, async_s3fs_cursor): query_id, future = async_s3fs_cursor.execute( """ @@ -181,12 +93,6 @@ def test_no_ops(self): cursor.close() conn.close() - def test_show_columns(self, async_s3fs_cursor): - query_id, future = async_s3fs_cursor.execute("SHOW COLUMNS IN one_row") - result_set = future.result() - assert result_set.description == [("field", "string", None, None, 0, 0, "UNKNOWN")] - assert result_set.fetchall() == [("number_of_rows ",)] - def test_empty_result(self, async_s3fs_cursor): query_id, future = async_s3fs_cursor.execute("SELECT * FROM one_row WHERE 1 = 2") result_set = future.result() diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 09abbb4c5..bd15736b2 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -205,20 +205,6 @@ def test_empty_result(self, s3fs_cursor): assert s3fs_cursor.fetchmany(10) == [] assert s3fs_cursor.fetchall() == [] - def test_query_id(self, s3fs_cursor): - assert not s3fs_cursor.query_id - s3fs_cursor.execute("SELECT * FROM one_row") - assert s3fs_cursor.query_id is not None - - def test_description_without_execute(self, s3fs_cursor): - assert s3fs_cursor.description is None - - def test_description_with_select(self, s3fs_cursor): - s3fs_cursor.execute("SELECT * FROM one_row") - assert s3fs_cursor.description == [ - ("number_of_rows", "integer", None, None, 10, 0, "UNKNOWN") - ] - def test_description_with_ctas(self, s3fs_cursor): table_name = ( f"test_description_with_ctas_{''.join(random.choices(string.ascii_lowercase, k=10))}" @@ -257,16 +243,14 @@ def test_description_with_create_table(self, s3fs_cursor): assert s3fs_cursor.fetchone() is None s3fs_cursor.execute(f"DROP TABLE {ENV.schema}.{table_name}") - @pytest.mark.skip(reason="Requires insert_test table to exist in test environment") - def test_executemany(self, s3fs_cursor): - rows = [(1, "foo"), (2, "bar"), (3, "jim")] + def test_executemany(self, s3fs_cursor, empty_table): + rows = [(1, "foo"), (2, "bar"), (3, "jim o'rourke")] s3fs_cursor.executemany( - f"INSERT INTO {ENV.schema}.insert_test (a, b) VALUES (%(a)s, %(b)s)", - [{"a": row[0], "b": row[1]} for row in rows], + f"INSERT INTO {empty_table} (a, b) VALUES (%(a)d, %(b)s)", + [{"a": a, "b": b} for a, b in rows], ) - s3fs_cursor.execute(f"SELECT * FROM {ENV.schema}.insert_test ORDER BY a") - assert s3fs_cursor.fetchall() == rows - s3fs_cursor.execute(f"DELETE FROM {ENV.schema}.insert_test WHERE a IN (1, 2, 3)") + s3fs_cursor.execute(f"SELECT * FROM {empty_table}") + assert sorted(s3fs_cursor.fetchall()) == list(rows) def test_on_start_query_execution(self): callback_query_id = None @@ -307,154 +291,26 @@ def callback(query_id): cursor.execute("SELECT * FROM one_row", on_start_query_execution=callback) assert callback_query_id == cursor.query_id + @pytest.mark.parametrize( + "s3fs_cursor", + [ + {"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, + {"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, + ], + indirect=["s3fs_cursor"], + ) def test_contain_tab_character(self, s3fs_cursor): - """Test that tab characters in result data are handled correctly. - - S3FS cursor uses tab as delimiter for parsing Athena's CSV output. - This test verifies that data containing tab characters is correctly - parsed when Athena properly quotes such fields. - """ - # Test with tab character in string using CHR(9) - s3fs_cursor.execute("SELECT 'before' || CHR(9) || 'after' AS col_with_tab") - result = s3fs_cursor.fetchone() - assert result == ("before\tafter",) - - # Test with multiple columns where one contains tab + """Both readers read values with tabs, newlines, and commas, which Athena quotes.""" s3fs_cursor.execute( """ SELECT - 'normal' AS col1, - 'has' || CHR(9) || 'tab' AS col2, - 'also_normal' AS col3 + 'before' || CHR(9) || 'after' AS col_tab + ,'line1' || CHR(10) || 'line2' AS col_newline + ,'a' || CHR(9) || 'b' || CHR(10) || 'c' AS col_tab_newline + ,'a,b,c' AS col_comma """ ) - result = s3fs_cursor.fetchone() - assert result == ("normal", "has\ttab", "also_normal") - - # Test with newline character as well - s3fs_cursor.execute("SELECT 'line1' || CHR(10) || 'line2' AS col_with_newline") - result = s3fs_cursor.fetchone() - assert result == ("line1\nline2",) - - # Test with both tab and newline - s3fs_cursor.execute( - "SELECT 'a' || CHR(9) || 'b' || CHR(10) || 'c' AS col_with_special_chars" - ) - result = s3fs_cursor.fetchone() - assert result == ("a\tb\nc",) - - @pytest.mark.parametrize( - "csv_reader_class", - [EmptyStringAsNullCSVReader, AthenaCSVReader], - ) - def test_basic_query_with_reader(self, csv_reader_class): - """Both readers should work for basic queries.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": csv_reader_class}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT * FROM one_row") - result = cursor.fetchall() - assert result == [(1,)] - - @pytest.mark.parametrize( - "csv_reader_class", - [EmptyStringAsNullCSVReader, AthenaCSVReader], - ) - def test_multiple_columns_with_reader(self, csv_reader_class): - """Test multiple columns work with both readers.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": csv_reader_class}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute( - """ - SELECT - 1 AS col_int, - 'text' AS col_string, - 1.5 AS col_double - """ - ) - result = cursor.fetchone() - assert result == (1, "text", 1.5) - - def test_null_with_empty_string_as_null_reader(self): - """EmptyStringAsNullCSVReader: NULL is returned as None.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT NULL AS null_col") - result = cursor.fetchone() - assert result == (None,) - - def test_null_with_athena_reader(self): - """AthenaCSVReader: NULL is returned as None.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": AthenaCSVReader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT NULL AS null_col") - result = cursor.fetchone() - assert result == (None,) - - def test_empty_string_with_empty_string_as_null_reader(self): - """EmptyStringAsNullCSVReader: Empty string becomes None (loses distinction).""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT '' AS empty_col") - result = cursor.fetchone() - # EmptyStringAsNullCSVReader treats empty string same as NULL - assert result == (None,) - - def test_empty_string_with_athena_reader(self): - """AthenaCSVReader: Empty string is preserved as empty string.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": AthenaCSVReader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT '' AS empty_col") - result = cursor.fetchone() - # AthenaCSVReader preserves empty string as '' - assert result == ("",) + assert s3fs_cursor.fetchall() == [("before\tafter", "line1\nline2", "a\tb\nc", "a,b,c")] @pytest.mark.parametrize( ("csv_reader", "expected_empty"), @@ -490,50 +346,6 @@ def test_null_vs_empty_string(self, csv_reader, expected_empty): # Empty string behavior depends on reader assert result[1] == expected_empty - def test_mixed_values_with_athena_reader(self): - """AthenaCSVReader: Mixed NULL, empty string, and regular values.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": AthenaCSVReader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute( - """ - SELECT - NULL AS null_col, - '' AS empty_col, - 'text' AS text_col, - NULL AS null_col2 - """ - ) - result = cursor.fetchone() - assert result == (None, "", "text", None) - - @pytest.mark.parametrize( - "csv_reader_class", - [EmptyStringAsNullCSVReader, AthenaCSVReader], - ) - def test_quoted_string_with_comma(self, csv_reader_class): - """Both readers should handle strings containing commas.""" - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": csv_reader_class}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT 'a,b,c' AS col_with_commas") - result = cursor.fetchone() - assert result == ("a,b,c",) - @pytest.mark.parametrize( "s3fs_cursor", [ @@ -556,7 +368,7 @@ def test_fetch_all_rows(self, s3fs_cursor): def test_ambiguous_time_zone_value(self, s3fs_cursor): """A local time that its zone repeats raises DataError instead of choosing one.""" s3fs_cursor.execute(AMBIGUOUS_TIME_ZONE_QUERY) - with pytest.raises(DataError, match="is ambiguous"): + with pytest.raises(DataError): s3fs_cursor.fetchall() @pytest.mark.parametrize( diff --git a/tests/pyathena/s3fs/test_reader.py b/tests/pyathena/s3fs/test_reader.py index e6d193fde..ed43ab286 100644 --- a/tests/pyathena/s3fs/test_reader.py +++ b/tests/pyathena/s3fs/test_reader.py @@ -5,7 +5,6 @@ # # SPDX-License-Identifier: MIT -from collections.abc import Iterator from io import StringIO from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader @@ -14,18 +13,6 @@ class TestEmptyStringAsNullCSVReader: """Tests for EmptyStringAsNullCSVReader using Python's standard csv module.""" - def test_basic_parsing(self): - data = StringIO("a,b,c\n1,2,3\n") - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [["a", "b", "c"], ["1", "2", "3"]] - - def test_tab_delimiter(self): - data = StringIO("a\tb\tc\n1\t2\t3\n") - reader = EmptyStringAsNullCSVReader(data, delimiter="\t") - rows = list(reader) - assert rows == [["a", "b", "c"], ["1", "2", "3"]] - def test_empty_field_returns_none(self): """EmptyStringAsNullCSVReader returns None for both NULL and empty string.""" # Both ,, (NULL) and ,"", (empty string) become None @@ -34,30 +21,6 @@ def test_empty_field_returns_none(self): rows = list(reader) assert rows == [["a", None, "b", None, "c"]] - def test_quoted_field_with_comma(self): - data = StringIO('"a,b",c\n') - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [["a,b", "c"]] - - def test_quoted_field_with_newline(self): - data = StringIO('"line1\nline2",b\n') - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [["line1\nline2", "b"]] - - def test_escaped_quote(self): - data = StringIO('"a""b",c\n') - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [['a"b', "c"]] - - def test_empty_file(self): - data = StringIO("") - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [] - def test_empty_line(self): """Empty line returns single None (e.g., SELECT NULL).""" # Python's csv.reader returns [] for empty lines; we normalize to [None] @@ -66,28 +29,12 @@ def test_empty_line(self): rows = list(reader) assert rows == [[None]] - def test_implements_iterator_protocol(self): - """EmptyStringAsNullCSVReader implements collections.abc.Iterator.""" - data = StringIO("a,b\n") - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - assert isinstance(reader, Iterator) - - def test_close(self): - """close() releases file resources.""" - data = StringIO("a,b\n") - reader = EmptyStringAsNullCSVReader(data, delimiter=",") - reader.close() - # After close, iteration should stop immediately - rows = list(reader) - assert rows == [] - def test_context_manager(self): - """Reader can be used as context manager.""" + """Leaving the context closes the file, and iteration stops.""" data = StringIO("a,b\n1,2\n") with EmptyStringAsNullCSVReader(data, delimiter=",") as reader: - rows = list(reader) - assert rows == [["a", "b"], ["1", "2"]] - # After exiting context, reader should be closed + assert next(reader) == ["a", "b"] + assert data.closed assert list(reader) == [] @@ -129,13 +76,6 @@ def test_null_at_end(self): rows = list(reader) assert rows == [["a", "b", None]] - def test_no_trailing_null(self): - """No trailing NULL when line ends with value.""" - data = StringIO("a,b,c\n") - reader = AthenaCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [["a", "b", "c"]] - def test_empty_string_at_start(self): """Empty string at the start of a line.""" data = StringIO('"",b,c\n') @@ -212,13 +152,6 @@ def test_single_null_empty_line(self): rows = list(reader) assert rows == [[None]] - def test_quoted_field_with_newline(self): - """Quoted field containing a newline character.""" - data = StringIO('"line1\nline2",b\n') - reader = AthenaCSVReader(data, delimiter=",") - rows = list(reader) - assert rows == [["line1\nline2", "b"]] - def test_quoted_field_with_multiple_newlines(self): """Quoted field containing multiple newline characters.""" data = StringIO('"line1\nline2\nline3",b,c\n') @@ -240,26 +173,10 @@ def test_multiple_rows_with_multiline_field(self): rows = list(reader) assert rows == [["a", "b"], ["multi\nline", "c"], ["d", "e"]] - def test_implements_iterator_protocol(self): - """AthenaCSVReader implements collections.abc.Iterator.""" - data = StringIO("a,b\n") - reader = AthenaCSVReader(data, delimiter=",") - assert isinstance(reader, Iterator) - - def test_close(self): - """close() releases file resources.""" - data = StringIO("a,b\n") - reader = AthenaCSVReader(data, delimiter=",") - reader.close() - # After close, iteration should stop immediately - rows = list(reader) - assert rows == [] - def test_context_manager(self): - """Reader can be used as context manager.""" + """Leaving the context closes the file, and iteration stops.""" data = StringIO("a,b\n1,2\n") with AthenaCSVReader(data, delimiter=",") as reader: - rows = list(reader) - assert rows == [["a", "b"], ["1", "2"]] - # After exiting context, reader should be closed + assert next(reader) == ["a", "b"] + assert data.closed assert list(reader) == [] diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 0c72aa5b7..fc6e84d11 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -38,10 +38,21 @@ def close(self) -> None: self._reader.close() -def _result_set(data, csv_reader, suffix=".csv", column_types=None, type_hints=None, stream=None): +def _result_set( + data, + csv_reader, + suffix=".csv", + column_types=None, + type_hints=None, + stream=None, + block_size=None, + filesystem_class=None, +): stream = stream or BytesIO(data.encode()) - filesystem = MagicMock() - filesystem._open.return_value = stream + if filesystem_class is None: + filesystem = MagicMock() + filesystem._open.return_value = stream + filesystem_class = MagicMock(return_value=filesystem) connection = MagicMock() connection.s3_client.head_object.return_value = {"ContentLength": len(data.encode())} response = { @@ -70,8 +81,9 @@ def _result_set(data, csv_reader, suffix=".csv", column_types=None, type_hints=N ), arraysize=1, retry_config=RetryConfig(), + block_size=block_size, csv_reader=csv_reader, - filesystem_class=MagicMock(return_value=filesystem), + filesystem_class=filesystem_class, result_set_type_hints=type_hints, ) return result_set, stream @@ -81,13 +93,6 @@ class _UserS3FileSystem(S3FileSystem): """A user's subclass, whose own cache settings apply.""" -_PREFETCH_OPTIONS = { - "block_size": 32 * 2**20, - "cache_type": "background", - "cache_options": {"maxblocks": 2}, -} - - class _FailingCSVReader(_TupleCSVReader): """A reader that fails to read the header row.""" @@ -97,21 +102,37 @@ def __next__(self) -> tuple[str | None, ...]: class TestAthenaS3FSResultSet: @pytest.mark.parametrize( - ("block_size", "filesystem_class", "expected"), + ("block_size", "filesystem_class", "prefetch"), [ - (None, S3FileSystem, _PREFETCH_OPTIONS), - (None, AioS3FileSystem, _PREFETCH_OPTIONS), - (5 * 2**20, S3FileSystem, {}), - (None, _UserS3FileSystem, {}), + (None, S3FileSystem, True), + (None, AioS3FileSystem, True), + (5 * 2**20, S3FileSystem, False), + (None, _UserS3FileSystem, False), ], ids=["default", "aio", "block_size", "subclass"], ) - def test_get_csv_open_options(self, block_size, filesystem_class, expected): - """PyAthena's filesystems prefetch the CSV result unless block_size is given.""" - result_set = AthenaS3FSResultSet.__new__(AthenaS3FSResultSet) - result_set._block_size = block_size - result_set._fs = filesystem_class.__new__(filesystem_class) - assert result_set._get_csv_open_options() == expected + def test_csv_open_options(self, block_size, filesystem_class, prefetch): + """PyAthena's filesystems prefetch the CSV result unless block_size is given. + + With block_size or a user's filesystem class, the file is opened with the + filesystem's own settings. + """ + data = "null_col,empty_col,value_col\n,,a\n" + with patch.object( + filesystem_class, "_open", autospec=True, return_value=BytesIO(data.encode()) + ) as open_file: + result_set, _ = _result_set( + data, None, block_size=block_size, filesystem_class=filesystem_class + ) + with result_set: + assert result_set.fetchall() == [(None, None, "a")] + options = open_file.call_args.kwargs + if prefetch: + assert options["cache_type"] == "background" + assert options["block_size"] == AthenaS3FSResultSet.CSV_BLOCK_SIZE + else: + assert "cache_type" not in options + assert "block_size" not in options def test_reader_closes_at_end_of_results(self): """The file is closed once every row has been read, before close().""" @@ -121,22 +142,11 @@ def test_reader_closes_at_end_of_results(self): assert result_set.fetchone() is None result_set.close() - def test_fetch_across_batches(self): - """The buffer is refilled one batch of arraysize rows at a time.""" - result_set, _ = _result_set("null_col,empty_col,value_col\n,,a\n,,b\n,,c\n", None) - assert result_set.fetchmany(2) == [(None, None, "a"), (None, None, "b")] - assert result_set.fetchone() == (None, None, "c") - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - assert result_set.fetchone() is None - assert result_set.rownumber == 3 - result_set.close() - def test_reader_failure_closes_file(self): """A failure to read the header row closes the file.""" data = "null_col,empty_col,value_col\n" stream = BytesIO(data.encode()) - with pytest.raises(OperationalError, match="unreadable header"): + with pytest.raises(OperationalError): _result_set(data, _FailingCSVReader, stream=stream) assert stream.closed From 5a8b81002b1ac0de856b1e6c6ed05545f5f39aa5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:14:03 +0900 Subject: [PATCH 4/8] Check that storage options reach chunked CSV reads and that fetches stop The offline storage-options test also reads a chunked CSV result, and checks that the given options reach every Polars read unmerged. The nested-value test fails if the rows end instead of raising, the mixed fetch case checks the row fetchone() returns after fetchmany(), and the duplicate-name query has a NULL timestamp. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 7 +++-- tests/pyathena/arrow/test_result_set.py | 6 +++- tests/pyathena/polars/test_result_set.py | 38 ++++++++++++++++-------- 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index c30c679cf..bcda053ef 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -968,12 +968,13 @@ def test_managed_results_match_result_file(self): indirect=["arrow_cursor"], ) def test_duplicate_column_names(self, arrow_cursor): - """Columns with the same name keep their own types; timestamps are truncated.""" + """Columns with the same name keep their own types, as do their timestamps and NULLs.""" arrow_cursor.execute( "SELECT 1 AS x, 2 AS x, CAST('01:02:03' AS TIME) AS x, 'a' AS y, " "json_parse('[1]') AS j, 'b' AS j, CAST('12:34:56' AS TIME) AS t, " "CAST('01:02:03' AS TIME) AS t, " - "CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP(9)) AS x" + "CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP(9)) AS x, " + "CAST(NULL AS TIMESTAMP) AS x" ) assert arrow_cursor.fetchall() == [ ( @@ -986,6 +987,7 @@ def test_duplicate_column_names(self, arrow_cursor): datetime(2017, 1, 1, 12, 34, 56).time(), datetime(2017, 1, 1, 1, 2, 3).time(), datetime(2020, 1, 2, 3, 4, 5, 123456), + None, ) ] assert arrow_cursor.as_arrow().column_names == [ @@ -998,6 +1000,7 @@ def test_duplicate_column_names(self, arrow_cursor): "t", "t", "x", + "x", ] def test_duplicate_column_names_with_newline(self, arrow_cursor): diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index f43c5e1e5..99f47452a 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -92,7 +92,11 @@ def test_fetch_across_empty_batches(self, method, chunks, expected): elif method == "iter": rows = list(result_set) else: - rows = [*result_set.fetchmany(2), *filter(None, [result_set.fetchone()])] + rows = result_set.fetchmany(2) + row = result_set.fetchone() + assert row == ((expected[2],) if len(expected) > 2 else None) + if row is not None: + rows.append(row) rows += result_set.fetchall() assert rows == [(value,) for value in expected] assert result_set.rownumber == len(expected) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 12723f9e8..7b8d99cad 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -300,28 +300,39 @@ def test_parquet_chunks_raise_when_read_fails_partway(self, tmp_path, read_optio if not read_options: assert sum(chunk.height for chunk in chunks) == _ROWS_BEFORE_FAILURE - @pytest.mark.parametrize("chunksize", [None, 1]) - def test_unload_storage_options_replace_defaults(self, tmp_path, chunksize): - """storage_options given to execute() reach every Polars read of an UNLOAD result. + @pytest.mark.parametrize( + ("unload", "chunksize"), + [(False, 1), (True, None), (True, 1)], + ids=["csv-chunks", "unload", "unload-chunks"], + ) + def test_storage_options_replace_defaults(self, tmp_path, unload, chunksize): + """storage_options given to execute() reach every Polars read as given. - Polars ignores them for local files, so the reads are observed. + These reads use Polars' native object store, which ignores the options for + local files, so the reads are observed. """ - pl.DataFrame({"a": [1, 2]}).write_parquet(tmp_path / "0.parquet") storage_options = {"aws_region": "us-west-2"} + if unload: + pl.DataFrame({"a": [1, 2]}).write_parquet(tmp_path / "0.parquet") + files = {"column_types": [], "data_manifest": [str(tmp_path / "0.parquet")]} + else: + (tmp_path / "result.csv").write_text("a\n1\n2\n") + files = { + "column_types": [("a", "integer")], + "output_location": str(tmp_path / "result.csv"), + } with ( patch("polars.read_parquet", wraps=pl.read_parquet) as read_parquet, patch("polars.scan_parquet", wraps=pl.scan_parquet) as scan_parquet, + patch("polars.scan_csv", wraps=pl.scan_csv) as scan_csv, ): result_set = _local_result_set( - [], - data_manifest=[str(tmp_path / "0.parquet")], - chunksize=chunksize, - storage_options=storage_options, + **files, chunksize=chunksize, storage_options=storage_options ) assert result_set.fetchall() == [(1,), (2,)] - calls = read_parquet.call_args_list + scan_parquet.call_args_list - # The schema scan, and the whole read or the scan of the chunks. - assert len(calls) == 2 + calls = read_parquet.call_args_list + scan_parquet.call_args_list + scan_csv.call_args_list + # An UNLOAD result's schema is scanned before its files are read. + assert len(calls) == (2 if unload else 1) assert all(call.kwargs["storage_options"] is storage_options for call in calls) @pytest.mark.parametrize("chunksize", [None, 10_000]) @@ -513,8 +524,9 @@ def test_itertuples_nested_value_python_does_not_hold(self): ({"d": date(1970, 1, 1)}, 2), ] # Polars raises the failure within a struct as a BaseException of its own. - with pytest.raises(BaseException): # noqa: B017, PT011 + with pytest.raises(BaseException) as error: # noqa: PT011 next(rows) + assert not isinstance(error.value, StopIteration) assert list(rows) == [] def test_itertuples_large_chunk(self): From 6b8f9ea66f66a752bba7cb832ab6e9b5a357a6d4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:18:20 +0900 Subject: [PATCH 5/8] Read the local Polars results through GetQueryResults and as_polars() The local result sets get their columns from a mocked GetQueryResults response instead of a patched _pre_fetch(), and the tests read them with as_polars(). The live storage-options test keeps the whole CSV read through fsspec, the only read that the offline test does not observe. The S3FS special values are columns of the NULL and empty-string query, and the Arrow reader tests share the fixture that records opened files. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_reader.py | 51 +++++++++---------- tests/pyathena/polars/test_cursor.py | 38 ++++---------- tests/pyathena/polars/test_result_set.py | 63 ++++++++++++----------- tests/pyathena/s3fs/test_cursor.py | 64 ++++++++---------------- tests/pyathena/s3fs/test_result_set.py | 2 +- 5 files changed, 88 insertions(+), 130 deletions(-) diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index 4fc676e73..7fb57644d 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -62,6 +62,25 @@ def max_prefetch_size(self, request, monkeypatch): monkeypatch.setattr(ParquetChunkReader, "MAX_PREFETCH_SIZE", -1) return request.param + @pytest.fixture + def opened_files(self, monkeypatch): + """A local filesystem that records the files it opens, which are opened ahead at once. + + A file opened ahead at once is never cancelled. + """ + fs = LocalFileSystem() + opened = [] + original_open = fs.open + + def open_(path, mode="rb", **kwargs): + f = original_open(path, mode, **kwargs) + opened.append(f) + return f + + monkeypatch.setattr(fs, "open", open_) + monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) + return fs, opened + @pytest.fixture def paths(self, tmp_path): """Three files of 3, 0, and 7 rows; the last has row groups of 2 rows.""" @@ -110,26 +129,15 @@ def test_columns_not_top_level(self, paths, max_prefetch_size, columns): ): next(reader) - def test_releases_files_opened_ahead(self, paths, monkeypatch): + @pytest.mark.parametrize("max_prefetch_size", ["stream"], indirect=True) + def test_releases_files_opened_ahead(self, paths, max_prefetch_size, opened_files): """Files opened ahead are closed when reading fails.""" - monkeypatch.setattr(ParquetChunkReader, "MAX_PREFETCH_SIZE", -1) # The third file fails while the files after it may be opened ahead. with open(paths[2], "rb") as f: valid = f.read() with open(paths[2], "wb") as f: f.write(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:]) - fs = LocalFileSystem() - opened = [] - original_open = fs.open - - def open_(path, mode="rb", **kwargs): - f = original_open(path, mode, **kwargs) - opened.append(f) - return f - - monkeypatch.setattr(fs, "open", open_) - # Open ahead at once, so that a file opened ahead is never cancelled. - monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) + fs, opened = opened_files with pytest.raises((OSError, ValueError)): list(ParquetChunkReader(fs, [*paths, paths[0], paths[0]], 2)) # A file after the failed third one was opened ahead. @@ -167,20 +175,9 @@ def test_invalid_chunksize(self, paths, chunksize): with pytest.raises(ValueError): # noqa: PT011 ParquetChunkReader(LocalFileSystem(), paths, chunksize) - def test_close(self, paths, max_prefetch_size, monkeypatch): + def test_close(self, paths, max_prefetch_size, opened_files): """close() closes the file being read and the files opened ahead.""" - fs = LocalFileSystem() - opened = [] - original_open = fs.open - - def open_(path, mode="rb", **kwargs): - f = original_open(path, mode, **kwargs) - opened.append(f) - return f - - monkeypatch.setattr(fs, "open", open_) - # Open ahead at once, so that a file opened ahead is never cancelled. - monkeypatch.setattr("pyathena.arrow.reader.ThreadPoolExecutor", ImmediateExecutor) + fs, opened = opened_files reader = ParquetChunkReader(fs, paths, 2) assert next(reader).num_rows == 2 reader.close() diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index c9b0f38a6..59709e140 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -37,11 +37,6 @@ # Polars 2.0 reads INT96 at the time unit of the read schema; Polars 1.x at nanoseconds. _POLARS_2 = int(pl.__version__.split(".")[0]) >= 2 -_INVALID_OBJECT_STORE_OPTIONS = { - "aws_access_key_id": "invalid", - "aws_secret_access_key": "invalid", - "aws_region": ENV.region_name, -} class TestPolarsCursor: @@ -147,33 +142,20 @@ def test_as_polars_with_read_kwargs(self, polars_cursor): ).as_polars() assert df.to_dicts() == [{"number_of_rows": "1"}] - @pytest.mark.parametrize( - ("polars_cursor", "chunksize", "storage_options"), - [ - # A whole CSV result file is read through fsspec with PyAthena's filesystem. - ( - {"cursor_kwargs": {"unload": False}}, - None, - {"key": "invalid", "secret": "invalid", "skip_instance_cache": True}, - ), - # The other results are read with Polars' native object store. - ({"cursor_kwargs": {"unload": False}}, 1, _INVALID_OBJECT_STORE_OPTIONS), - ({"cursor_kwargs": {"unload": True}}, None, _INVALID_OBJECT_STORE_OPTIONS), - ], - indirect=["polars_cursor"], - ids=["csv", "csv-chunks", "unload"], - ) - def test_storage_options_replace_defaults(self, polars_cursor, chunksize, storage_options): + def test_storage_options_replace_defaults(self, polars_cursor): """storage_options given to execute() replace PyAthena's S3 settings for the read. - Their invalid credentials fail the read. + A whole CSV result file is read through fsspec with PyAthena's filesystem, which + the invalid credentials of the options fail. """ with pytest.raises(OperationalError): - # A whole result is read in execute(), and chunks as they are iterated. - list( - polars_cursor.execute( - "SELECT * FROM one_row", chunksize=chunksize, storage_options=storage_options - ).iter_chunks() + polars_cursor.execute( + "SELECT * FROM one_row", + storage_options={ + "key": "invalid", + "secret": "invalid", + "skip_instance_cache": True, + }, ) @pytest.mark.parametrize( diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 7b8d99cad..394301aa6 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -42,8 +42,8 @@ def _local_result_set( ) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose result files are local files. - The given columns stand in for the GetQueryResults metadata, and the HEAD request of - the CSV result file finds it not empty. The connection's S3 client returns the data + GetQueryResults returns the given columns, and the HEAD request of the CSV result + file finds it not empty. The connection's S3 client returns the data manifest of an UNLOAD result, and PyAthena's filesystem reads the local files. Args: @@ -73,16 +73,19 @@ def _local_result_set( connection.s3_client.get_object.return_value = { "Body": BytesIO("\n".join(data_manifest or []).encode()) } - metadata = tuple( - {"Name": name, "Type": type_, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} - for name, type_ in column_types - ) - - def pre_fetch(self: AthenaPolarsResultSet) -> None: - self._metadata = metadata - + connection.client.get_query_results.return_value = { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [ + {"Name": name, "Type": type_, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} + for name, type_ in column_types + ] + }, + "Rows": [], + } + } with ( - patch.object(AthenaPolarsResultSet, "_pre_fetch", pre_fetch), + # The HEAD request parses the output location, which is a local path here. patch.object(AthenaPolarsResultSet, "_get_content_length", return_value=1), patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), ): @@ -98,18 +101,6 @@ def pre_fetch(self: AthenaPolarsResultSet) -> None: ) -def _read_all(result_set: AthenaPolarsResultSet) -> pl.DataFrame: - """Read the rows of a result set into one DataFrame. - - Args: - result_set: A result set read whole or in chunks. - - Returns: - The DataFrame of the result, or its chunks concatenated. - """ - return pl.concat(list(result_set.iter_chunks())) - - def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose chunks are ``chunks``, without AWS. @@ -350,7 +341,7 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, chunksize): has_header=False, schema_overrides={first: pl.Utf8}, ) - df = _read_all(result_set) + df = result_set.as_polars() assert df.to_dict(as_series=False) == {first: ["1", "2"], second: ["x", "y"]} @pytest.mark.parametrize("chunksize", [None, 10_000]) @@ -366,7 +357,7 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, chunksize): new_columns=["z", "w"], schema_overrides={"z": pl.Utf8}, ) - df = _read_all(result_set) + df = result_set.as_polars() assert df.to_dict(as_series=False) == {"z": ["001"], "w": ["x"]} @pytest.mark.parametrize( @@ -379,8 +370,16 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, chunksize): (10_000, {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), (None, {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), (10_000, {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), - (None, {"with_column_names": lambda names: [n.upper() for n in names]}, None), - (10_000, {"with_column_names": lambda names: [n.upper() for n in names]}, None), + ( + None, + {"with_column_names": lambda names: [n.upper() for n in names]}, + OperationalError, + ), + ( + 10_000, + {"with_column_names": lambda names: [n.upper() for n in names]}, + OperationalError, + ), ], ) def test_read_csv_truncates_timestamps(self, tmp_path, chunksize, kwargs, expected): @@ -395,11 +394,11 @@ def test_read_csv_truncates_timestamps(self, tmp_path, chunksize, kwargs, expect path = tmp_path / "result.csv" path.write_bytes(b'"t","v"\n"2020-01-02 03:04:05.123456789012","x"\n,"y"\n') columns = [("t", "timestamp"), ("v", "varchar")] - if expected is None: + if expected is OperationalError: with pytest.raises(OperationalError): - _read_all(_local_result_set(columns, str(path), chunksize=chunksize, **kwargs)) + _local_result_set(columns, str(path), chunksize=chunksize, **kwargs).as_polars() return - df = _read_all(_local_result_set(columns, str(path), chunksize=chunksize, **kwargs)) + df = _local_result_set(columns, str(path), chunksize=chunksize, **kwargs).as_polars() for name, values in expected.items(): assert df[name].equals(values, check_dtypes=True) if not expected: @@ -417,12 +416,12 @@ def test_read_csv_invalid_timestamps(self, tmp_path, chunksize, ignore_errors): path.write_bytes(b'"t"\n"2020-01-02 03:04:05.123456789012"\n"invalid"\n') if not ignore_errors: with pytest.raises(OperationalError): - _read_all(_local_result_set([("t", "timestamp")], str(path), chunksize=chunksize)) + _local_result_set([("t", "timestamp")], str(path), chunksize=chunksize).as_polars() return result_set = _local_result_set( [("t", "timestamp")], str(path), chunksize=chunksize, ignore_errors=True ) - assert _read_all(result_set)["t"].equals( + assert result_set.as_polars()["t"].equals( pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), check_dtypes=True, ) diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index bd15736b2..12953a661 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -250,7 +250,7 @@ def test_executemany(self, s3fs_cursor, empty_table): [{"a": a, "b": b} for a, b in rows], ) s3fs_cursor.execute(f"SELECT * FROM {empty_table}") - assert sorted(s3fs_cursor.fetchall()) == list(rows) + assert sorted(s3fs_cursor.fetchall()) == rows def test_on_start_query_execution(self): callback_query_id = None @@ -292,59 +292,39 @@ def callback(query_id): assert callback_query_id == cursor.query_id @pytest.mark.parametrize( - "s3fs_cursor", - [ - {"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, - {"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, - ], - indirect=["s3fs_cursor"], - ) - def test_contain_tab_character(self, s3fs_cursor): - """Both readers read values with tabs, newlines, and commas, which Athena quotes.""" - s3fs_cursor.execute( - """ - SELECT - 'before' || CHR(9) || 'after' AS col_tab - ,'line1' || CHR(10) || 'line2' AS col_newline - ,'a' || CHR(9) || 'b' || CHR(10) || 'c' AS col_tab_newline - ,'a,b,c' AS col_comma - """ - ) - assert s3fs_cursor.fetchall() == [("before\tafter", "line1\nline2", "a\tb\nc", "a,b,c")] - - @pytest.mark.parametrize( - ("csv_reader", "expected_empty"), + ("s3fs_cursor", "expected_empty"), [ # EmptyStringAsNullCSVReader: empty string becomes None - (EmptyStringAsNullCSVReader, None), - (AthenaCSVReader, ""), # AthenaCSVReader: empty string is preserved + ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), + # AthenaCSVReader: empty string is preserved + ({"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, ""), ], + indirect=["s3fs_cursor"], ) - def test_null_vs_empty_string(self, csv_reader, expected_empty): + def test_null_vs_empty_string(self, s3fs_cursor, expected_empty): """ Test NULL vs empty string handling with different CSV readers. EmptyStringAsNullCSVReader: Both NULL and empty string become None. AthenaCSVReader: NULL is None, empty string is preserved as ''. + Both readers read values with tabs, newlines, and commas, which Athena quotes. See docs/null_handling.md for details. """ - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": csv_reader}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT NULL AS null_col, '' AS empty_col") - result = cursor.fetchone() - # NULL is always None - assert result[0] is None - # Empty string behavior depends on reader - assert result[1] == expected_empty + s3fs_cursor.execute( + """ + SELECT + NULL AS null_col + ,'' AS empty_col + ,'before' || CHR(9) || 'after' AS col_tab + ,'line1' || CHR(10) || 'line2' AS col_newline + ,'a' || CHR(9) || 'b' || CHR(10) || 'c' AS col_tab_newline + ,'a,b,c' AS col_comma + """ + ) + assert s3fs_cursor.fetchall() == [ + (None, expected_empty, "before\tafter", "line1\nline2", "a\tb\nc", "a,b,c") + ] @pytest.mark.parametrize( "s3fs_cursor", diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index fc6e84d11..404152e3c 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -48,8 +48,8 @@ def _result_set( block_size=None, filesystem_class=None, ): - stream = stream or BytesIO(data.encode()) if filesystem_class is None: + stream = stream or BytesIO(data.encode()) filesystem = MagicMock() filesystem._open.return_value = stream filesystem_class = MagicMock(return_value=filesystem) From 738069eda8dac11aa4141e631454f2334111d8b7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:34:00 +0900 Subject: [PATCH 6/8] Check the read options and conversions that the removed tests protected ArrowCursor.as_polars() is checked on two rows and columns. The Polars storage-options test reads with explicit options while PyAthena's own credential lookup fails, and compares the options by value. The S3FS open options are checked against the documented 32 MiB blocks and two cached blocks, and the Parquet chunk reader test checks that a streamed file is opened with the requested cache type. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 6 ++++-- tests/pyathena/arrow/test_reader.py | 6 +++++- tests/pyathena/polars/test_result_set.py | 19 +++++++++++++------ tests/pyathena/s3fs/test_result_set.py | 6 ++++-- 4 files changed, 26 insertions(+), 11 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index bcda053ef..957815784 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -566,8 +566,10 @@ def test_complex_unload_as_arrow(self, arrow_cursor): ] def test_as_polars(self, arrow_cursor): - df = arrow_cursor.execute("SELECT * FROM one_row").as_polars() - assert df.to_dicts() == [{"number_of_rows": 1}] + df = arrow_cursor.execute( + "SELECT * FROM (VALUES (1, 'a'), (2, 'b')) AS t(k, s) ORDER BY k" + ).as_polars() + assert df.to_dicts() == [{"k": 1, "s": "a"}, {"k": 2, "s": "b"}] def test_cancel(self, arrow_cursor): def cancel(c): diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index 7fb57644d..cc4cdfd1f 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -42,11 +42,13 @@ def __init__(self, block_size): super().__init__() self.block_size = block_size self.fetched = [] + self.cache_types = [] def size(self, path): return os.path.getsize(path) def _open(self, path, mode="rb", block_size=None, cache_type="bytes", **kwargs): + self.cache_types.append(cache_type) return _RecordingFile( self, path, block_size=block_size or self.block_size, cache_type=cache_type, **kwargs ) @@ -154,7 +156,7 @@ def test_prefetched_file(self, paths): @pytest.mark.parametrize("cache_type", [None, "bytes"], ids=["default", "bytes"]) def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): - """A streamed file is read one row group at a time, when needed.""" + """A streamed file is read one row group at a time, when needed, with the cache type.""" monkeypatch.setattr(ParquetChunkReader, "MAX_PREFETCH_SIZE", -1) path = tmp_path / "groups.parquet" columns = {f"c{i}": pa.array([os.urandom(256) for _ in range(300)]) for i in range(3)} @@ -166,6 +168,8 @@ def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): assert sum(table.num_rows for table in reader) == 250 # The later row groups are fetched after the first chunk. assert len(fs.fetched) > fetched + # Without a cache type, only the ranges that pyarrow requests are read. + assert fs.cache_types == [cache_type or "none"] def test_no_files(self): assert list(ParquetChunkReader(LocalFileSystem(), [], 4)) == [] diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 394301aa6..4031f88be 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -38,6 +38,7 @@ def _local_result_set( output_location: str = "s3://bucket/result.csv", data_manifest: list[str] | None = None, chunksize: int | None = None, + credentials_error: Exception | None = None, **kwargs: Any, ) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose result files are local files. @@ -51,6 +52,8 @@ def _local_result_set( output_location: The local path of the CSV result file. data_manifest: The local paths of the Parquet files of an UNLOAD result. chunksize: The chunk size of the result set. + credentials_error: An error that the session raises when credentials are looked + up, or None for a session without credentials. **kwargs: The read arguments given to ``execute()``. Returns: @@ -70,6 +73,7 @@ def _local_result_set( ) connection = MagicMock(region_name=None) connection.session.get_credentials.return_value = None + connection.session.get_credentials.side_effect = credentials_error connection.s3_client.get_object.return_value = { "Body": BytesIO("\n".join(data_manifest or []).encode()) } @@ -299,8 +303,9 @@ def test_parquet_chunks_raise_when_read_fails_partway(self, tmp_path, read_optio def test_storage_options_replace_defaults(self, tmp_path, unload, chunksize): """storage_options given to execute() reach every Polars read as given. - These reads use Polars' native object store, which ignores the options for - local files, so the reads are observed. + They are read without PyAthena's credentials, whose lookup fails here. These + reads use Polars' native object store, which ignores the options for local + files, so the reads are observed. """ storage_options = {"aws_region": "us-west-2"} if unload: @@ -318,13 +323,15 @@ def test_storage_options_replace_defaults(self, tmp_path, unload, chunksize): patch("polars.scan_csv", wraps=pl.scan_csv) as scan_csv, ): result_set = _local_result_set( - **files, chunksize=chunksize, storage_options=storage_options + **files, + chunksize=chunksize, + credentials_error=OSError("no credentials"), + storage_options=storage_options, ) assert result_set.fetchall() == [(1,), (2,)] calls = read_parquet.call_args_list + scan_parquet.call_args_list + scan_csv.call_args_list - # An UNLOAD result's schema is scanned before its files are read. - assert len(calls) == (2 if unload else 1) - assert all(call.kwargs["storage_options"] is storage_options for call in calls) + assert calls + assert all(call.kwargs["storage_options"] == storage_options for call in calls) @pytest.mark.parametrize("chunksize", [None, 10_000]) def test_csv_read_kwargs_replace_defaults(self, tmp_path, chunksize): diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 404152e3c..19d848364 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -112,7 +112,8 @@ class TestAthenaS3FSResultSet: ids=["default", "aio", "block_size", "subclass"], ) def test_csv_open_options(self, block_size, filesystem_class, prefetch): - """PyAthena's filesystems prefetch the CSV result unless block_size is given. + """PyAthena's filesystems read the CSV result in 32 MiB blocks with the background + cache, which keeps two blocks, unless block_size is given. With block_size or a user's filesystem class, the file is opened with the filesystem's own settings. @@ -129,7 +130,8 @@ def test_csv_open_options(self, block_size, filesystem_class, prefetch): options = open_file.call_args.kwargs if prefetch: assert options["cache_type"] == "background" - assert options["block_size"] == AthenaS3FSResultSet.CSV_BLOCK_SIZE + assert options["block_size"] == 32 * 2**20 + assert options["cache_options"] == {"maxblocks": 2} else: assert "cache_type" not in options assert "block_size" not in options From 558f3686581337e89e5c98e78eabed5856199eef Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:43:29 +0900 Subject: [PATCH 7/8] Test the cleanup and routing that the follow-up review found untested The Arrow CSV result file is closed when pyarrow's S3 filesystem fails to get its size, through the result set's constructor. The Polars storage-options test fails when a read goes through PyAthena's filesystem, PolarsDataFrameIterator keeps the rows of a chunk without column names, and the chunk reader test checks the cache type of every streamed open without counting them. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_reader.py | 2 +- tests/pyathena/arrow/test_result_set.py | 41 ++++++++++++++++++++++++ tests/pyathena/polars/test_result_set.py | 21 +++++++++--- 3 files changed, 59 insertions(+), 5 deletions(-) diff --git a/tests/pyathena/arrow/test_reader.py b/tests/pyathena/arrow/test_reader.py index cc4cdfd1f..51629d25c 100644 --- a/tests/pyathena/arrow/test_reader.py +++ b/tests/pyathena/arrow/test_reader.py @@ -169,7 +169,7 @@ def test_streamed_file_row_groups(self, tmp_path, monkeypatch, cache_type): # The later row groups are fetched after the first chunk. assert len(fs.fetched) > fetched # Without a cache type, only the ranges that pyarrow requests are read. - assert fs.cache_types == [cache_type or "none"] + assert set(fs.cache_types) == {cache_type or "none"} def test_no_files(self): assert list(ParquetChunkReader(LocalFileSystem(), [], 4)) == [] diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 99f47452a..609ab6e95 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -14,6 +14,7 @@ from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet from pyathena.connection import Connection +from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig @@ -47,6 +48,46 @@ def test_fetch_after_close(self): assert result_set.fetchmany() == [] assert result_set.fetchall() == [] + def test_csv_file_closed_when_stream_fails(self): + """The CSV result file is closed when reading it cannot start, as when HEAD fails. + + No AWS calls; the Athena and S3 clients and pyarrow's S3 filesystem are mocked. + """ + connection = MagicMock() + connection.client.get_query_results.return_value = { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [ + { + "Name": "a", + "Type": "integer", + "Precision": 10, + "Scale": 0, + "Nullable": "UNKNOWN", + } + ] + }, + "Rows": [], + } + } + connection.s3_client.head_object.return_value = {"ContentLength": 4} + with patch("pyarrow.fs.S3FileSystem") as file_system: + file = file_system.return_value.open_input_file.return_value + file.size.side_effect = OSError("HEAD failed") + with pytest.raises(OperationalError): + AthenaArrowResultSet( + connection=connection, + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location="s3://bucket/query.csv", + substatement_type="SELECT", + ), + arraysize=1, + retry_config=RetryConfig(), + ) + file.close.assert_called_once_with() + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter", "mixed"]) @pytest.mark.parametrize( ("chunks", "expected"), diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 4031f88be..960c47bf5 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -39,6 +39,7 @@ def _local_result_set( data_manifest: list[str] | None = None, chunksize: int | None = None, credentials_error: Exception | None = None, + file_system: Any = None, **kwargs: Any, ) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose result files are local files. @@ -54,6 +55,7 @@ def _local_result_set( chunksize: The chunk size of the result set. credentials_error: An error that the session raises when credentials are looked up, or None for a session without credentials. + file_system: Creates PyAthena's filesystem; by default, one of local files. **kwargs: The read arguments given to ``execute()``. Returns: @@ -91,7 +93,10 @@ def _local_result_set( with ( # The HEAD request parses the output location, which is a local path here. patch.object(AthenaPolarsResultSet, "_get_content_length", return_value=1), - patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), + patch( + "pyathena.polars.result_set.S3FileSystem", + file_system or (lambda **_: LocalFileSystem()), + ), ): return AthenaPolarsResultSet( connection=connection, @@ -303,9 +308,9 @@ def test_parquet_chunks_raise_when_read_fails_partway(self, tmp_path, read_optio def test_storage_options_replace_defaults(self, tmp_path, unload, chunksize): """storage_options given to execute() reach every Polars read as given. - They are read without PyAthena's credentials, whose lookup fails here. These - reads use Polars' native object store, which ignores the options for local - files, so the reads are observed. + They are read without PyAthena's credentials, whose lookup fails here, or its + filesystem. These reads use Polars' native object store, which ignores the + options for local files, so the reads are observed. """ storage_options = {"aws_region": "us-west-2"} if unload: @@ -326,6 +331,7 @@ def test_storage_options_replace_defaults(self, tmp_path, unload, chunksize): **files, chunksize=chunksize, credentials_error=OSError("no credentials"), + file_system=MagicMock(side_effect=AssertionError("PyAthena's filesystem")), storage_options=storage_options, ) assert result_set.fetchall() == [(1,), (2,)] @@ -544,6 +550,13 @@ def test_itertuples_large_chunk(self): assert rows == [(i, f"v{i}") for i in range(size)] + def test_itertuples_without_column_names(self): + """Without column names, each row of a chunk is empty.""" + df = pl.DataFrame({"n": [1, 2]}) + + assert list(PolarsDataFrameIterator(df, {}, []).itertuples()) == [(), ()] + assert list(PolarsDataFrameIterator(df, {}, []).iterrows()) == [(0, {}), (1, {})] + def test_iterrows_continuous_row_numbers(self): """Row numbers continue across chunks, with each row keyed by the column names.""" chunks = [pl.DataFrame({"n": [1, 2]}), pl.DataFrame({"n": [3]})] From 5b64fc7814e499d6ee4632b26062d4f17394cc54 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:50:17 +0900 Subject: [PATCH 8/8] Read the S3FS test metadata through GetQueryResults The S3FS result-set helper gets its columns from a mocked GetQueryResults response instead of a patched _get_query_results(), as the Polars helper does, and the Arrow close test checks that the file was closed without counting the calls. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_result_set.py | 2 +- tests/pyathena/s3fs/test_result_set.py | 41 ++++++++++++++----------- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 609ab6e95..26b539e4c 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -86,7 +86,7 @@ def test_csv_file_closed_when_stream_fails(self): arraysize=1, retry_config=RetryConfig(), ) - file.close.assert_called_once_with() + assert file.close.called @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter", "mixed"]) @pytest.mark.parametrize( diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 19d848364..deb9da9fa 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -55,11 +55,17 @@ def _result_set( filesystem_class = MagicMock(return_value=filesystem) connection = MagicMock() connection.s3_client.head_object.return_value = {"ContentLength": len(data.encode())} - response = { + connection.client.get_query_results.return_value = { "ResultSet": { "ResultSetMetadata": { "ColumnInfo": [ - {"Name": name, "Type": col_type} + { + "Name": name, + "Type": col_type, + "Precision": 0, + "Scale": 0, + "Nullable": "UNKNOWN", + } for name, col_type in zip( ("null_col", "empty_col", "value_col"), column_types or ("varchar", "varchar", "varchar"), @@ -70,22 +76,21 @@ def _result_set( "Rows": [], } } - with patch.object(AthenaS3FSResultSet, "_get_query_results", return_value=response): - result_set = AthenaS3FSResultSet( - connection=connection, - converter=DefaultS3FSTypeConverter(), - query_execution=MagicMock( - state=AthenaQueryExecution.STATE_SUCCEEDED, - output_location=f"s3://bucket/query{suffix}", - substatement_type="SELECT", - ), - arraysize=1, - retry_config=RetryConfig(), - block_size=block_size, - csv_reader=csv_reader, - filesystem_class=filesystem_class, - result_set_type_hints=type_hints, - ) + result_set = AthenaS3FSResultSet( + connection=connection, + converter=DefaultS3FSTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location=f"s3://bucket/query{suffix}", + substatement_type="SELECT", + ), + arraysize=1, + retry_config=RetryConfig(), + block_size=block_size, + csv_reader=csv_reader, + filesystem_class=filesystem_class, + result_set_type_hints=type_hints, + ) return result_set, stream