From b22ef484b61fcd7ada4a45c1ab64ffe3520c83ba Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:57:10 +0900 Subject: [PATCH 1/7] Remove and rewrite redundant PandasCursor tests Drop the parameter cases that repeat the same read: the pyarrow engine case of UNLOAD, which "auto" already reads with pyarrow, the single-chunk chunksize of 1,000,000, and the unload and engine cases of DDL results. Remove the tests that set private chunking attributes or compare private predicates with constants, and the binary and pyarrow tests whose values other tests check. The binary stream close tests run offline through a cursor whose Athena and S3 clients and filesystem are local, and check that the files they opened are closed. The auto-optimized chunk sizes, the storage options, and the pyarrow engine fallbacks are checked through the values and chunk sizes that the cursor returns. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 914 +++++++++------------------ 1 file changed, 309 insertions(+), 605 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 1058386de..632cb819c 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -7,23 +7,21 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from decimal import Decimal -from unittest.mock import MagicMock, PropertyMock, patch +from typing import Any +from unittest.mock import MagicMock, patch import numpy as np import pandas as pd import pytest -from pandas.io.parsers import TextFileReader +from fsspec.implementations.local import LocalFileSystem from pyathena.error import DatabaseError, DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem +from pyathena.formatter import DefaultParameterFormatter from pyathena.model import AthenaQueryExecution from pyathena.pandas.converter import DefaultPandasTypeConverter from pyathena.pandas.cursor import PandasCursor -from pyathena.pandas.result_set import ( - AthenaPandasResultSet, - PandasDataFrameIterator, - _read_csv_with_pyarrow, -) +from pyathena.pandas.result_set import AthenaPandasResultSet from pyathena.util import RetryConfig from tests import ENV from tests.pyathena import temporal @@ -43,6 +41,105 @@ def _pandas_converter_without_bigint_dtype(): return converter +class _ResultFileSystem(LocalFileSystem): + """A filesystem that serves one local file for every path and keeps the files it opens.""" + + def __init__(self, path: str, **kwargs: Any) -> None: + """Initialize the filesystem. + + Args: + path: The local file that every path opens. + **kwargs: Ignored filesystem options. + """ + super().__init__() + self.path = path + self.opened: list[Any] = [] + + def _open(self, path: str, mode: str = "rb", **kwargs: Any) -> Any: + """Open the local file instead of the given path. + + Args: + path: The path to open, which is ignored. + mode: The mode to open the file in. + **kwargs: Ignored open options. + + Returns: + The opened local file. + """ + file = super()._open(self.path, mode) + self.opened.append(file) + return file + + +@contextlib.contextmanager +def _offline_cursor(tmp_path, columns, rows, **cursor_kwargs): + """Create a PandasCursor whose queries return a local CSV result file, without AWS. + + The Athena and S3 clients of the connection return the query execution, the + result metadata, and the size of the file, and PyAthena's filesystem reads the + local file. + + Args: + tmp_path: The directory of the result file. + columns: The names and Athena types of the result columns. + rows: The rows of the result file, as text. + **cursor_kwargs: Arguments for the cursor. + + Yields: + The cursor and the filesystem that reads the result file. + """ + path = tmp_path / "result.csv" + with path.open("w", newline="") as file: + writer = csv.writer(file, quoting=csv.QUOTE_ALL, lineterminator="\n") + writer.writerow([name for name, _ in columns]) + writer.writerows(rows) + connection = MagicMock() + connection.client.start_query_execution.return_value = {"QueryExecutionId": "query_id"} + connection.client.get_query_execution.return_value = { + "QueryExecution": { + "QueryExecutionId": "query_id", + "Query": "SELECT 1", + "StatementType": "DML", + "Status": {"State": AthenaQueryExecution.STATE_SUCCEEDED}, + "ResultConfiguration": {"OutputLocation": "s3://bucket/result.csv"}, + } + } + connection.client.get_query_results.return_value = { + "ResultSet": { + "Rows": [{"Data": [{"VarCharValue": name} for name, _ in columns]}], + "ResultSetMetadata": { + "ColumnInfo": [ + {"Name": name, "Type": type_, "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"} + for name, type_ in columns + ] + }, + } + } + connection.s3_client.head_object.return_value = {"ContentLength": path.stat().st_size} + fs = _ResultFileSystem(str(path), skip_instance_cache=True) + with ( + patch("pyathena.pandas.result_set.S3FileSystem", lambda **_: fs), + PandasCursor( + connection=connection, + converter=DefaultPandasTypeConverter(), + formatter=DefaultParameterFormatter(), + retry_config=RetryConfig(), + **cursor_kwargs, + ) as cursor, + ): + yield cursor, fs + + +# A value that makes a result file larger than the smallest that the PyArrow engine reads. +_PADDING = "x" * 120 + + +class _QuoteNoneDialect(csv.excel): + """A CSV dialect that reads quotes as text.""" + + quoting = csv.QUOTE_NONE + + def _pandas_value(column: temporal.Column): """The value that a DataFrame holds for a temporal column.""" if column.value is None: @@ -67,7 +164,6 @@ def test_binary_null_vs_empty(self, pandas_cursor, engine, chunksize): ) AS t(id, value, label, text_value) ORDER BY id""" pandas_cursor.execute(query, engine=engine, chunksize=chunksize) rows = pandas_cursor.fetchall() - assert pandas_cursor.result_set._csv_stream.closed assert [row[:3] for row in rows] == [ (1, None, "null"), (2, b"", "empty"), @@ -78,92 +174,83 @@ def test_binary_null_vs_empty(self, pandas_cursor, engine, chunksize): assert pd.isna(rows[1][3]) assert rows[2][3] == "NULL" - @pytest.mark.parametrize("chunksize", [None, 2]) - def test_binary_as_pandas(self, pandas_cursor, chunksize): - pandas_cursor.execute( - "SELECT * FROM (VALUES (1, CAST(NULL AS VARBINARY)), (2, X''), (3, X'00ff')) " - "AS t(id, value) ORDER BY id", - chunksize=chunksize, - storage_options={"connection": pandas_cursor.connection, "default_cache_type": "none"}, - ) - result = pandas_cursor.as_pandas() - df = pd.concat(list(result), ignore_index=True) if chunksize else result - assert df["value"].tolist() == [None, b"", b"\x00\xff"] - assert pandas_cursor.result_set._csv_stream.closed - @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize( "close_method", ["iterator", "generator", "context", "exhaust", "nrows", "get_chunk", "cursor", "execute"], ) - def test_binary_csv_stream_close(self, pandas_cursor, engine, close_method): - pandas_cursor.execute( - "SELECT X'00' AS value, rpad('x', 4096, 'x') AS padding " - "FROM UNNEST(sequence(1, 100)) AS t(id)", - engine=engine, + def test_binary_csv_stream_close(self, tmp_path, engine, close_method): + """The binary CSV stream of a chunked read is closed however the reading ends. + + Later reads get no chunks, and the fetch methods no rows. No AWS calls; the + result file is a local file larger than pandas' read buffer. + """ + with _offline_cursor( + tmp_path, + [("value", "varbinary"), ("padding", "varchar")], + [("00", "x" * 4096)] * 100, chunksize=1, - nrows=2 if close_method == "nrows" else None, - ) - stream = pandas_cursor.result_set._csv_stream - source = stream.buffer.raw._reader._file - dataframe_iterator = pandas_cursor.as_pandas() - chunks = pandas_cursor.iter_chunks() if close_method == "generator" else dataframe_iterator - assert next(chunks)["value"].tolist() == [b"\x00"] - assert not stream.closed - assert not source.closed - - if close_method in ("iterator", "generator"): + ) as (cursor, fs): + cursor.execute("SELECT 1", engine=engine, nrows=2 if close_method == "nrows" else None) + dataframe_iterator = cursor.as_pandas() + chunks = cursor.iter_chunks() if close_method == "generator" else dataframe_iterator + assert next(chunks)["value"].tolist() == [b"\x00"] + opened = list(fs.opened) + assert opened + assert not any(file.closed for file in opened) + + if close_method in ("iterator", "generator"): + chunks.close() + elif close_method == "context": + with chunks: + assert next(chunks)["value"].tolist() == [b"\x00"] + elif close_method == "exhaust": + assert sum(len(chunk) for chunk in chunks) == 99 + elif close_method == "nrows": + assert sum(len(chunk) for chunk in chunks) == 1 + elif close_method == "get_chunk": + assert len(chunks.get_chunk(99)) == 99 + with pytest.raises(StopIteration): + chunks.get_chunk(1) + elif close_method == "cursor": + cursor.close() + else: + cursor.execute("SELECT 1") + + assert all(file.closed for file in opened) chunks.close() - elif close_method == "context": - with chunks: - assert next(chunks)["value"].tolist() == [b"\x00"] - elif close_method == "exhaust": - assert sum(len(chunk) for chunk in chunks) == 99 - elif close_method == "nrows": - assert sum(len(chunk) for chunk in chunks) == 1 - elif close_method == "get_chunk": - assert len(chunks.get_chunk(99)) == 99 with pytest.raises(StopIteration): - chunks.get_chunk(1) - elif close_method == "cursor": - pandas_cursor.close() - else: - pandas_cursor.execute("SELECT 1") - - assert stream.closed - assert source.closed - chunks.close() - with pytest.raises(StopIteration): - next(dataframe_iterator) - with pytest.raises(StopIteration): - dataframe_iterator.get_chunk() - if close_method != "execute": - assert pandas_cursor.fetchone() is None + next(dataframe_iterator) + with pytest.raises(StopIteration): + dataframe_iterator.get_chunk() + if close_method != "execute": + assert cursor.fetchone() is None @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize("read_method", ["next", "get_chunk"]) - def test_binary_csv_stream_close_on_read_error(self, pandas_cursor, engine, read_method): - pandas_cursor.execute( - "SELECT X'00' AS value, rpad('x', 4096, 'x') AS padding " - "FROM UNNEST(sequence(1, 100)) AS t(id)", - engine=engine, + def test_binary_csv_stream_close_on_read_error(self, tmp_path, engine, read_method): + """A failed chunk read closes the binary CSV stream (GH-1225). + + The read raises the conversion error, and later reads and fetches raise + ProgrammingError. No AWS calls; the result file is a local file. + """ + with _offline_cursor( + tmp_path, + [("value", "varbinary"), ("padding", "varchar")], + [("00", "x" * 4096)] * 100, chunksize=1, - dtype={"padding": "int64"}, - ) - stream = pandas_cursor.result_set._csv_stream - source = stream.buffer.raw._reader._file - chunks = pandas_cursor.as_pandas() - assert not stream.closed - assert not source.closed - read_chunk = chunks.__next__ if read_method == "next" else chunks.get_chunk - with pytest.raises(ValueError, match=r"invalid literal|Unable to convert column"): - read_chunk() - assert stream.closed - assert source.closed - # GH-1225: reads and fetches after the failed read raise. - for later in (lambda: next(chunks), chunks.get_chunk, pandas_cursor.fetchone): - with pytest.raises(ProgrammingError, match=r"after a failed read\."): - later() + ) as (cursor, fs): + cursor.execute("SELECT 1", engine=engine, dtype={"padding": "int64"}) + chunks = cursor.as_pandas() + assert fs.opened + assert not any(file.closed for file in fs.opened) + read_chunk = chunks.__next__ if read_method == "next" else chunks.get_chunk + with pytest.raises(ValueError): # noqa: PT011 + read_chunk() + assert all(file.closed for file in fs.opened) + for later in (lambda: next(chunks), chunks.get_chunk, cursor.fetchone): + with pytest.raises(ProgrammingError): + later() def test_binary_converter_override(self, pandas_cursor): pandas_cursor.execute( @@ -176,17 +263,26 @@ def test_binary_converter_override(self, pandas_cursor): pandas_cursor.execute("SELECT X'00ff' AS value", converters=None) assert pandas_cursor.fetchone() == ("00 ff",) - @pytest.mark.parametrize("engine", ["c", "pyarrow"]) - def test_binary_without_converter(self, engine): + @pytest.mark.parametrize( + ("engine", "removed", "second_name"), + [("c", None, "padding"), ("pyarrow", None, "padding"), ("c", "varbinary", "value")], + ids=["c", "pyarrow", "duplicate_name"], + ) + def test_binary_without_converter(self, engine, removed, second_name): + """Without a varbinary converter, binary values are the hex text of the result. + + The PyArrow engine reads the result only without converters; a column with the + name of the binary column keeps its own converter. + """ converter = DefaultPandasTypeConverter() - for type_ in list(converter.mappings): + for type_ in [removed] if removed else list(converter.mappings): converter.remove(type_) with ( connect(cursor_class=PandasCursor, converter=converter) as conn, conn.cursor() as cursor, ): cursor.execute( - "SELECT X'00ff' AS value, %(padding)s AS padding", + f"SELECT X'00ff' AS value, %(padding)s AS {second_name}", {"padding": "x" * 200}, engine=engine, ) @@ -211,33 +307,22 @@ def test_binary_dataframe_column_names(self, pandas_cursor, engine, read_options ) assert pandas_cursor.as_pandas().iloc[0].tolist() == expected - @pytest.mark.parametrize("read_options", [{"quoting": 3}, {"quotechar": "'"}]) + @pytest.mark.parametrize( + "read_options", + [{"quoting": 3}, {"quotechar": "'"}, {"dialect": _QuoteNoneDialect}], + ids=["read_options0", "read_options1", "dialect"], + ) def test_binary_custom_quoting(self, pandas_cursor, read_options): - pandas_cursor.execute("SELECT X'00ff' AS value", **read_options) - assert pandas_cursor.as_pandas().iloc[0].tolist() == ['"00 ff"'] + """Options that change how pandas reads quotes leave binary values as raw text. - @pytest.mark.parametrize("engine", ["c", "python"]) - def test_binary_custom_dialect(self, pandas_cursor, engine): - dialect = csv.excel() - dialect.quoting = csv.QUOTE_NONE + A binary NULL is missing, as pandas reads an empty field. + """ pandas_cursor.execute( - "SELECT CAST(NULL AS VARBINARY) AS value, 'text' AS label", - engine=engine, - dialect=dialect, + "SELECT X'00ff' AS value, CAST(NULL AS VARBINARY) AS null_value", **read_options ) row = pandas_cursor.as_pandas().iloc[0].tolist() - assert pd.isna(row[0]) - assert row[1] == '"text"' - - def test_binary_duplicate_name_without_converter(self): - converter = DefaultPandasTypeConverter() - converter.remove("varbinary") - with ( - connect(cursor_class=PandasCursor, converter=converter) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT true AS value, X'00ff' AS value") - assert cursor.as_pandas().iloc[0].tolist() == [True, "00 ff"] + assert row[0] == '"00 ff"' + assert pd.isna(row[1]) def test_binary_single_null(self, pandas_cursor): pandas_cursor.execute("SELECT CAST(NULL AS VARBINARY) AS value") @@ -248,9 +333,7 @@ def test_binary_single_null(self, pandas_cursor): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -262,38 +345,16 @@ def test_fetchone(self, pandas_cursor, parquet_engine, chunksize): assert pandas_cursor.fetchone() is None @pytest.mark.parametrize( - ("pandas_cursor", "chunksize"), - [ - ({"cursor_kwargs": {"unload": False}}, None), - ({"cursor_kwargs": {"unload": False}}, 1_000), - ({"cursor_kwargs": {"unload": True}}, None), - ], + "pandas_cursor", + [{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}], indirect=["pandas_cursor"], ) - def test_result_set_file_system(self, pandas_cursor, chunksize): + def test_result_set_file_system(self, pandas_cursor): # GH-978: the filesystems that read the results were kept in the fsspec # instance cache with the connection, so the connection was never freed. - # The result set reads through its own filesystem instead of creating - # another one from storage_options. - with patch.object( - S3FileSystem, "__init__", autospec=True, side_effect=S3FileSystem.__init__ - ) as init: - pandas_cursor.execute("SELECT * FROM one_row", chunksize=chunksize) - assert pandas_cursor.fetchall() == [(1,)] - assert init.call_count == 1 - assert not cached_file_systems(pandas_cursor.connection) - if not pandas_cursor.result_set.is_unload: - assert pandas_cursor.result_set._csv_stream.closed - - @pytest.mark.parametrize( - "pandas_cursor", - [{"endpoint_url": f"https://athena.{ENV.region_name}.amazonaws.com"}], - indirect=True, - ) - def test_athena_endpoint_url(self, pandas_cursor): - # GH-576: the S3 requests were sent to Athena's endpoint_url. pandas_cursor.execute("SELECT * FROM one_row") assert pandas_cursor.fetchall() == [(1,)] + assert not cached_file_systems(pandas_cursor.connection) def test_result_sets_share_s3_client(self, pandas_cursor): # GH-1011: each result set and its filesystem built their own S3 client, @@ -310,92 +371,36 @@ def test_result_sets_share_s3_client(self, pandas_cursor): assert all(fs._client is conn.s3_client for fs in file_systems) @pytest.mark.parametrize( - ("query", "expected", "binary"), - [ - ("SELECT * FROM one_row", [(1,)], False), - ("SELECT X'01' AS value", [(b"\x01",)], True), - ], + ("query", "expected"), + [("SELECT * FROM one_row", [(1,)]), ("SELECT X'01' AS value", [(b"\x01",)])], ids=["plain", "binary"], ) - @pytest.mark.parametrize("with_options", [False, True], ids=["none", "options"]) - def test_csv_storage_options(self, pandas_cursor, query, expected, binary, with_options): - # Given storage_options, even None, the CSV output is opened through fsspec - # with them, as pandas does, instead of the result set's filesystem. - storage_options = ( - { - "connection": pandas_cursor.connection, - "default_cache_type": "none", - "skip_instance_cache": True, - } - if with_options - else None + def test_csv_storage_options(self, pandas_cursor, query, expected): + """storage_options given to execute() open the CSV result through fsspec. + + They replace the result set's filesystem, so invalid credentials fail the read. + """ + pandas_cursor.execute( + query, + storage_options={"connection": pandas_cursor.connection, "skip_instance_cache": True}, ) - with patch.object( - S3FileSystem, "open", autospec=True, side_effect=S3FileSystem.open - ) as open_: - pandas_cursor.execute(query, storage_options=storage_options) - assert pandas_cursor.fetchall() == expected - file_systems = [c.args[0] for c in open_.call_args_list] - assert not [fs for fs in file_systems if fs is pandas_cursor.result_set._fs] - if with_options: - assert file_systems - assert all(fs.default_cache_type == "none" for fs in file_systems) - if not binary: - # pandas opens and closes the file itself. - assert pandas_cursor.result_set._csv_stream is None - - @pytest.mark.parametrize("infer_string", [True, False]) - def test_pyarrow_engine_string_values(self, pandas_cursor, infer_string): - # The PyArrow engine keeps the text of strings that look like numbers, and - # a NULL string missing, as the C engine does. - query = """ - SELECT * FROM (VALUES - (1, '1', 'abcdefghijklmnopqrstuvwxyz'), - (2, CAST(NULL AS VARCHAR), 'abcdefghijklmnopqrstuvwxyz'), - (3, 'nan', 'abcdefghijklmnopqrstuvwxyz'), - (4, '007', 'abcdefghijklmnopqrstuvwxyz'), - (5, '1e3', 'abcdefghijklmnopqrstuvwxyz') - ) AS t(id, v, padding) ORDER BY id - """ - with pd.option_context("future.infer_string", infer_string): - with patch( - "pyathena.pandas.result_set._read_csv_with_pyarrow", wraps=_read_csv_with_pyarrow - ) as read_csv_with_pyarrow: - pandas_cursor.execute(query, engine="pyarrow") - actual = pandas_cursor.as_pandas()["v"] - read_csv_with_pyarrow.assert_called_once() - pandas_cursor.execute(query, engine="c") - expected = pandas_cursor.as_pandas()["v"] - pd.testing.assert_series_equal(actual, expected, check_exact=True) - assert actual.iloc[[0, 2, 3, 4]].tolist() == ["1", "nan", "007", "1e3"] - assert pd.isna(actual.iloc[1]) - - def test_pyarrow_engine_multiline_values_across_blocks(self, pandas_cursor): - # The 2.4 MB result spans several 1 MiB pyarrow read blocks, and its - # values contain a newline and quotes. - with patch( - "pyathena.pandas.result_set._read_csv_with_pyarrow", wraps=_read_csv_with_pyarrow - ) as read_csv_with_pyarrow: + assert pandas_cursor.fetchall() == expected + with pytest.raises(OperationalError): pandas_cursor.execute( - """ - SELECT array_join(repeat('x', 600), '') || chr(10) - || '"' || array_join(repeat('y', 598), '') || '"' AS v - FROM UNNEST(sequence(1, 2000)) AS t(i) - """, - engine="pyarrow", + query, + storage_options={ + "key": "invalid", + "secret": "invalid", + "skip_instance_cache": True, + }, ) - df = pandas_cursor.as_pandas() - read_csv_with_pyarrow.assert_called_once() - assert df["v"].tolist() == ["x" * 600 + '\n"' + "y" * 598 + '"'] * 2000 @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine", "chunksize"), [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -463,9 +468,7 @@ def test_get_chunk(self, pandas_cursor, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -480,9 +483,7 @@ def test_fetchall(self, pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -496,9 +497,7 @@ def test_iterator(self, pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -676,10 +675,7 @@ def test_complex(self, pandas_cursor, chunksize): @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], + [({"cursor_kwargs": {"unload": True}}, "auto")], indirect=["pandas_cursor"], ) def test_complex_unload_pyarrow(self, pandas_cursor, parquet_engine): @@ -796,9 +792,7 @@ def test_fetch_no_data(self, pandas_cursor): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -817,9 +811,7 @@ def test_as_pandas(self, pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -966,9 +958,7 @@ def test_complex_as_pandas(self, pandas_cursor, chunksize): @pytest.mark.parametrize( ("pandas_cursor", "chunksize", "options", "sizes"), [ - ({"cursor_kwargs": {"unload": True}}, 1, {}, [1, 1, 1]), ({"cursor_kwargs": {"unload": True}}, 2, {}, [2, 1]), - ({"cursor_kwargs": {"unload": True}}, 10, {}, [3]), ( {"cursor_kwargs": {"unload": True}}, 1, @@ -1005,7 +995,7 @@ def test_unload_chunks(self, pandas_cursor, chunksize, options, sizes): @pytest.mark.parametrize( "pandas_cursor", [{"cursor_kwargs": {"unload": True}}], indirect=["pandas_cursor"] ) - @pytest.mark.parametrize("columns", [["missing"], ["k", "missing"]]) + @pytest.mark.parametrize("columns", [["k", "missing"]]) def test_unload_chunks_unknown_columns(self, pandas_cursor, columns): """Columns that are not in the result raise with chunks as without them.""" query = "SELECT * FROM (VALUES (1, 'a')) AS t(k, s)" @@ -1057,10 +1047,7 @@ def test_unload_chunks_complex(self, pandas_cursor): @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], + [({"cursor_kwargs": {"unload": True}}, "auto")], indirect=["pandas_cursor"], ) def test_complex_unload_as_pandas_pyarrow(self, pandas_cursor, parquet_engine): @@ -1202,205 +1189,59 @@ def test_no_ops(self): cursor.close() conn.close() - def test_auto_determine_chunksize(self): - """Test _auto_determine_chunksize method behavior.""" - - # Mock the parent class initialization - with patch("pyathena.pandas.result_set.AthenaResultSet.__init__"): - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - - # Small file - no chunking - chunksize = result_set._auto_determine_chunksize(10 * 1024 * 1024) # 10MB - assert chunksize is None - - # Medium file - 50K chunks - chunksize = result_set._auto_determine_chunksize(150 * 1024 * 1024) # 150MB - assert chunksize == 50_000 - - # Large file - 100K chunks - chunksize = result_set._auto_determine_chunksize(300 * 1024 * 1024) # 300MB - assert chunksize == 100_000 - - def test_chunksize_configuration_customization(self): - """Test that users can customize chunksize configuration.""" - - # Mock the parent class initialization - with patch("pyathena.pandas.result_set.AthenaResultSet.__init__"): - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - - # Test default behavior - chunksize = result_set._auto_determine_chunksize(200 * 1024 * 1024) # 200MB - assert chunksize == 100_000 # Default AUTO_CHUNK_SIZE_LARGE - - # Customize chunk size - result_set.AUTO_CHUNK_SIZE_LARGE = 250_000 - chunksize = result_set._auto_determine_chunksize(200 * 1024 * 1024) # 200MB - assert chunksize == 250_000 - - # Customize threshold - result_set.LARGE_FILE_THRESHOLD_BYTES = 100 * 1024 * 1024 # 100MB - chunksize = result_set._auto_determine_chunksize(80 * 1024 * 1024) # 80MB - assert chunksize is None # Now below threshold - - def test_get_csv_engine_explicit_specification(self): - """Test _get_csv_engine respects explicit engine specification.""" - - # Mock the parent class initialization - with patch("pyathena.pandas.result_set.AthenaResultSet.__init__"): - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._chunksize = None # Default values - result_set._quoting = 1 - result_set._metadata = None - result_set._query_execution = None - result_set._kwargs = {} - - # Test C engine specification - result_set._engine = "c" - engine = result_set._get_csv_engine() - assert engine == "c" - - # Test Python engine specification - result_set._engine = "python" - engine = result_set._get_csv_engine() - assert engine == "python" - - # Test PyArrow specification (when available and compatible) - result_set._engine = "pyarrow" - with ( - patch.object(result_set, "_get_available_engine", return_value="pyarrow"), - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine() - assert engine == "pyarrow" - - # Test PyArrow with column names that repeat, which it does not rename - with ( - patch.object(result_set, "_get_available_engine", return_value="pyarrow"), - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - patch.object( - type(result_set), - "description", - new_callable=PropertyMock, - return_value=[("x", "integer"), ("x", "integer")], - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - - # Test PyArrow with read options that rename the columns - result_set._kwargs = {"names": ["b", "a"]} - with ( - patch.object(result_set, "_get_available_engine", return_value="pyarrow"), - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - result_set._kwargs = {} - - # Test PyArrow with incompatible chunksize (via parameter) - with ( - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine(chunksize=1000) - assert engine == "c" - - # Test PyArrow with incompatible chunksize (via instance variable) - result_set._chunksize = 1000 - with ( - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - - # Test PyArrow with incompatible quoting - result_set._chunksize = None - result_set._quoting = 0 # Non-default quoting - with ( - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - - # Test PyArrow with incompatible converters - result_set._quoting = 1 # Reset to default - with ( - patch.object( - type(result_set), - "converters", - new_callable=PropertyMock, - return_value={"col1": str}, - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - - # Test PyArrow specification (when unavailable) - with ( - patch.object(result_set, "_get_available_engine", side_effect=ImportError), - patch.object( - type(result_set), "converters", new_callable=PropertyMock, return_value={} - ), - ): - engine = result_set._get_csv_engine() - assert engine == "c" - @pytest.mark.parametrize( - ("keep_default_na", "na_values", "kwargs", "expected"), + ("columns", "rows", "read_options", "expected", "sizes"), [ - (False, ("",), {}, True), - (False, ("",), {"dtype": {"a": "str"}, "parse_dates": ["b"]}, True), - (True, ("",), {}, False), - (False, ("", "NA"), {}, False), - (False, np.array([""]), {}, False), - (False, ("",), {"storage_options": None}, False), - (False, ("",), {"on_bad_lines": "skip"}, False), - (False, ("",), {"dtype": "str"}, False), - (False, ("",), {"parse_dates": "d"}, False), - (False, ("",), {"parse_dates": ("d",)}, False), - ], - ids=[ - "default", - "dtype_parse_dates", - "keep_default_na", - "na_values", - "na_values_array", - "storage_options", - "on_bad_lines", - "single_dtype", - "parse_dates_string", - "parse_dates_tuple", + ( + [("a", "integer"), ("s", "varchar")], + [("1", _PADDING), ("2", _PADDING)], + {"chunksize": 1}, + [(1, _PADDING), (2, _PADDING)], + [1, 1], + ), + ( + [("s", "varchar")], + [(_PADDING,)], + {"quoting": csv.QUOTE_NONE}, + [(f'"{_PADDING}"',)], + [1], + ), + ( + [("x", "integer"), ("x", "varchar")], + [("1", f"007{_PADDING}")], + {}, + [(1, f"007{_PADDING}")], + [1], + ), + ( + [("b", "varbinary"), ("s", "varchar")], + [("00 ff", _PADDING)], + {}, + [(b"\x00\xff", _PADDING)], + [1], + ), ], + ids=["chunksize", "quoting", "duplicate_names", "converters"], ) - def test_reads_csv_with_pyarrow(self, keep_default_na, na_values, kwargs, expected): - # PyAthena reads the CSV result for the PyArrow engine only with the options - # that _read_csv_with_pyarrow() reproduces; pandas reads it otherwise. - with patch("pyathena.pandas.result_set.AthenaResultSet.__init__"): - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._keep_default_na = keep_default_na - result_set._na_values = na_values - result_set._kwargs = kwargs - assert result_set._reads_csv_with_pyarrow() is expected + def test_pyarrow_engine_incompatible_options( + self, tmp_path, columns, rows, read_options, expected, sizes + ): + """engine="pyarrow" reads with the C engine what the PyArrow engine cannot read. + + That is a result in chunks, with another quoting, with columns of the same name, + or with converters. No AWS calls; the result file is a local file. + """ + with _offline_cursor(tmp_path, columns, rows) as (cursor, _): + cursor.execute("SELECT 1", engine="pyarrow", **read_options) + chunks = list(cursor.iter_chunks()) + assert [len(chunk) for chunk in chunks] == sizes + assert [tuple(row) for chunk in chunks for row in chunk.itertuples(index=False)] == expected @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine", "chunksize"), [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), - ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -1414,9 +1255,6 @@ def test_show_columns(self, pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), - ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["pandas_cursor"], ) @@ -1442,10 +1280,7 @@ def test_empty_result_ddl(self, pandas_cursor, parquet_engine, chunksize): @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], + [({"cursor_kwargs": {"unload": True}}, "auto")], indirect=["pandas_cursor"], ) def test_empty_result_dml_unload(self, pandas_cursor, parquet_engine): @@ -1480,7 +1315,6 @@ def test_ctas(self, pandas_cursor): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1510,7 +1344,6 @@ def test_integer_na_values(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1529,7 +1362,6 @@ def test_float_na_values(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1548,7 +1380,6 @@ def test_boolean_na_values(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1563,11 +1394,7 @@ def test_executemany(self, pandas_cursor, parquet_engine, empty_table): @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], + [({"cursor_kwargs": {"unload": False}}, "auto")], indirect=["pandas_cursor"], ) def test_executemany_fetch(self, pandas_cursor, parquet_engine): @@ -1583,7 +1410,6 @@ def test_executemany_fetch(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1601,7 +1427,6 @@ def test_not_skip_blank_lines(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1643,7 +1468,6 @@ def test_null_vs_empty_string(self, pandas_cursor, parquet_engine): [ ({"cursor_kwargs": {"unload": False}}, "auto"), ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), ], indirect=["pandas_cursor"], ) @@ -1771,83 +1595,60 @@ def test_unload_with_filesystem_options(self, pandas_cursor, option): df = pandas_cursor.execute("SELECT * FROM one_row", **kwargs).as_pandas() assert df.to_dict("records") == [{"number_of_rows": 1}] - def test_pandas_cursor_iter_chunks_with_chunksize(self, pandas_cursor): - """Test PandasCursor iter_chunks method with chunksize set.""" - cursor = pandas_cursor - cursor._chunksize = 1000 - cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) as t(number)") - - chunk_count = 0 - for chunk in cursor.iter_chunks(): - chunk_count += 1 - assert isinstance(chunk, pd.DataFrame) - assert len(chunk) > 0 - - assert chunk_count >= 1 - @pytest.mark.parametrize( - ("pandas_cursor", "chunked"), + ("pandas_cursor", "size_options", "read", "expected"), [ - ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, False), - ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, True), + ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, {}, "iter_chunks", [5]), + ( + {"cursor_kwargs": {"auto_optimize_chunksize": True}}, + {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, + "iter_chunks", + [2, 2, 1], + ), + ( + {"cursor_kwargs": {"auto_optimize_chunksize": True}}, + {"AUTO_CHUNK_THRESHOLD_LARGE": 0, "AUTO_CHUNK_SIZE_LARGE": 2}, + "iter_chunks", + [2, 2, 1], + ), + ( + {"cursor_kwargs": {"auto_optimize_chunksize": True}}, + {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, + "as_pandas", + None, + ), + ( + {"cursor_kwargs": {"auto_optimize_chunksize": True, "chunksize": 3}}, + {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, + "iter_chunks", + [3, 2], + ), ], indirect=["pandas_cursor"], + ids=["small", "medium", "large", "medium-as_pandas", "explicit_chunksize"], ) def test_pandas_cursor_auto_optimize_chunksize_enabled( - self, pandas_cursor, chunked, monkeypatch + self, pandas_cursor, size_options, read, expected, monkeypatch ): - """Test that as_pandas() returns the whole result with auto_optimize_chunksize.""" - if chunked: - # Make the five-row result exceed the threshold and read it two rows at a time. + """auto_optimize_chunksize chunks a CSV result by its size; a chunksize wins. + + as_pandas() joins the chunks into the whole result. + """ + if size_options: + # Make the five-row result exceed the size threshold, one byte per row. monkeypatch.setattr(AthenaPandasResultSet, "LARGE_FILE_THRESHOLD_BYTES", 0) monkeypatch.setattr(AthenaPandasResultSet, "ESTIMATED_BYTES_PER_ROW", 1) - monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_THRESHOLD_MEDIUM", 0) - monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_SIZE_MEDIUM", 2) + for name, value in size_options.items(): + monkeypatch.setattr(AthenaPandasResultSet, name, value) pandas_cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) AS t(number)") - reader = pandas_cursor.result_set._df_iter._reader - assert isinstance(reader, TextFileReader) is chunked - - df = pandas_cursor.as_pandas() - assert isinstance(df, pd.DataFrame) - assert df["number"].tolist() == [1, 2, 3, 4, 5] - assert df.index.tolist() == [0, 1, 2, 3, 4] - - def test_pandas_cursor_auto_optimize_chunksize_disabled(self, pandas_cursor): - """Test PandasCursor with auto_optimize_chunksize disabled (default).""" - cursor = pandas_cursor - cursor._chunksize = None # No explicit chunksize - cursor._auto_optimize_chunksize = False # Default behavior - cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) as t(number)") - - # Should return DataFrame (no chunking) - result = cursor.as_pandas() - assert isinstance(result, pd.DataFrame) - - def test_pandas_cursor_explicit_chunksize_overrides_auto_optimize(self, pandas_cursor): - """Test that explicit chunksize takes precedence over auto-optimization.""" - cursor = pandas_cursor - cursor._chunksize = 1000 # Explicit chunksize - cursor._auto_optimize_chunksize = True # Should be ignored - cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) as t(number)") - - # Should return iterator due to explicit chunksize - result = cursor.as_pandas() - assert isinstance(result, PandasDataFrameIterator) - - def test_pandas_cursor_iter_chunks_without_chunksize(self, pandas_cursor): - """Test PandasCursor iter_chunks method without chunksize (single DataFrame).""" - cursor = pandas_cursor - cursor._chunksize = None - cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) as t(number)") - - chunk_count = 0 - for chunk in cursor.iter_chunks(): - chunk_count += 1 - assert isinstance(chunk, pd.DataFrame) - assert len(chunk) > 0 - - # Should yield exactly one chunk (the entire DataFrame) - assert chunk_count == 1 + if read == "as_pandas": + df = pandas_cursor.as_pandas() + assert df["number"].tolist() == [1, 2, 3, 4, 5] + assert df.index.tolist() == [0, 1, 2, 3, 4] + else: + chunks = list(pandas_cursor.iter_chunks()) + assert [len(chunk) for chunk in chunks] == expected + assert pd.concat(chunks)["number"].tolist() == [1, 2, 3, 4, 5] def test_pandas_cursor_whole_result_reused(self, pandas_cursor): """Test that as_pandas() and iter_chunks() do not consume the fetched rows.""" @@ -1863,100 +1664,6 @@ def test_pandas_cursor_whole_result_reused(self, pandas_cursor): assert pandas_cursor.as_pandas() is df assert pandas_cursor.fetchall() == [(2,), (3,)] - def test_pandas_cursor_chunked_vs_regular_same_data(self, pandas_cursor): - """Test that chunked and regular reading produce the same data.""" - query = "SELECT * FROM many_rows LIMIT 100" # Use a reasonable size for testing - - # Regular reading (no chunksize) - regular_cursor = pandas_cursor - regular_cursor._chunksize = None - regular_cursor.execute(query) - regular_df = regular_cursor.as_pandas() - - # Chunked reading - reuse the same cursor but set chunksize - chunked_cursor = pandas_cursor - chunked_cursor._chunksize = 25 - chunked_cursor.execute(query) - chunked_dfs = list(chunked_cursor.iter_chunks()) - - # Combine chunks - combined_df = pd.concat(chunked_dfs, ignore_index=True) - - # Should have the same data - pd.testing.assert_frame_equal(regular_df.sort_index(), combined_df.sort_index()) - - # Should have multiple chunks - assert len(chunked_dfs) > 1 - - # Each chunk should be <= chunksize - for chunk in chunked_dfs: - assert len(chunk) <= 25 - - def test_pandas_cursor_actual_chunking_behavior(self, pandas_cursor): - """Test that chunking actually limits memory usage by checking chunk sizes.""" - # Use a query that returns more rows than chunksize - cursor = pandas_cursor - cursor._chunksize = 10 - cursor.execute("SELECT * FROM many_rows LIMIT 50") - - result = cursor.as_pandas() - assert isinstance(result, PandasDataFrameIterator) - - chunk_sizes = [] - total_rows = 0 - - for chunk in result: - chunk_size = len(chunk) - chunk_sizes.append(chunk_size) - total_rows += chunk_size - - # Each chunk should not exceed chunksize (except possibly the last one) - assert chunk_size <= 10 - - # Should have processed all 50 rows - assert total_rows == 50 - - # Should have multiple chunks - assert len(chunk_sizes) > 1 - - # Most chunks should be exactly chunksize (10), last might be smaller - # All but the last chunk should be exactly 10 - for size in chunk_sizes[:-1]: - assert size == 10 - # Last chunk should be <= 10 - assert chunk_sizes[-1] <= 10 - - def test_pandas_cursor_iter_chunks_consistency(self, pandas_cursor): - """Test that iter_chunks() and direct iteration give same results.""" - cursor = pandas_cursor - cursor._chunksize = 15 - cursor.execute("SELECT * FROM many_rows LIMIT 45") - - # Method 1: Using iter_chunks() - chunks_via_method = list(cursor.iter_chunks()) - - # Method 2: Direct iteration on as_pandas() - cursor.execute("SELECT * FROM many_rows LIMIT 45") # Re-execute - chunks_via_direct = list(cursor.as_pandas()) - - # Should have same number of chunks - assert len(chunks_via_method) == len(chunks_via_direct) - - # Each corresponding chunk should be identical - for chunk1, chunk2 in zip(chunks_via_method, chunks_via_direct, strict=False): - pd.testing.assert_frame_equal(chunk1, chunk2) - - @pytest.mark.parametrize( - "pandas_cursor", [{"cursor_kwargs": {"chunksize": 2}}], indirect=["pandas_cursor"] - ) - def test_get_chunk_time(self, pandas_cursor): - """get_chunk() converts time columns as iteration does.""" - pandas_cursor.execute( - "SELECT * FROM (VALUES (1, CAST('12:34:56' AS TIME)), (2, NULL)) AS t(i, v) ORDER BY i" - ) - chunk = pandas_cursor.as_pandas().get_chunk() - assert chunk["v"].tolist() == [datetime(2017, 1, 1, 12, 34, 56).time(), None] - @pytest.mark.parametrize( "pandas_cursor", [ @@ -2234,7 +1941,7 @@ def test_integer_and_json_with_null(self, pandas_cursor, kwargs): SELECT * FROM (VALUES (1, BIGINT '9007199254740993', json_parse('9007199254740993'), X'01', json_parse('9007199254740995'), CAST('01:02:03' AS TIME)), - (2, NULL, NULL, NULL, json_parse('9007199254740997'), CAST('04:05:06' AS TIME)) + (2, NULL, NULL, NULL, json_parse('9007199254740997'), CAST(NULL AS TIME)) ) AS t(col_int, col_bigint, col_json, col_binary, col_json_not_null, col_time) ORDER BY col_int """, @@ -2265,7 +1972,4 @@ def column(name): assert column("col_json").tolist() == [9007199254740993, None] assert column("col_json_not_null").tolist() == [9007199254740995, 9007199254740997] assert column("col_binary").tolist() == [b"\x01", None] - assert column("col_time").tolist() == [ - datetime(2017, 1, 1, 1, 2, 3).time(), - datetime(2017, 1, 1, 4, 5, 6).time(), - ] + assert column("col_time").tolist() == [datetime(2017, 1, 1, 1, 2, 3).time(), None] From e6b5f651afbae246875ff846cc58eeb7de5f5552 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:57:10 +0900 Subject: [PATCH 2/7] Remove async and aio pandas tests that repeat the sync result set AsyncPandasCursor and AioPandasCursor read through the same AthenaPandasResultSet as PandasCursor, so their copies of the NA, DDL, and many-rows tests are removed, and their parameter cases keep unload False and True. The forwarding of engine and na_values joins test_read_options. The binary NULL and empty tests stay, because each cursor chooses its own default converter. test_cancel_initial and test_invalid_arraysize of AioPandasCursor stay until #1245 adds their shared aio tests. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/pandas/test_cursor.py | 65 +---- tests/pyathena/pandas/test_async_cursor.py | 291 +-------------------- 2 files changed, 8 insertions(+), 348 deletions(-) diff --git a/tests/pyathena/aio/pandas/test_cursor.py b/tests/pyathena/aio/pandas/test_cursor.py index 5fd863a62..b38b399f3 100644 --- a/tests/pyathena/aio/pandas/test_cursor.py +++ b/tests/pyathena/aio/pandas/test_cursor.py @@ -12,10 +12,7 @@ from pyathena.aio.pandas.cursor import AioPandasCursor from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution -from pyathena.pandas.result_set import AthenaPandasResultSet from pyathena.util import RetryConfig -from tests import ENV -from tests.pyathena.aio.conftest import _aio_connect class TestAioPandasCursor: @@ -35,17 +32,15 @@ async def test_binary_null_vs_empty(self, aio_pandas_cursor): ] async def test_fetchone(self, aio_pandas_cursor): - await aio_pandas_cursor.execute("SELECT * FROM one_row") + assert await aio_pandas_cursor.execute("SELECT * FROM one_row") is aio_pandas_cursor + assert aio_pandas_cursor.description == [ + ("number_of_rows", "integer", None, None, 10, 0, "UNKNOWN") + ] assert aio_pandas_cursor.rownumber == 0 assert await aio_pandas_cursor.fetchone() == (1,) assert aio_pandas_cursor.rownumber == 1 assert await aio_pandas_cursor.fetchone() is None - async def test_fetchmany(self, aio_pandas_cursor): - await aio_pandas_cursor.execute("SELECT * FROM many_rows LIMIT 15") - assert len(await aio_pandas_cursor.fetchmany(10)) == 10 - assert len(await aio_pandas_cursor.fetchmany(10)) == 5 - async def test_fetchall(self, aio_pandas_cursor): await aio_pandas_cursor.execute("SELECT * FROM one_row") assert await aio_pandas_cursor.fetchall() == [(1,)] @@ -59,10 +54,6 @@ async def test_as_pandas(self, aio_pandas_cursor): assert df.columns.tolist() == ["number_of_rows"] assert df["number_of_rows"].iloc[0] == 1 - async def test_execute_returns_self(self, aio_pandas_cursor): - result = await aio_pandas_cursor.execute("SELECT * FROM one_row") - assert result is aio_pandas_cursor - async def test_no_result_set_raises(self, aio_pandas_cursor): with pytest.raises(ProgrammingError): await aio_pandas_cursor.fetchone() @@ -73,53 +64,16 @@ async def test_no_result_set_raises(self, aio_pandas_cursor): with pytest.raises(ProgrammingError): aio_pandas_cursor.as_pandas() - async def test_context_manager(self): - from pyathena.aio.pandas.cursor import AioPandasCursor - - conn = await _aio_connect(schema_name=ENV.schema, cursor_class=AioPandasCursor) - try: - async with conn.cursor() as cursor: - await cursor.execute("SELECT * FROM one_row") - assert await cursor.fetchone() == (1,) - finally: - conn.close() - - async def test_arraysize_default(self, aio_pandas_cursor): - assert aio_pandas_cursor.arraysize == AthenaPandasResultSet.DEFAULT_FETCH_SIZE - async def test_invalid_arraysize(self, aio_pandas_cursor): aio_pandas_cursor.arraysize = 10000 assert aio_pandas_cursor.arraysize == 10000 with pytest.raises(ProgrammingError): aio_pandas_cursor.arraysize = -1 - async def test_description(self, aio_pandas_cursor): - await aio_pandas_cursor.execute("SELECT CAST(1 AS INT) AS foobar FROM one_row") - assert await aio_pandas_cursor.fetchall() == [(1,)] - assert aio_pandas_cursor.description == [ - ("foobar", "integer", None, None, 10, 0, "UNKNOWN") - ] - - async def test_description_initial(self, aio_pandas_cursor): - assert aio_pandas_cursor.description is None - async def test_cancel_initial(self, aio_pandas_cursor): with pytest.raises(ProgrammingError): await aio_pandas_cursor.cancel() - async def test_executemany_fetch(self, aio_pandas_cursor): - await aio_pandas_cursor.executemany( - "SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)] - ) - with pytest.raises(ProgrammingError): - await aio_pandas_cursor.fetchall() - with pytest.raises(ProgrammingError): - await aio_pandas_cursor.fetchmany() - with pytest.raises(ProgrammingError): - await aio_pandas_cursor.fetchone() - with pytest.raises(ProgrammingError): - aio_pandas_cursor.as_pandas() - @pytest.mark.parametrize( "aio_pandas_cursor", [{"cursor_kwargs": {"unload": True}}], @@ -127,19 +81,10 @@ async def test_executemany_fetch(self, aio_pandas_cursor): ) async def test_fetchone_unload(self, aio_pandas_cursor): await aio_pandas_cursor.execute("SELECT * FROM one_row") + assert aio_pandas_cursor.as_pandas()["number_of_rows"].tolist() == [1] assert await aio_pandas_cursor.fetchone() == (1,) assert await aio_pandas_cursor.fetchone() is None - @pytest.mark.parametrize( - "aio_pandas_cursor", - [{"cursor_kwargs": {"unload": True}}], - indirect=["aio_pandas_cursor"], - ) - async def test_as_pandas_unload(self, aio_pandas_cursor): - await aio_pandas_cursor.execute("SELECT * FROM one_row") - df = aio_pandas_cursor.as_pandas() - assert len(df) == 1 - @pytest.mark.parametrize( "execute_kwargs", [ diff --git a/tests/pyathena/pandas/test_async_cursor.py b/tests/pyathena/pandas/test_async_cursor.py index 86fa37146..a4ee52b4a 100644 --- a/tests/pyathena/pandas/test_async_cursor.py +++ b/tests/pyathena/pandas/test_async_cursor.py @@ -1,20 +1,15 @@ import contextlib -import math -import random -import string import time from datetime import datetime from random import randint from unittest.mock import MagicMock, patch -import numpy as np import pandas as pd import pytest from pyathena.error import NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.pandas.async_cursor import AsyncPandasCursor -from pyathena.pandas.result_set import AthenaPandasResultSet from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig from tests import ENV @@ -42,9 +37,7 @@ def test_binary_null_vs_empty(self, async_pandas_cursor): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -63,9 +56,7 @@ def test_fetchone(self, async_pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -105,9 +96,7 @@ def test_get_chunk(self, async_pandas_cursor, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -128,9 +117,7 @@ def test_fetchall(self, async_pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -149,9 +136,7 @@ def test_iterator(self, async_pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -179,9 +164,7 @@ def test_invalid_arraysize(self, async_pandas_cursor): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -206,10 +189,7 @@ def test_description(self, async_pandas_cursor, parquet_engine, chunksize): ("async_pandas_cursor", "parquet_engine", "chunksize"), [ ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -290,10 +270,7 @@ def test_query_execution(self, async_pandas_cursor, parquet_engine, chunksize): ("async_pandas_cursor", "parquet_engine", "chunksize"), [ ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -317,8 +294,6 @@ def test_poll(self, async_pandas_cursor, parquet_engine, chunksize): ("async_pandas_cursor", "parquet_engine", "chunksize"), [ ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], @@ -340,9 +315,7 @@ def test_bad_query(self, async_pandas_cursor, parquet_engine, chunksize): [ ({"cursor_kwargs": {"unload": False}}, "auto", None), ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), ], indirect=["async_pandas_cursor"], ) @@ -359,28 +332,6 @@ def test_as_pandas(self, async_pandas_cursor, parquet_engine, chunksize): assert df.shape[1] == 1 assert [(row["number_of_rows"],) for _, row in df.iterrows()] == [(1,)] - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine", "chunksize"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), - ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), - ], - indirect=["async_pandas_cursor"], - ) - def test_many_as_pandas(self, async_pandas_cursor, parquet_engine, chunksize): - query_id, future = async_pandas_cursor.execute( - "SELECT * FROM many_rows", engine=parquet_engine, chunksize=chunksize - ) - df = future.result().as_pandas() - if chunksize: - df = pd.concat((d for d in df), ignore_index=True) - assert df.shape[0] == 10000 - assert df.shape[1] == 1 - assert [(row["a"],) for _, row in df.iterrows()] == [(i,) for i in range(10000)] - def test_cancel(self, async_pandas_cursor): query_id, future = async_pandas_cursor.execute( """ @@ -409,232 +360,13 @@ def test_no_ops(self): cursor.close() conn.close() - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine", "chunksize"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), - ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), - ], - indirect=["async_pandas_cursor"], - ) - def test_show_columns(self, async_pandas_cursor, parquet_engine, chunksize): - query_id, future = async_pandas_cursor.execute( - "SHOW COLUMNS IN one_row", engine=parquet_engine, chunksize=chunksize - ) - result_set = future.result() - assert result_set.description == [("field", "string", None, None, 0, 0, "UNKNOWN")] - assert result_set.fetchall() == [("number_of_rows ",)] - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine", "chunksize"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto", None), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000), - ({"cursor_kwargs": {"unload": False}}, "auto", 1_000_000), - ({"cursor_kwargs": {"unload": True}}, "auto", None), - ({"cursor_kwargs": {"unload": True}}, "pyarrow", None), - ], - indirect=["async_pandas_cursor"], - ) - def test_empty_result_ddl(self, async_pandas_cursor, parquet_engine, chunksize): - table = "test_pandas_cursor_empty_result_" + "".join( - random.choices(string.ascii_lowercase + string.digits, k=10) - ) - query_id, future = async_pandas_cursor.execute( - f""" - CREATE EXTERNAL TABLE IF NOT EXISTS - {ENV.schema}.{table} (number_of_rows INT) - ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' - LINES TERMINATED BY '\n' STORED AS TEXTFILE - LOCATION '{ENV.s3_staging_dir}{ENV.schema}/{table}/' - """, - engine=parquet_engine, - chunksize=chunksize, - ) - df = future.result().as_pandas() - if chunksize: - df = pd.concat((d for d in df), ignore_index=True) - assert df.shape[0] == 0 - assert df.shape[1] == 0 - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_empty_result_dml_unload(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - """ - SELECT * FROM one_row LIMIT 0 - """, - engine=parquet_engine, - ) - df = future.result().as_pandas() - assert df.shape[0] == 0 - assert df.shape[1] == 0 - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_integer_na_values(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - """ - SELECT * FROM integer_na_values - """, - engine=parquet_engine, - ) - df = future.result().as_pandas() - if async_pandas_cursor._unload: - rows = [ - ( - True if math.isnan(row["a"]) else row["a"], - True if math.isnan(row["b"]) else row["b"], - ) - for _, row in df.iterrows() - ] - # If the UNLOAD option is enabled, it is converted to float for some reason. - assert rows == [(1.0, 2.0), (1.0, True), (True, True)] - else: - rows = [(row["a"], row["b"]) for _, row in df.iterrows()] - assert rows == [(1, 2), (1, pd.NA), (pd.NA, pd.NA)] - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_float_na_values(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - """ - SELECT * FROM (VALUES (0.33), (NULL)) AS t (col) - """, - engine=parquet_engine, - ) - df = future.result().as_pandas() - rows = [(row["col"],) for _, row in df.iterrows()] - np.testing.assert_equal(rows, [(0.33,), (np.nan,)]) - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_boolean_na_values(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - """ - SELECT * FROM boolean_na_values - """, - engine=parquet_engine, - ) - df = future.result().as_pandas() - rows = [(row["a"], row["b"]) for _, row in df.iterrows()] - assert rows == [(True, False), (False, None), (None, None)] - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_not_skip_blank_lines(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - """ - SELECT col FROM (VALUES (1), (NULL)) AS t (col) - """, - engine=parquet_engine, - ) - result_set = future.result() - assert len(result_set.fetchall()) == 2 - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_empty_and_null_string(self, async_pandas_cursor, parquet_engine): - # TODO https://github.com/pyathena-dev/PyAthena/issues/118 - query = """ - SELECT * FROM (VALUES ('', 'a'), ('N/A', 'a'), ('NULL', 'a'), (NULL, 'a')) - AS t (col1, col2) - """ - query_id, future = async_pandas_cursor.execute(query, engine=parquet_engine) - result_set = future.result() - if async_pandas_cursor._unload: - # NULL and empty characters are correctly converted when the UNLOAD option is enabled. - np.testing.assert_equal( - result_set.fetchall(), - [("", "a"), ("N/A", "a"), ("NULL", "a"), (np.nan, "a")], - ) - else: - np.testing.assert_equal( - result_set.fetchall(), - [(np.nan, "a"), ("N/A", "a"), ("NULL", "a"), (np.nan, "a")], - ) - query_id, future = async_pandas_cursor.execute(query, na_values=None, engine=parquet_engine) - result_set = future.result() - if async_pandas_cursor._unload: - # NULL and empty characters are correctly converted when the UNLOAD option is enabled. - np.testing.assert_equal( - result_set.fetchall(), - [("", "a"), ("N/A", "a"), ("NULL", "a"), (np.nan, "a")], - ) - else: - assert result_set.fetchall() == [ - ("", "a"), - ("N/A", "a"), - ("NULL", "a"), - ("", "a"), - ] - - @pytest.mark.parametrize( - ("async_pandas_cursor", "parquet_engine"), - [ - ({"cursor_kwargs": {"unload": False}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "auto"), - ({"cursor_kwargs": {"unload": True}}, "pyarrow"), - ], - indirect=["async_pandas_cursor"], - ) - def test_null_decimal_value(self, async_pandas_cursor, parquet_engine): - query_id, future = async_pandas_cursor.execute( - "SELECT CAST(null AS DECIMAL) AS col_decimal", - engine=parquet_engine, - ) - result_set = future.result() - assert result_set.fetchall() == [(None,)] - @pytest.mark.parametrize( "execute_kwargs", [ {}, { + "engine": "python", + "na_values": None, "block_size": 2048, "cache_type": "none", "s3_max_workers": 3, @@ -648,6 +380,7 @@ def test_read_options(self, execute_kwargs): No AWS calls; the query and its result set are mocked. """ cursor_kwargs = { + "engine": "c", "block_size": 1024, "cache_type": "bytes", "auto_optimize_chunksize": True, @@ -674,21 +407,3 @@ def test_read_options(self, 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 - - @pytest.mark.parametrize( - "async_pandas_cursor", - [{"cursor_kwargs": {"auto_optimize_chunksize": True}}], - indirect=True, - ) - def test_auto_optimize_chunksize(self, async_pandas_cursor, monkeypatch): - """auto_optimize_chunksize given to the cursor chunks the CSV result.""" - # Make the five-row result exceed the threshold and read it two rows at a time. - monkeypatch.setattr(AthenaPandasResultSet, "LARGE_FILE_THRESHOLD_BYTES", 0) - monkeypatch.setattr(AthenaPandasResultSet, "ESTIMATED_BYTES_PER_ROW", 1) - monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_THRESHOLD_MEDIUM", 0) - monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_SIZE_MEDIUM", 2) - _, future = async_pandas_cursor.execute( - "SELECT number FROM (VALUES (1), (2), (3), (4), (5)) AS t(number)" - ) - result_set = future.result() - assert [df["number"].tolist() for df in result_set.iter_chunks()] == [[1, 2], [3, 4], [5]] From 3eb41841065a6c8821c1b91e27aafe941b697463 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:57:10 +0900 Subject: [PATCH 3/7] Read the pandas result-set tests through the result set's constructor The offline result-set tests build AthenaPandasResultSet over local result files, with the columns from a mocked GetQueryResults response, instead of asserting private predicates and the exact arguments of internal read_csv() and read_parquet() calls. They check the values, chunk sizes, and open options of the reads, and the pyarrow engine options move here from the cursor tests. The to_sql partition test compares its rows instead of asserting a non-empty tuple, and the unload converter's identity test is removed. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_converter.py | 12 +- tests/pyathena/pandas/test_result_set.py | 650 +++++++++++------------ tests/pyathena/pandas/test_util.py | 7 +- 3 files changed, 306 insertions(+), 363 deletions(-) diff --git a/tests/pyathena/pandas/test_converter.py b/tests/pyathena/pandas/test_converter.py index 6b820098b..512bd79a2 100644 --- a/tests/pyathena/pandas/test_converter.py +++ b/tests/pyathena/pandas/test_converter.py @@ -5,10 +5,7 @@ # # SPDX-License-Identifier: MIT -from pyathena.pandas.converter import ( - DefaultPandasTypeConverter, - DefaultPandasUnloadTypeConverter, -) +from pyathena.pandas.converter import DefaultPandasTypeConverter class TestDefaultPandasTypeConverter: @@ -23,10 +20,3 @@ def test_convert_delegates_to_mapping(self): assert converter.convert("boolean", "true") is True assert converter.convert("varchar", "hello") == "hello" assert converter.convert("varchar", None) is None - - -class TestDefaultPandasUnloadTypeConverter: - def test_convert_delegates_to_default(self): - """convert() dispatches through the default converter instead of returning None.""" - converter = DefaultPandasUnloadTypeConverter() - assert converter.convert("varchar", "hello") == "hello" diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index e1f3b47de..c284a97e5 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -10,13 +10,14 @@ import io import threading from collections import defaultdict -from unittest.mock import ANY, MagicMock, PropertyMock, patch, sentinel +from unittest.mock import MagicMock, PropertyMock, patch import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq import pytest +from fsspec.implementations.dirfs import DirFileSystem from fsspec.implementations.local import LocalFileSystem from pandas.testing import assert_frame_equal @@ -27,7 +28,6 @@ from pyathena.pandas.result_set import ( AthenaPandasResultSet, PandasDataFrameIterator, - _JSONConverter, _no_trunc_date, _read_csv_with_pyarrow, ) @@ -94,6 +94,7 @@ def test_iterrows_continuous_row_numbers(self): ("csv", "read_csv_kwargs"), [ ("id,kind\n10,a\n11,a\n12,b\n13,b\n14,c\n", {}), + ("id,kind\n10,a\n11,b\n", {}), ("id,kind\n10,a\n11,a\n12,b\n13,b\n14,c\n", {"index_col": "id"}), ("id,kind\n10,c\n11,c\n12,a\n13,a\n14,b\n", {"dtype": {"kind": "category"}}), ("id,kind\n10,\n11,\n12,b\n13,a\n14,\n", {"dtype": {"kind": "category"}}), @@ -108,6 +109,7 @@ def test_iterrows_continuous_row_numbers(self): ], ids=[ "default_index", + "single_chunk", "index_col", "category", "category_null_chunk", @@ -134,13 +136,6 @@ def test_as_pandas_remaining_chunks(self): assert df.index.tolist() == [2, 3, 4] assert df_iter.as_pandas().empty - def test_as_pandas_single_dataframe(self): - """A DataFrame that was read at once is returned as is.""" - df = pd.DataFrame({"n": [1, 2]}) - df_iter = PandasDataFrameIterator(df, _no_trunc_date) - - assert df_iter.as_pandas() is df - @pytest.mark.parametrize( "read", [PandasDataFrameIterator.__next__, PandasDataFrameIterator.get_chunk], @@ -167,17 +162,6 @@ def chunks(): later(df_iter) -_FS = sentinel.pyathena_fs -_USER_FS = sentinel.user_fs - - -_READ_MODES = [ - pytest.param(None, {}, id="dataframe"), - pytest.param(1, {}, id="chunksize"), - pytest.param(None, {"iterator": True}, id="iterator"), -] - - class TestAthenaPandasResultSet: def test_fetch_across_chunks(self): """The fetch methods continue across chunks, reading a chunk when its rows are needed.""" @@ -212,27 +196,6 @@ def chunks(): assert result_set.fetchone() is None assert result_set.rownumber == 10 - def test_fetchall_across_chunks(self): - """fetchall() returns the remaining rows of all chunks, skipping empty chunks.""" - chunks = [ - pd.DataFrame({"n": [0, 1]}), - pd.DataFrame({"n": []}, dtype="int64"), - pd.DataFrame({"n": [2, 3]}, index=[2, 3]), - ] - result_set = _fetching_result_set(iter(chunks), arraysize=3) - assert result_set.fetchone() == (0,) - assert result_set.fetchall() == [(1,), (2,), (3,)] - assert result_set.rownumber == 4 - assert result_set.fetchall() == [] - assert result_set.rownumber == 4 - - def test_fetch_duplicate_column_names(self): - """Columns with the same label keep their own values by position.""" - df = pd.DataFrame([[1, "a", 2], [3, "b", 4]], columns=["x", "y", "x"]) - result_set = _fetching_result_set(df, arraysize=3) - assert result_set.fetchone() == (1, "a", 2) - assert result_set.fetchmany() == [(3, "b", 4)] - @pytest.mark.parametrize( ("seconds", "method"), [ @@ -312,37 +275,6 @@ def test_fetch_after_close(self): assert result_set.fetchall() == [] assert result_set.rownumber is None - @pytest.mark.parametrize("engine", ["auto", "c", "python", "pyarrow"]) - @pytest.mark.parametrize( - ("output_location", "file_size_bytes", "pyarrow_engine"), - [ - ("s3://bucket/result.txt", None, "c"), - ("s3://bucket/result.txt", 99, "c"), - ("s3://bucket/result.txt", 100, "c"), - ("s3://bucket/result.txt", 101, "c"), - ("s3://bucket/result.csv", None, "pyarrow"), - ("s3://bucket/result.csv", 99, "c"), - ("s3://bucket/result.csv", 100, "pyarrow"), - ("s3://bucket/result.csv", 101, "pyarrow"), - (None, None, "pyarrow"), - ], - ) - def test_get_csv_engine_result_format( - self, engine, output_location, file_size_bytes, pyarrow_engine - ): - """PyArrow falls back for DDL text while C and Python choices are preserved.""" - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._query_execution = MagicMock(output_location=output_location) - result_set._metadata = None - result_set._converter = DefaultPandasTypeConverter() - result_set._engine = engine - result_set._chunksize = None - result_set._quoting = csv.QUOTE_ALL - result_set._kwargs = {} - - expected = {"auto": "c", "c": "c", "python": "python", "pyarrow": pyarrow_engine}[engine] - assert result_set._get_csv_engine(file_size_bytes) == expected - @pytest.mark.parametrize("infer_string", [True, False]) @pytest.mark.parametrize("read_csv_kwargs", [{}, {"keep_default_na": True}]) @pytest.mark.parametrize( @@ -368,143 +300,85 @@ def test_get_csv_engine_result_format( ids=["show_tables", "show_columns", "describe"], ) def test_read_csv_ddl_preserves_numeric_looking_names( - self, data, names, expected_names, read_csv_kwargs, infer_string + self, tmp_path, data, names, expected_names, read_csv_kwargs, infer_string ): - """Reading DDL results preserves numeric-looking names and their padding.""" - data = data.encode() - assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES - - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._query_execution = MagicMock( - output_location="s3://bucket/result.txt", substatement_type=None - ) - result_set._converter = DefaultPandasTypeConverter() - result_set._engine = "pyarrow" - result_set._chunksize = None - result_set._auto_optimize_chunksize = False - result_set._quoting = csv.QUOTE_ALL - result_set._keep_default_na = False - result_set._na_values = ("",) - result_set._kwargs = read_csv_kwargs - result_set._block_size = None - result_set._cache_type = None - result_set._fs = MagicMock() - result_set._fs.open.return_value = stream = io.BytesIO(data) - description = [(name, "string", None, None, 0, 0, "UNKNOWN") for name in names] - - with ( - pd.option_context("future.infer_string", infer_string), - patch.object( - AthenaPandasResultSet, - "description", - new_callable=PropertyMock, - return_value=description, - ), - patch.object(result_set, "_get_content_length", return_value=len(data)), - patch("pandas.read_csv", wraps=pd.read_csv) as read_csv, - ): - expected = pd.read_csv(io.BytesIO(data), **result_set._get_csv_read_options("c", None)) - read_csv.reset_mock() - actual = result_set._read_csv() + """A DDL result read with engine="pyarrow" keeps numeric-looking names and their padding. + It is read as with the C engine, and the result file is closed. + """ + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + types = [(name, "string") for name in names] + frames = [] + with pd.option_context("future.infer_string", infer_string): + for engine in ("pyarrow", "c"): + fs = _BucketFileSystem(tmp_path / engine) + frames.append( + _local_result_set( + fs, types, data, name="result.txt", engine=engine, **read_csv_kwargs + ).as_pandas() + ) + assert all(file.closed for _, file in fs.opened) + actual, expected = frames assert_frame_equal(actual, expected, check_exact=True) assert actual.iloc[:, 0].tolist() == expected_names - read_csv.assert_called_once() - assert read_csv.call_args.kwargs["engine"] == "c" - result_set._fs.open.assert_called_once_with( - "s3://bucket/result.txt", mode="rb", **result_set._get_csv_open_options() - ) - assert stream.closed - @pytest.mark.parametrize( - ("chunksize", "execute_kwargs", "expected"), - [ - (None, {}, False), - (10, {}, True), - (10, {"columns": ["a", "b"]}, True), - (10, {"columns": ("a",)}, True), - (10, {"dtype_backend": "numpy_nullable"}, True), - (10, {"dtype_backend": "pyarrow", "columns": ["a"]}, True), - (10, {"columns": []}, False), - (10, {"columns": ["a", "a"]}, False), - (10, {"columns": "a"}, False), - (10, {"columns": [0]}, False), - (10, {"columns": ["s.a"]}, False), - (10, {"dtype_backend": None}, False), - (10, {"dtype_backend": "invalid"}, False), - (10, {"filters": [("a", "=", 1)]}, False), - (10, {"storage_options": {}}, False), - ], - ) - def test_reads_parquet_in_chunks(self, chunksize, execute_kwargs, expected): - """Only chunksize with supported read options reads UNLOAD results in chunks.""" - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._chunksize = chunksize - result_set._kwargs = execute_kwargs - assert result_set._reads_parquet_in_chunks() is expected + def test_read_small_csv_with_pyarrow_engine(self, tmp_path): + """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine.""" + data = '"d","t"\n"2024-01-01","2024-01-01 00:00:00.000"\n' + assert len(data) < AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + types = [("d", "date"), ("t", "timestamp")] + actual, expected = ( + _local_result_set( + _BucketFileSystem(tmp_path / engine), types, data, engine=engine + ).as_pandas() + for engine in ("pyarrow", "c") + ) + assert_frame_equal(actual, expected, check_exact=True) @pytest.mark.parametrize( - ("block_size", "cache_type", "expected"), + ("kwargs", "expected"), [ - ( - None, - None, - { - "block_size": 32 * 2**20, - "cache_type": "background", - "cache_options": {"maxblocks": 2}, - }, + pytest.param({}, ["NA", "null", "x", None, "y", "z", "w"], id="default"), + pytest.param( + {"keep_default_na": True}, + [None, None, "x", None, "y", "z", "w"], + id="keep_default_na", ), - (5 * 2**20, None, {}), - (None, "bytes", {}), - (5 * 2**20, "readahead", {}), - ], - ) - def test_get_csv_open_options(self, block_size, cache_type, expected): - """CSV results are prefetched unless the cursor sets block_size or cache_type.""" - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) - result_set._block_size = block_size - result_set._cache_type = cache_type - assert result_set._get_csv_open_options() == expected - - @pytest.mark.parametrize( - ("execute_kwargs", "path", "filesystem_kwargs"), - [ - ({}, "bucket/unload/", {"filesystem": _FS}), - ({"filesystem": _USER_FS}, "bucket/unload/", {"filesystem": _USER_FS}), - ({"filesystem": None}, "s3://bucket/unload/", {"filesystem": None}), - ({"storage_options": {"anon": True}}, "s3://bucket/unload/", {}), - ({"storage_options": None}, "s3://bucket/unload/", {}), + pytest.param( + {"on_bad_lines": "skip"}, ["NA", "null", None, "y", "z", "w"], id="on_bad_lines" + ), + pytest.param( + {"dtype": "str"}, ["NA", "null", "x", None, "y", "z", "w"], id="single_dtype" + ), + pytest.param({"parse_dates": ("d",)}, OperationalError, id="parse_dates_tuple"), ], - ids=["default", "filesystem", "filesystem-none", "storage-options", "storage-options-none"], ) - def test_read_parquet_filesystem(self, execute_kwargs, path, filesystem_kwargs): - """filesystem or storage_options given to execute() replace PyAthena's filesystem. + def test_read_csv_with_pyarrow_engine_options(self, tmp_path, kwargs, expected): + """With engine="pyarrow", options that PyAthena's own read does not reproduce reach pandas. - No AWS calls; the manifest and pandas.read_parquet are mocked. + pandas applies them, or rejects the ones that it does not accept. """ - result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) # bypass __init__ - result_set._unload_location = None - result_set._engine = "pyarrow" - result_set._fs = _FS - result_set._kwargs = dict(execute_kwargs) - with ( - patch.object( - AthenaPandasResultSet, - "_read_data_manifest", - return_value=["s3://bucket/unload/0.parquet"], - ), - patch("pandas.read_parquet") as read_parquet, - ): - result_set._read_parquet("pyarrow") - assert read_parquet.call_args.args == (path,) - assert read_parquet.call_args.kwargs == { - "engine": "pyarrow", - "use_threads": True, - "coerce_int96_timestamp_unit": "us", - **execute_kwargs, - **filesystem_kwargs, - } + rows = ["NA", "null", "x", "", "y", "z", "w"] + dates = [f"2024-01-0{i + 1}" for i in range(len(rows))] + lines = [f'"{v}","{d}"' for v, d in zip(rows, dates, strict=True)] + if "on_bad_lines" in kwargs: + # The row of "x" has a field too many. + lines[2] += ',"extra"' + del dates[2] + data = '"v","d"\n' + "".join(f"{line}\n" for line in lines) + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + types = [("v", "varchar"), ("d", "date")] + fs = _BucketFileSystem(tmp_path) + if expected is OperationalError: + with pytest.raises(OperationalError): + _local_result_set(fs, types, data, engine="pyarrow", **kwargs) + return + df = _local_result_set(fs, types, data, engine="pyarrow", **kwargs).as_pandas() + assert [None if pd.isna(v) else v for v in df["v"]] == expected + if "dtype" in kwargs: + assert df["d"].tolist() == dates + else: + assert df["d"].tolist() == [pd.Timestamp(d) for d in dates] def test_iter_parquet_chunks_null_in_later_file(self, tmp_path): """Top-level integer and boolean columns with NULL anywhere convert alike in every chunk.""" @@ -522,31 +396,6 @@ def test_iter_parquet_chunks_null_in_later_file(self, tmp_path): assert chunk.dtypes.tolist() == whole.dtypes.tolist() assert_frame_equal(pd.concat(chunks), whole, check_exact=True) - @pytest.mark.parametrize("write_statistics", [True, False]) - def test_iter_parquet_chunks_large_integers(self, tmp_path, write_statistics): - """Integers without NULL stay exact, with or without Parquet statistics.""" - chunks, whole = _read_unload_result( - tmp_path, - [pa.table({"i": pa.array([9007199254740993, 1], pa.int64())})], - chunksize=1, - write_statistics=write_statistics, - ) - assert [chunk["i"].dtype for chunk in chunks] == ["int64", "int64"] - assert [int(chunk["i"].iloc[0]) for chunk in chunks] == [9007199254740993, 1] - assert_frame_equal(pd.concat(chunks), whole, check_exact=True) - - def test_iter_parquet_chunks_nested_values(self, tmp_path): - """Nested values convert per chunk; with dtype_backend="pyarrow" they do not.""" - table = pa.table({"li": pa.array([[1], [None]], pa.list_(pa.int64()))}) - chunks, _ = _read_unload_result(tmp_path, [table], chunksize=1) - assert [chunk["li"].iloc[0].dtype for chunk in chunks] == ["int64", "float64"] - chunks, whole = _read_unload_result( - tmp_path / "pyarrow", [table], chunksize=1, dtype_backend="pyarrow" - ) - for chunk in chunks: - assert chunk.dtypes.tolist() == whole.dtypes.tolist() - assert_frame_equal(pd.concat(chunks), whole, check_exact=True) - def test_read_parquet_footers_stopped(self, tmp_path): """GH-1202: the footer reads run in the stop scope of their caller. @@ -589,33 +438,95 @@ def test_as_pandas_parquet_chunks_engine(self, tmp_path): ): result_set._as_pandas() - def test_iter_parquet_chunks_reader_error(self, tmp_path): - """Errors from creating the chunk reader are raised as OperationalError.""" - with pytest.raises(OperationalError, match="chunksize must be positive"): - _read_unload_result(tmp_path, [pa.table({"i": [1]})], chunksize=0) + @pytest.mark.parametrize( + ("kwargs", "sizes", "columns"), + [ + pytest.param({}, [1, 1, 1], ["a", "s"], id="default"), + pytest.param({"columns": ["a"]}, [1, 1, 1], ["a"], id="columns"), + pytest.param({"dtype_backend": "pyarrow"}, [1, 1, 1], ["a", "s"], id="dtype_backend"), + # Other options read the result whole, where pandas applies or rejects them. + pytest.param({"filters": [("a", "=", 1)]}, [1], ["a", "s"], id="filters"), + pytest.param({"columns": ["s.x"]}, [3], ["x"], id="nested_column"), + pytest.param({"columns": ["a", "a"]}, [3], ["a", "a"], id="repeated_columns"), + pytest.param({"dtype_backend": None}, None, None, id="invalid_dtype_backend"), + ], + ) + def test_unload_chunks_with_read_options(self, tmp_path, kwargs, sizes, columns): + """Only chunksize with supported read options reads an UNLOAD result in chunks.""" + tables = [ + pa.table({"a": [1, 2], "s": [{"x": 1}, {"x": 2}]}), + pa.table({"a": [3], "s": [{"x": 3}]}), + ] + fs = _BucketFileSystem(tmp_path) + if sizes is None: + with pytest.raises(OperationalError): + _local_result_set(fs, tables=tables, chunksize=1, **kwargs) + return + chunks = list(_local_result_set(fs, tables=tables, chunksize=1, **kwargs).iter_chunks()) + assert [len(chunk) for chunk in chunks] == sizes + assert [chunk.columns.tolist() for chunk in chunks] == [columns] * len(sizes) + + @pytest.mark.parametrize( + ("option", "value"), + [ + ("filesystem", "user"), + ("storage_options", {}), + ("storage_options", None), + ("filesystem", None), + ], + ids=["filesystem", "storage_options", "storage_options-none", "filesystem-none"], + ) + def test_unload_filesystem_options(self, tmp_path, option, value): + """filesystem or storage_options given to execute(), even None, replace PyAthena's. + + A given filesystem reads its own files. Otherwise pandas opens the S3 files + itself, so the read is observed. + """ + fs = _BucketFileSystem(tmp_path / "pyathena") + tables = [pa.table({"a": [1]})] + if value == "user": + value = _BucketFileSystem(tmp_path / "user") + _write_unload_tables(value, [pa.table({"a": [2]})]) + assert _local_result_set(fs, tables=tables, filesystem=value).fetchall() == [(2,)] + return + with patch("pandas.read_parquet", return_value=pd.DataFrame({"a": [2]})) as read_parquet: + assert _local_result_set(fs, tables=tables, **{option: value}).fetchall() == [(2,)] + # The S3 location, which pandas opens with the given options instead of PyAthena's. + assert read_parquet.call_args.args[0] == "s3://bucket/unload/" + assert read_parquet.call_args.kwargs[option] == value + assert read_parquet.call_args.kwargs.get("filesystem") is None @pytest.mark.parametrize( - ("data", "types"), + ("data", "types", "chunksize"), [ - pytest.param('"v"\n"a"\n"b"\n', {"v": "varchar"}, id="text"), - pytest.param('"b"\n"00 ff"\n', {"b": "varbinary"}, id="binary"), + pytest.param('"v"\n"a"\n"b"\n', [("v", "varchar")], None, id="text"), + pytest.param('"v"\n"a"\n"b"\n', [("v", "varchar")], 1, id="text-chunks"), + pytest.param('"b"\n"00 ff"\n', [("b", "varbinary")], None, id="binary"), ], ) - @pytest.mark.parametrize(("chunksize", "kwargs"), _READ_MODES) - def test_read_csv_open_options(self, data, types, chunksize, kwargs): - """Every kind of CSV read opens the result file with the same options.""" - fs = MagicMock() - with patch.object( - AthenaPandasResultSet, - "_get_csv_open_options", - autospec=True, - return_value={"cache_type": "sentinel"}, - ) as open_options: - result, _ = _read_csv_result(data, types, chunksize=chunksize, fs=fs, **kwargs) - if chunksize or kwargs: - result.close() - open_options.assert_called_once_with(ANY) - assert fs.open.call_args.kwargs["cache_type"] == "sentinel" + @pytest.mark.parametrize( + ("kwargs", "prefetch"), + [ + pytest.param({}, True, id="default"), + pytest.param({"block_size": 5 * 2**20}, False, id="block_size"), + pytest.param({"cache_type": "bytes"}, False, id="cache_type"), + ], + ) + def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefetch): + """A CSV result is read in CSV_BLOCK_SIZE blocks with the background cache. + + With block_size or cache_type, the file is opened with the filesystem's settings. + """ + fs = _BucketFileSystem(tmp_path) + result_set = _local_result_set(fs, types, data, chunksize=chunksize, **kwargs) + assert len(result_set.fetchall()) == len(data.splitlines()) - 1 + [(options, _)] = fs.opened + if prefetch: + assert options["cache_type"] == "background" + assert options["block_size"] == AthenaPandasResultSet.CSV_BLOCK_SIZE + else: + assert "cache_type" not in options + assert "block_size" not in options @pytest.mark.filterwarnings("error::pandas.errors.DtypeWarning") @pytest.mark.parametrize( @@ -625,8 +536,6 @@ def test_read_csv_open_options(self, data, types, chunksize, kwargs): ("engine", "kwargs", "access"), [ pytest.param("c", {}, "whole", id="c"), - pytest.param("c", {"low_memory": True}, "whole", id="low-memory"), - pytest.param("c", {"low_memory": False}, "whole", id="no-low-memory"), pytest.param("auto", {}, "whole", id="auto"), pytest.param("pyarrow", {}, "whole", id="pyarrow-fallback"), pytest.param("python", {}, "whole", id="python"), @@ -641,9 +550,7 @@ def test_read_csv_json_null_without_dtype_warning( ): """JSON NULLs in a later parser block keep exact values without a warning.""" data = "n,j\n" + f"1,{json_value}\n" * 400000 + "2,\n" * 10 - result, options = _read_csv_result( - data, {"n": "integer", "j": "json"}, engine=engine, **kwargs - ) + result, _ = _read_csv_result(data, {"n": "integer", "j": "json"}, engine=engine, **kwargs) if access != "whole": with result: first = result.get_chunk(400005) if access == "get_chunk" else next(result) @@ -652,9 +559,6 @@ def test_read_csv_json_null_without_dtype_warning( else: df = result values = df.index if kwargs.get("index_col") == "j" else df["j"] - assert options["engine"] == ("python" if engine == "python" else "c") - assert options.get("low_memory") == kwargs.get("low_memory") - assert options["converters"]["j"] is _JSONConverter.keep_text assert len(df) == 400010 assert values.dtype == object assert type(values[0]) is type(expected) @@ -676,7 +580,6 @@ def test_read_csv_json_numbers_in_a_parser_block_without_null(self, kwargs): @pytest.mark.filterwarnings("error::pandas.errors.DtypeWarning") @pytest.mark.parametrize("engine", ["c", "python"]) - @pytest.mark.parametrize("low_memory", [None, True, False]) @pytest.mark.parametrize( ("text", "dtype", "values"), [ @@ -687,75 +590,37 @@ def test_read_csv_json_numbers_in_a_parser_block_without_null(self, kwargs): pytest.param('"{""a"": 1}"\n3,[1]', object, [{"a": 1}, [1]], id="containers"), ], ) - def test_read_csv_json_without_null(self, engine, low_memory, text, dtype, values): + def test_read_csv_json_without_null(self, engine, text, dtype, values): """JSON without NULL gets the dtype that pandas infers for a converter's values.""" - kwargs = {} if low_memory is None else {"low_memory": low_memory} - if engine == "python" and low_memory is not None: - pytest.skip("low_memory is a C engine option") - df, options = _read_csv_result( - f"n,j\n1,{text}\n", {"n": "integer", "j": "json"}, engine=engine, **kwargs - ) - assert options.get("low_memory") == low_memory + df, _ = _read_csv_result(f"n,j\n1,{text}\n", {"n": "integer", "j": "json"}, engine=engine) assert df["j"].dtype == dtype assert df["j"].tolist() == values - def test_read_csv_json_multi_index_level(self): - """A json column in a MultiIndex is converted.""" - df, _ = _read_csv_result( - "n,j\n1,true\n2,false\n3,\n", {"n": "integer", "j": "json"}, index_col=["n", "j"] - ) - level = df.index.get_level_values("j") - assert level[:2].tolist() == [True, False] - assert pd.isna(level[2]) - @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize( - ("dtype", "n_dtype"), + ("dtype", "n_dtype", "m_dtype"), [ - pytest.param({"n": "int32"}, "int32", id="mapping"), - pytest.param(defaultdict(lambda: "str", {"n": "int32"}), "int32", id="defaultdict"), - pytest.param("str", "str", id="scalar"), - pytest.param({"n": "int32", "j": "Int64"}, "int32", id="json-column"), + pytest.param({"n": "int32"}, "int32", "int64", id="mapping"), + pytest.param( + defaultdict(lambda: "str", {"n": "int32"}), "int32", "str", id="defaultdict" + ), + pytest.param("str", "str", "str", id="scalar"), + pytest.param({"n": "int32", "j": "Int64"}, "int32", "int64", id="json-column"), ], ) - def test_read_csv_json_with_dtype(self, engine, dtype, n_dtype): - """A dtype given to execute() applies to the other columns, not to json columns.""" - df, options = _read_csv_result( - "n,j\n1,2\n3,\n", {"n": "integer", "j": "json"}, engine=engine, dtype=dtype + def test_read_csv_json_with_dtype(self, engine, dtype, n_dtype, m_dtype): + """A dtype given to execute() replaces PyAthena's and does not apply to json columns.""" + df, _ = _read_csv_result( + "n,m,j\n1,1,2\n3,4,\n", + {"n": "integer", "m": "integer", "j": "json"}, + engine=engine, + dtype=dtype, ) - assert options["dtype"] is dtype assert df["n"].dtype == n_dtype + assert df["m"].dtype == m_dtype assert df["j"].tolist() == [2, None] - @pytest.mark.parametrize( - ("types", "engine", "kwargs"), - [ - pytest.param({"n": "integer", "j": "integer"}, "c", {}, id="no-json"), - pytest.param({"n": "integer", "j": "json"}, "c", {}, id="json"), - pytest.param( - {"n": "integer", "j": "json"}, "c", {"usecols": ["n"]}, id="json-excluded" - ), - pytest.param( - {"n": "integer", "j": "json"}, "c", {"converters": {}}, id="custom-converters" - ), - pytest.param( - {"n": "integer", "j": "json"}, - "c", - {"converters": {"j": int}}, - id="custom-json-converter", - ), - pytest.param({"n": "integer", "j": "json"}, "python", {}, id="python"), - pytest.param({"n": "integer", "j": "integer"}, "pyarrow", {}, id="pyarrow"), - ], - ) - def test_read_csv_keeps_low_memory_default(self, types, engine, kwargs): - """No column type, converter, or engine sets a low_memory option.""" - df, options = _read_csv_result("n,j\n" + "1,2\n" * 30, types, engine=engine, **kwargs) - assert options["engine"] == engine - assert "low_memory" not in options - assert len(df) == 30 - @pytest.mark.filterwarnings("ignore:Both a converter and dtype:pandas.errors.ParserWarning") @pytest.mark.parametrize( ("kwargs", "expected", "dtype"), @@ -799,31 +664,31 @@ def test_read_csv_single_json_column_with_null(self, engine): df, _ = _read_csv_result("j\n1\n\n", {"j": "json"}, engine=engine) assert df["j"].tolist() == [1, None] - def test_read_csv_json_name_in_multi_index_columns(self): - """Columns read with a header of several rows are left as pandas reads them.""" - df, _ = _read_csv_result("j,n\n1,2\n3,4\n", {"j": "json", "n": "integer"}, header=[0, 1]) - assert df.to_dict("list") == {("j", "1"): [3], ("n", "2"): [4]} + @pytest.mark.parametrize("chunksize", [None, 1]) + def test_read_csv_json_conversion_error(self, chunksize): + """A json value that the converter cannot convert raises OperationalError. - def test_read_csv_json_conversion_error(self): - """A json value that the converter cannot convert raises OperationalError.""" - with pytest.raises(OperationalError, match="Expecting property name"): - _read_csv_result("n,j\n1,{\n", {"n": "integer", "j": "json"}) + A chunked read raises when it reads the chunk. + """ - def test_read_csv_json_conversion_error_in_chunk(self): - """A json value in a chunk that the converter cannot convert raises OperationalError.""" - result, _ = _read_csv_result("n,j\n1,{\n", {"n": "integer", "j": "json"}, chunksize=1) - with result, pytest.raises(OperationalError, match="Expecting property name"): - next(result) + def read(): + result, _ = _read_csv_result( + "n,j\n1,{\n", {"n": "integer", "j": "json"}, chunksize=chunksize + ) + if chunksize: + with result: + next(result) - @pytest.mark.parametrize("converters", [{}, {"j": int}]) - def test_read_csv_json_custom_converters(self, converters): + with pytest.raises(OperationalError): + read() + + @pytest.mark.parametrize(("converters", "expected"), [({}, ["[1]"]), ({"j": len}, [3])]) + def test_read_csv_json_custom_converters(self, converters, expected): """Converters given to execute() replace the json conversion.""" - df, options = _read_csv_result( - "n,j\n1,2\n", {"n": "integer", "j": "json"}, converters=converters + df, _ = _read_csv_result( + 'n,j\n1,"[1]"\n', {"n": "integer", "j": "json"}, converters=converters ) - assert options["converters"] is converters - assert "j" not in (options["dtype"] or {}) - assert df["j"].tolist() == [2] + assert df["j"].tolist() == expected # A CSV result of Athena with each type the PyArrow engine reads, then a row of NULLs. @@ -901,7 +766,7 @@ def _pyarrow_read_csv_kwargs(types, tab_separated=False, dtype_overrides=None, * return result_set._get_csv_read_options("pyarrow", None) -def _unload_result_set(tmp_path, tables, chunksize, write_statistics=True, **kwargs): +def _unload_result_set(tmp_path, tables, chunksize, **kwargs): """Write Parquet files and create a result set of an UNLOAD result in them, without AWS. Returns: @@ -911,7 +776,7 @@ def _unload_result_set(tmp_path, tables, chunksize, write_statistics=True, **kwa paths = [] for index, table in enumerate(tables): path = tmp_path / f"{index}.parquet" - pq.write_table(table, path, write_statistics=write_statistics) + pq.write_table(table, path) paths.append(str(path)) result_set = AthenaPandasResultSet.__new__(AthenaPandasResultSet) # bypass __init__ result_set._fs = LocalFileSystem() @@ -926,7 +791,7 @@ def _unload_result_set(tmp_path, tables, chunksize, write_statistics=True, **kwa return result_set, paths -def _read_unload_result(tmp_path, tables, chunksize, write_statistics=True, **kwargs): +def _read_unload_result(tmp_path, tables, chunksize, **kwargs): """Read Parquet files in chunks through the result set, without AWS. The column metadata must be what the dataset of the files reports, read without @@ -936,9 +801,7 @@ def _read_unload_result(tmp_path, tables, chunksize, write_statistics=True, **kw The chunks, and the whole result read by pandas.read_parquet() with the same options. """ - result_set, paths = _unload_result_set( - tmp_path, tables, chunksize, write_statistics=write_statistics, **kwargs - ) + result_set, paths = _unload_result_set(tmp_path, tables, chunksize, **kwargs) with ( patch.object(AthenaPandasResultSet, "_read_data_manifest", return_value=paths), patch.object( @@ -955,7 +818,7 @@ def _read_unload_result(tmp_path, tables, chunksize, write_statistics=True, **kw return chunks, whole -def _read_csv_result(data, types, engine="c", chunksize=None, fs=None, **kwargs): +def _read_csv_result(data, types, engine="c", chunksize=None, **kwargs): """Read an in-memory CSV through the result set and the real pandas parser. ``types`` maps column names to Athena types, or lists (name, type) pairs for @@ -974,7 +837,7 @@ def _read_csv_result(data, types, engine="c", chunksize=None, fs=None, **kwargs) result_set._csv_stream = None result_set._block_size = None result_set._cache_type = None - result_set._fs = fs or MagicMock() + result_set._fs = MagicMock() # Binary columns are read through a text stream. result_set._fs.open.side_effect = lambda path, mode="rb", **_: ( io.StringIO(data, newline="") if "t" in mode else io.BytesIO(data.encode()) @@ -1016,6 +879,101 @@ def _read_csv_result(data, types, engine="c", chunksize=None, fs=None, **kwargs) ), options +class _BucketFileSystem(DirFileSystem): + """The local files under a directory, addressed by S3 paths, with the files it opened.""" + + def __init__(self, root): + super().__init__(path=str(root), fs=LocalFileSystem(), skip_instance_cache=True) + # The keyword arguments of each open() and the file that it returned. + self.opened = [] + + def _join(self, path): + if isinstance(path, str): + path = path.removeprefix("s3://") + elif isinstance(path, list): + path = [p.removeprefix("s3://") for p in path] + return super()._join(path) + + def open(self, path, *args, **kwargs): + file = super().open(path, *args, **kwargs) + self.opened.append((kwargs, file)) + return file + + +def _write_unload_tables(fs, tables): + """Write the Parquet files of an UNLOAD result under ``s3://bucket/unload/``. + + Returns: + The data manifest of the files. + """ + fs.makedirs("bucket/unload", exist_ok=True) + manifest = [f"s3://bucket/unload/{index}.parquet" for index in range(len(tables))] + for path, table in zip(manifest, tables, strict=True): + with fs.open(path, "wb") as file: + pq.write_table(table, file) + fs.opened.clear() + return manifest + + +def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", **kwargs): + """Create a result set of a succeeded query over local result files, without AWS. + + The CSV ``data`` is the result file ``s3://bucket/``; ``tables`` are the files + of an UNLOAD result, which the connection's S3 client lists in the data manifest. + ``fs`` holds the files and stands in for PyAthena's filesystem, ``types`` lists the + (name, Athena type) pairs of the GetQueryResults metadata, and ``kwargs`` are the + result set's options. + """ + unload = tables is not None + location = f"s3://bucket/{name}" + if unload: + manifest = _write_unload_tables(fs, tables) + length = 1 + else: + manifest = [] + fs.makedirs("bucket", exist_ok=True) + fs.pipe_file(location, data.encode()) + length = len(data.encode()) + 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": location}, + "Statistics": {"DataManifestLocation": "s3://bucket/manifest.csv"}, + } + } + ) + connection = MagicMock() + connection.s3_client.get_object.side_effect = lambda **_: { + "Body": io.BytesIO("\n".join(manifest).encode()) + } + connection.s3_client.meta.config.max_pool_connections = 2 + connection.s3_client.head_object.return_value = {"ContentLength": length} + connection.client.get_query_results.return_value = { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [ + {"Name": n, "Type": t, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} + for n, t in types + ] + }, + "Rows": [], + } + } + with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): + return AthenaPandasResultSet( + connection=connection, + converter=DefaultPandasTypeConverter(), + query_execution=query_execution, + arraysize=1000, + retry_config=RetryConfig(), + unload=unload, + **kwargs, + ) + + def _string_series(values, infer_string): """Express the expected string dtype for the active pandas option.""" return pd.Series(values, dtype="str" if infer_string else object) @@ -1190,6 +1148,15 @@ def _types_frame(infer_string, parse_time=True): [], id="unparsed_dates", ), + pytest.param( + '"v"\n"2024-01-01"\n\n', + {"v": "varchar"}, + {"parse_dates": ["v"]}, + # The string dtype applies again after parse_dates, and NULL stays missing. + lambda infer: pd.DataFrame({"v": _string_series(["2024-01-01", float("nan")], infer)}), + [], + id="parsed_dates", + ), pytest.param( "id \tint \t \nname \tstring \t \n", {"col_name": "varchar"}, @@ -1244,19 +1211,6 @@ def test_read_csv_with_pyarrow_matches_pandas( ) -@pytest.mark.parametrize("infer_string", [True, False]) -def test_read_csv_with_pyarrow_string_dtype_after_dates(infer_string): - # A string dtype applies again after parse_dates, as with pandas' PyArrow - # engine, and keeps NULL missing. - with pd.option_context("future.infer_string", infer_string): - df = _read_csv_with_pyarrow( - io.BytesIO(b'"v"\n"2024-01-01"\n\n'), - _pyarrow_read_csv_kwargs({"v": "varchar"}, parse_dates=["v"]), - ) - assert df["v"].tolist()[0] == "2024-01-01" - assert pd.isna(df["v"].tolist()[1]) - - def test_read_csv_with_pyarrow_ignores_unused_dtype_entries(): # As with pandas' PyArrow engine, a dtype entry for a column that is not in # the result is not validated. diff --git a/tests/pyathena/pandas/test_util.py b/tests/pyathena/pandas/test_util.py index 4564f3fcb..4fe9bc7d6 100644 --- a/tests/pyathena/pandas/test_util.py +++ b/tests/pyathena/pandas/test_util.py @@ -545,9 +545,8 @@ def test_to_sql_session_credentials_and_s3_config(monkeypatch, executor_class, c cursor = conn.cursor() cursor.execute(f"SELECT * FROM {table_name} ORDER BY col_int") assert cursor.fetchall() == [(1,), (2,)] - # The bucket resource, and with threads also the workers' resources. - expected = 1 if executor_class is SpawnProcessPoolExecutor else 3 - assert len(resources.call_args_list) == expected + # Every S3 resource of this process uses s3_config. + assert resources.call_args_list assert all(c.kwargs["config"].max_pool_connections == 37 for c in resources.call_args_list) @@ -610,7 +609,7 @@ def test_to_sql_with_multiple_partitions(cursor): compression="snappy", ) cursor.execute(f"SHOW PARTITIONS {table_name}") - assert sorted(cursor.fetchall()), [(f"col_int={i}/col_string=a",) for i in range(5)] + [ + assert sorted(cursor.fetchall()) == [(f"col_int={i}/col_string=a",) for i in range(5)] + [ (f"col_int={i}/col_string=b",) for i in range(5, 10) ] cursor.execute(f"SELECT COUNT(*) FROM {table_name}") From ecf6f4d064a81a4aa263da6741b956257f0c0e89 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:12:05 +0900 Subject: [PATCH 4/7] Check the pandas reads against explicit values and close whole binary reads The DDL and small-file pyarrow tests compare the result with explicit names, values, and dtypes instead of a read through the C engine of the same result set. A whole binary read closes its stream after the read and after a failed read, the duplicate-name binary case keeps a boolean column's own converter, the pyarrow engine options include na_values, and the CSV test helper no longer records the internal read options. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 33 +++++++++-- tests/pyathena/pandas/test_result_set.py | 74 +++++++++++------------- 2 files changed, 62 insertions(+), 45 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 632cb819c..7576c9f32 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -226,6 +226,25 @@ def test_binary_csv_stream_close(self, tmp_path, engine, close_method): if close_method != "execute": assert cursor.fetchone() is None + @pytest.mark.parametrize("engine", ["c", "python"]) + @pytest.mark.parametrize("fails", [False, True], ids=["read", "failed-read"]) + def test_binary_csv_stream_close_after_whole_read(self, tmp_path, engine, fails): + """The binary CSV stream of a whole read is closed once the read ends or fails. + + No AWS calls; the result file is a local file larger than pandas' read buffer. + """ + with _offline_cursor( + tmp_path, [("value", "varbinary"), ("padding", "varchar")], [("00", "x" * 4096)] * 100 + ) as (cursor, fs): + if fails: + with pytest.raises(OperationalError): + cursor.execute("SELECT 1", engine=engine, dtype={"padding": "int64"}) + else: + cursor.execute("SELECT 1", engine=engine) + assert cursor.fetchone() == (b"\x00", "x" * 4096) + assert fs.opened + assert all(file.closed for file in fs.opened) + @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize("read_method", ["next", "get_chunk"]) def test_binary_csv_stream_close_on_read_error(self, tmp_path, engine, read_method): @@ -264,11 +283,15 @@ def test_binary_converter_override(self, pandas_cursor): assert pandas_cursor.fetchone() == ("00 ff",) @pytest.mark.parametrize( - ("engine", "removed", "second_name"), - [("c", None, "padding"), ("pyarrow", None, "padding"), ("c", "varbinary", "value")], + ("engine", "removed", "second_column", "second_value"), + [ + ("c", None, "%(padding)s AS padding", "x" * 200), + ("pyarrow", None, "%(padding)s AS padding", "x" * 200), + ("c", "varbinary", "true AS value", True), + ], ids=["c", "pyarrow", "duplicate_name"], ) - def test_binary_without_converter(self, engine, removed, second_name): + def test_binary_without_converter(self, engine, removed, second_column, second_value): """Without a varbinary converter, binary values are the hex text of the result. The PyArrow engine reads the result only without converters; a column with the @@ -282,11 +305,11 @@ def test_binary_without_converter(self, engine, removed, second_name): conn.cursor() as cursor, ): cursor.execute( - f"SELECT X'00ff' AS value, %(padding)s AS {second_name}", + f"SELECT X'00ff' AS value, {second_column}", {"padding": "x" * 200}, engine=engine, ) - assert cursor.fetchone() == ("00 ff", "x" * 200) + assert cursor.fetchone() == ("00 ff", second_value) @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize( diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index c284a97e5..4e4ee3f31 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -304,36 +304,35 @@ def test_read_csv_ddl_preserves_numeric_looking_names( ): """A DDL result read with engine="pyarrow" keeps numeric-looking names and their padding. - It is read as with the C engine, and the result file is closed. + The result file is closed after the read. """ assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [(name, "string") for name in names] - frames = [] + fs = _BucketFileSystem(tmp_path) with pd.option_context("future.infer_string", infer_string): - for engine in ("pyarrow", "c"): - fs = _BucketFileSystem(tmp_path / engine) - frames.append( - _local_result_set( - fs, types, data, name="result.txt", engine=engine, **read_csv_kwargs - ).as_pandas() - ) - assert all(file.closed for _, file in fs.opened) - actual, expected = frames - assert_frame_equal(actual, expected, check_exact=True) - assert actual.iloc[:, 0].tolist() == expected_names + df = _local_result_set( + fs, types, data, name="result.txt", engine="pyarrow", **read_csv_kwargs + ).as_pandas() + assert all(file.closed for _, file in fs.opened) + assert df.columns.tolist() == names + assert df.iloc[:, 0].tolist() == expected_names def test_read_small_csv_with_pyarrow_engine(self, tmp_path): - """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine.""" + """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine. + + The C engine reads timestamps at microseconds. + """ data = '"d","t"\n"2024-01-01","2024-01-01 00:00:00.000"\n' assert len(data) < AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [("d", "date"), ("t", "timestamp")] - actual, expected = ( - _local_result_set( - _BucketFileSystem(tmp_path / engine), types, data, engine=engine - ).as_pandas() - for engine in ("pyarrow", "c") - ) - assert_frame_equal(actual, expected, check_exact=True) + df = _local_result_set( + _BucketFileSystem(tmp_path), types, data, engine="pyarrow" + ).as_pandas() + assert df.to_dict("list") == { + "d": [pd.Timestamp("2024-01-01")], + "t": [pd.Timestamp("2024-01-01 00:00:00")], + } + assert df["t"].dtype == "datetime64[us]" @pytest.mark.parametrize( ("kwargs", "expected"), @@ -351,6 +350,9 @@ def test_read_small_csv_with_pyarrow_engine(self, tmp_path): {"dtype": "str"}, ["NA", "null", "x", None, "y", "z", "w"], id="single_dtype" ), pytest.param({"parse_dates": ("d",)}, OperationalError, id="parse_dates_tuple"), + pytest.param( + {"na_values": ("", "NA")}, [None, "null", "x", None, "y", "z", "w"], id="na_values" + ), ], ) def test_read_csv_with_pyarrow_engine_options(self, tmp_path, kwargs, expected): @@ -550,7 +552,7 @@ def test_read_csv_json_null_without_dtype_warning( ): """JSON NULLs in a later parser block keep exact values without a warning.""" data = "n,j\n" + f"1,{json_value}\n" * 400000 + "2,\n" * 10 - result, _ = _read_csv_result(data, {"n": "integer", "j": "json"}, engine=engine, **kwargs) + result = _read_csv_result(data, {"n": "integer", "j": "json"}, engine=engine, **kwargs) if access != "whole": with result: first = result.get_chunk(400005) if access == "get_chunk" else next(result) @@ -570,7 +572,7 @@ def test_read_csv_json_null_without_dtype_warning( def test_read_csv_json_numbers_in_a_parser_block_without_null(self, kwargs): """JSON integers keep their type when their parser block has a float but no NULL.""" data = "n,j\n1,1.5\n" + "1,7\n" * 400000 + "2,\n" * 10 - df, _ = _read_csv_result(data, {"n": "integer", "j": "json"}, **kwargs) + df = _read_csv_result(data, {"n": "integer", "j": "json"}, **kwargs) values = df["j"] assert values.dtype == object assert values[0] == 1.5 @@ -592,7 +594,7 @@ def test_read_csv_json_numbers_in_a_parser_block_without_null(self, kwargs): ) def test_read_csv_json_without_null(self, engine, text, dtype, values): """JSON without NULL gets the dtype that pandas infers for a converter's values.""" - df, _ = _read_csv_result(f"n,j\n1,{text}\n", {"n": "integer", "j": "json"}, engine=engine) + df = _read_csv_result(f"n,j\n1,{text}\n", {"n": "integer", "j": "json"}, engine=engine) assert df["j"].dtype == dtype assert df["j"].tolist() == values @@ -611,7 +613,7 @@ def test_read_csv_json_without_null(self, engine, text, dtype, values): ) def test_read_csv_json_with_dtype(self, engine, dtype, n_dtype, m_dtype): """A dtype given to execute() replaces PyAthena's and does not apply to json columns.""" - df, _ = _read_csv_result( + df = _read_csv_result( "n,m,j\n1,1,2\n3,4,\n", {"n": "integer", "m": "integer", "j": "json"}, engine=engine, @@ -640,7 +642,7 @@ def test_read_csv_json_column_with_the_same_name_as_another_column( self, kwargs, expected, dtype ): """A column renamed from the name of a json column is read with its own options.""" - df, _ = _read_csv_result("x,x\n1,002\n3,004\n", [("x", "json"), ("x", "integer")], **kwargs) + df = _read_csv_result("x,x\n1,002\n3,004\n", [("x", "json"), ("x", "integer")], **kwargs) assert df.to_dict("list") == expected assert df["x.1"].dtype == dtype @@ -648,7 +650,7 @@ def test_read_csv_json_column_with_the_same_name_as_another_column( @pytest.mark.parametrize("kwargs", [{}, {"index_col": "j"}]) def test_read_csv_json_parse_dates(self, engine, kwargs): """Dates that pandas parses from a json column with parse_dates are kept.""" - df, _ = _read_csv_result( + df = _read_csv_result( "n,j\n1,20240101\n", {"n": "integer", "j": "json"}, engine=engine, @@ -661,7 +663,7 @@ def test_read_csv_json_parse_dates(self, engine, kwargs): @pytest.mark.parametrize("engine", ["c", "python"]) def test_read_csv_single_json_column_with_null(self, engine): """NULL in a result with only a json column, an empty line, is None.""" - df, _ = _read_csv_result("j\n1\n\n", {"j": "json"}, engine=engine) + df = _read_csv_result("j\n1\n\n", {"j": "json"}, engine=engine) assert df["j"].tolist() == [1, None] @pytest.mark.parametrize("chunksize", [None, 1]) @@ -672,7 +674,7 @@ def test_read_csv_json_conversion_error(self, chunksize): """ def read(): - result, _ = _read_csv_result( + result = _read_csv_result( "n,j\n1,{\n", {"n": "integer", "j": "json"}, chunksize=chunksize ) if chunksize: @@ -685,7 +687,7 @@ def read(): @pytest.mark.parametrize(("converters", "expected"), [({}, ["[1]"]), ({"j": len}, [3])]) def test_read_csv_json_custom_converters(self, converters, expected): """Converters given to execute() replace the json conversion.""" - df, _ = _read_csv_result( + df = _read_csv_result( 'n,j\n1,"[1]"\n', {"n": "integer", "j": "json"}, converters=converters ) assert df["j"].tolist() == expected @@ -864,19 +866,11 @@ def _read_csv_result(data, types, engine="c", chunksize=None, **kwargs): return_value=None, ), patch.object(AthenaPandasResultSet, "_get_content_length", return_value=len(data)), - patch("pandas.read_csv", wraps=pd.read_csv) as read_csv, - patch( - "pyathena.pandas.result_set._read_csv_with_pyarrow", - wraps=_read_csv_with_pyarrow, - ) as read_pyarrow, ): result = result_set._read_csv() - options = read_pyarrow.call_args.args[1] if read_pyarrow.called else read_csv.call_args.kwargs if isinstance(result, pd.DataFrame): - return result_set._finish_csv_frame(result), options - return PandasDataFrameIterator( - result, result_set._finish_csv_frame, result_set._csv_stream - ), options + return result_set._finish_csv_frame(result) + return PandasDataFrameIterator(result, result_set._finish_csv_frame, result_set._csv_stream) class _BucketFileSystem(DirFileSystem): From 3c543d73a0afdea7440d99eaf3d2d71a58a580bf Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:20:29 +0900 Subject: [PATCH 5/7] Share the local result files of the offline pandas tests The offline cursor and result-set tests read their local result files through one LocalBucketFileSystem and build the GetQueryResults metadata with query_results_response() from tests/pyathena/util.py. The CSV engine fallback test moves to the result-set tests, where the engine is chosen, and the auto-optimized chunk sizes run offline except one live case. test_csv_open_options keeps five cases, and a user's filesystem for UNLOAD results has its own test. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 199 +++++++---------------- tests/pyathena/pandas/test_result_set.py | 168 +++++++++++-------- tests/pyathena/util.py | 57 ++++++- 3 files changed, 220 insertions(+), 204 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 7576c9f32..ffb6e6fd4 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1,5 +1,6 @@ import contextlib import csv +import io import math import random import string @@ -7,13 +8,11 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from decimal import Decimal -from typing import Any from unittest.mock import MagicMock, patch import numpy as np import pandas as pd import pytest -from fsspec.implementations.local import LocalFileSystem from pyathena.error import DatabaseError, DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem @@ -31,7 +30,9 @@ CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, RESULT_FILE_VALUES_QUERY, + LocalBucketFileSystem, cached_file_systems, + query_results_response, ) @@ -41,34 +42,9 @@ def _pandas_converter_without_bigint_dtype(): return converter -class _ResultFileSystem(LocalFileSystem): - """A filesystem that serves one local file for every path and keeps the files it opens.""" - - def __init__(self, path: str, **kwargs: Any) -> None: - """Initialize the filesystem. - - Args: - path: The local file that every path opens. - **kwargs: Ignored filesystem options. - """ - super().__init__() - self.path = path - self.opened: list[Any] = [] - - def _open(self, path: str, mode: str = "rb", **kwargs: Any) -> Any: - """Open the local file instead of the given path. - - Args: - path: The path to open, which is ignored. - mode: The mode to open the file in. - **kwargs: Ignored open options. - - Returns: - The opened local file. - """ - file = super()._open(self.path, mode) - self.opened.append(file) - return file +# The columns and rows of a binary result file larger than pandas' read buffer. +_BINARY_COLUMNS = [("value", "varbinary"), ("padding", "varchar")] +_BINARY_ROWS = [("00", "x" * 4096)] * 100 @contextlib.contextmanager @@ -88,11 +64,14 @@ def _offline_cursor(tmp_path, columns, rows, **cursor_kwargs): Yields: The cursor and the filesystem that reads the result file. """ - path = tmp_path / "result.csv" - with path.open("w", newline="") as file: - writer = csv.writer(file, quoting=csv.QUOTE_ALL, lineterminator="\n") - writer.writerow([name for name, _ in columns]) - writer.writerows(rows) + text = io.StringIO(newline="") + writer = csv.writer(text, quoting=csv.QUOTE_ALL, lineterminator="\n") + writer.writerow([name for name, _ in columns]) + writer.writerows(rows) + data = text.getvalue().encode() + fs = LocalBucketFileSystem(tmp_path) + fs.makedirs("bucket", exist_ok=True) + fs.pipe_file("s3://bucket/result.csv", data) connection = MagicMock() connection.client.start_query_execution.return_value = {"QueryExecutionId": "query_id"} connection.client.get_query_execution.return_value = { @@ -104,19 +83,8 @@ def _offline_cursor(tmp_path, columns, rows, **cursor_kwargs): "ResultConfiguration": {"OutputLocation": "s3://bucket/result.csv"}, } } - connection.client.get_query_results.return_value = { - "ResultSet": { - "Rows": [{"Data": [{"VarCharValue": name} for name, _ in columns]}], - "ResultSetMetadata": { - "ColumnInfo": [ - {"Name": name, "Type": type_, "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"} - for name, type_ in columns - ] - }, - } - } - connection.s3_client.head_object.return_value = {"ContentLength": path.stat().st_size} - fs = _ResultFileSystem(str(path), skip_instance_cache=True) + connection.client.get_query_results.return_value = query_results_response(columns) + connection.s3_client.head_object.return_value = {"ContentLength": len(data)} with ( patch("pyathena.pandas.result_set.S3FileSystem", lambda **_: fs), PandasCursor( @@ -130,10 +98,6 @@ def _offline_cursor(tmp_path, columns, rows, **cursor_kwargs): yield cursor, fs -# A value that makes a result file larger than the smallest that the PyArrow engine reads. -_PADDING = "x" * 120 - - class _QuoteNoneDialect(csv.excel): """A CSV dialect that reads quotes as text.""" @@ -187,15 +151,15 @@ def test_binary_csv_stream_close(self, tmp_path, engine, close_method): """ with _offline_cursor( tmp_path, - [("value", "varbinary"), ("padding", "varchar")], - [("00", "x" * 4096)] * 100, + _BINARY_COLUMNS, + _BINARY_ROWS, chunksize=1, ) as (cursor, fs): cursor.execute("SELECT 1", engine=engine, nrows=2 if close_method == "nrows" else None) dataframe_iterator = cursor.as_pandas() chunks = cursor.iter_chunks() if close_method == "generator" else dataframe_iterator assert next(chunks)["value"].tolist() == [b"\x00"] - opened = list(fs.opened) + opened = [file for _, file in fs.opened] assert opened assert not any(file.closed for file in opened) @@ -233,9 +197,7 @@ def test_binary_csv_stream_close_after_whole_read(self, tmp_path, engine, fails) No AWS calls; the result file is a local file larger than pandas' read buffer. """ - with _offline_cursor( - tmp_path, [("value", "varbinary"), ("padding", "varchar")], [("00", "x" * 4096)] * 100 - ) as (cursor, fs): + with _offline_cursor(tmp_path, _BINARY_COLUMNS, _BINARY_ROWS) as (cursor, fs): if fails: with pytest.raises(OperationalError): cursor.execute("SELECT 1", engine=engine, dtype={"padding": "int64"}) @@ -243,7 +205,7 @@ def test_binary_csv_stream_close_after_whole_read(self, tmp_path, engine, fails) cursor.execute("SELECT 1", engine=engine) assert cursor.fetchone() == (b"\x00", "x" * 4096) assert fs.opened - assert all(file.closed for file in fs.opened) + assert all(file.closed for _, file in fs.opened) @pytest.mark.parametrize("engine", ["c", "python"]) @pytest.mark.parametrize("read_method", ["next", "get_chunk"]) @@ -255,18 +217,18 @@ def test_binary_csv_stream_close_on_read_error(self, tmp_path, engine, read_meth """ with _offline_cursor( tmp_path, - [("value", "varbinary"), ("padding", "varchar")], - [("00", "x" * 4096)] * 100, + _BINARY_COLUMNS, + _BINARY_ROWS, chunksize=1, ) as (cursor, fs): cursor.execute("SELECT 1", engine=engine, dtype={"padding": "int64"}) chunks = cursor.as_pandas() assert fs.opened - assert not any(file.closed for file in fs.opened) + assert not any(file.closed for _, file in fs.opened) read_chunk = chunks.__next__ if read_method == "next" else chunks.get_chunk with pytest.raises(ValueError): # noqa: PT011 read_chunk() - assert all(file.closed for file in fs.opened) + assert all(file.closed for _, file in fs.opened) for later in (lambda: next(chunks), chunks.get_chunk, cursor.fetchone): with pytest.raises(ProgrammingError): later() @@ -1212,54 +1174,6 @@ def test_no_ops(self): cursor.close() conn.close() - @pytest.mark.parametrize( - ("columns", "rows", "read_options", "expected", "sizes"), - [ - ( - [("a", "integer"), ("s", "varchar")], - [("1", _PADDING), ("2", _PADDING)], - {"chunksize": 1}, - [(1, _PADDING), (2, _PADDING)], - [1, 1], - ), - ( - [("s", "varchar")], - [(_PADDING,)], - {"quoting": csv.QUOTE_NONE}, - [(f'"{_PADDING}"',)], - [1], - ), - ( - [("x", "integer"), ("x", "varchar")], - [("1", f"007{_PADDING}")], - {}, - [(1, f"007{_PADDING}")], - [1], - ), - ( - [("b", "varbinary"), ("s", "varchar")], - [("00 ff", _PADDING)], - {}, - [(b"\x00\xff", _PADDING)], - [1], - ), - ], - ids=["chunksize", "quoting", "duplicate_names", "converters"], - ) - def test_pyarrow_engine_incompatible_options( - self, tmp_path, columns, rows, read_options, expected, sizes - ): - """engine="pyarrow" reads with the C engine what the PyArrow engine cannot read. - - That is a result in chunks, with another quoting, with columns of the same name, - or with converters. No AWS calls; the result file is a local file. - """ - with _offline_cursor(tmp_path, columns, rows) as (cursor, _): - cursor.execute("SELECT 1", engine="pyarrow", **read_options) - chunks = list(cursor.iter_chunks()) - assert [len(chunk) for chunk in chunks] == sizes - assert [tuple(row) for chunk in chunks for row in chunk.itertuples(index=False)] == expected - @pytest.mark.parametrize( ("pandas_cursor", "parquet_engine", "chunksize"), [ @@ -1619,43 +1533,52 @@ def test_unload_with_filesystem_options(self, pandas_cursor, option): assert df.to_dict("records") == [{"number_of_rows": 1}] @pytest.mark.parametrize( - ("pandas_cursor", "size_options", "read", "expected"), + "pandas_cursor", [{"cursor_kwargs": {"auto_optimize_chunksize": True}}], indirect=True + ) + def test_pandas_cursor_auto_optimize_chunksize_enabled(self, pandas_cursor, monkeypatch): + """auto_optimize_chunksize chunks a CSV result of Athena by its size.""" + # Make the five-row result exceed the size threshold, one byte per row. + monkeypatch.setattr(AthenaPandasResultSet, "LARGE_FILE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(AthenaPandasResultSet, "ESTIMATED_BYTES_PER_ROW", 1) + monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_THRESHOLD_MEDIUM", 0) + monkeypatch.setattr(AthenaPandasResultSet, "AUTO_CHUNK_SIZE_MEDIUM", 2) + pandas_cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) AS t(number)") + chunks = list(pandas_cursor.iter_chunks()) + assert [len(chunk) for chunk in chunks] == [2, 2, 1] + assert pd.concat(chunks)["number"].tolist() == [1, 2, 3, 4, 5] + + @pytest.mark.parametrize( + ("cursor_kwargs", "size_options", "read", "expected"), [ - ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, {}, "iter_chunks", [5]), + ({}, {}, "iter_chunks", [5]), ( - {"cursor_kwargs": {"auto_optimize_chunksize": True}}, - {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, - "iter_chunks", - [2, 2, 1], - ), - ( - {"cursor_kwargs": {"auto_optimize_chunksize": True}}, + {}, {"AUTO_CHUNK_THRESHOLD_LARGE": 0, "AUTO_CHUNK_SIZE_LARGE": 2}, "iter_chunks", [2, 2, 1], ), ( - {"cursor_kwargs": {"auto_optimize_chunksize": True}}, + {}, {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, "as_pandas", None, ), ( - {"cursor_kwargs": {"auto_optimize_chunksize": True, "chunksize": 3}}, + {"chunksize": 3}, {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, "iter_chunks", [3, 2], ), ], - indirect=["pandas_cursor"], - ids=["small", "medium", "large", "medium-as_pandas", "explicit_chunksize"], + ids=["small", "large", "medium-as_pandas", "explicit_chunksize"], ) - def test_pandas_cursor_auto_optimize_chunksize_enabled( - self, pandas_cursor, size_options, read, expected, monkeypatch + def test_auto_optimize_chunksize( + self, tmp_path, monkeypatch, cursor_kwargs, size_options, read, expected ): """auto_optimize_chunksize chunks a CSV result by its size; a chunksize wins. - as_pandas() joins the chunks into the whole result. + as_pandas() joins the chunks into the whole result. No AWS calls; the result + file is a local file. """ if size_options: # Make the five-row result exceed the size threshold, one byte per row. @@ -1663,15 +1586,19 @@ def test_pandas_cursor_auto_optimize_chunksize_enabled( monkeypatch.setattr(AthenaPandasResultSet, "ESTIMATED_BYTES_PER_ROW", 1) for name, value in size_options.items(): monkeypatch.setattr(AthenaPandasResultSet, name, value) - pandas_cursor.execute("SELECT number FROM (VALUES (1), (2), (3), (4), (5)) AS t(number)") - if read == "as_pandas": - df = pandas_cursor.as_pandas() - assert df["number"].tolist() == [1, 2, 3, 4, 5] - assert df.index.tolist() == [0, 1, 2, 3, 4] - else: - chunks = list(pandas_cursor.iter_chunks()) - assert [len(chunk) for chunk in chunks] == expected - assert pd.concat(chunks)["number"].tolist() == [1, 2, 3, 4, 5] + rows = [(str(number),) for number in range(1, 6)] + with _offline_cursor( + tmp_path, [("number", "integer")], rows, auto_optimize_chunksize=True, **cursor_kwargs + ) as (cursor, _): + cursor.execute("SELECT 1") + if read == "as_pandas": + df = cursor.as_pandas() + assert df["number"].tolist() == [1, 2, 3, 4, 5] + assert df.index.tolist() == [0, 1, 2, 3, 4] + else: + chunks = list(cursor.iter_chunks()) + assert [len(chunk) for chunk in chunks] == expected + assert pd.concat(chunks)["number"].tolist() == [1, 2, 3, 4, 5] def test_pandas_cursor_whole_result_reused(self, pandas_cursor): """Test that as_pandas() and iter_chunks() do not consume the fetched rows.""" diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 4e4ee3f31..910a132f6 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -17,7 +17,6 @@ import pyarrow as pa import pyarrow.parquet as pq import pytest -from fsspec.implementations.dirfs import DirFileSystem from fsspec.implementations.local import LocalFileSystem from pandas.testing import assert_frame_equal @@ -32,7 +31,17 @@ _read_csv_with_pyarrow, ) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on -from tests.pyathena.util import ThrottledOperation, assert_fetches_fail +from tests.pyathena.util import ( + LocalBucketFileSystem, + ThrottledOperation, + assert_fetches_fail, + query_results_response, +) + +# A value that makes a result file larger than the smallest that the PyArrow engine reads. +_PADDING = "x" * 120 +# A CSV result file of two rows. +_TEXT_CSV = '"v"\n"a"\n"b"\n' def _typed_values(rows): @@ -308,7 +317,7 @@ def test_read_csv_ddl_preserves_numeric_looking_names( """ assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [(name, "string") for name in names] - fs = _BucketFileSystem(tmp_path) + fs = LocalBucketFileSystem(tmp_path) with pd.option_context("future.infer_string", infer_string): df = _local_result_set( fs, types, data, name="result.txt", engine="pyarrow", **read_csv_kwargs @@ -317,6 +326,62 @@ def test_read_csv_ddl_preserves_numeric_looking_names( assert df.columns.tolist() == names assert df.iloc[:, 0].tolist() == expected_names + @pytest.mark.parametrize( + ("columns", "rows", "read_options", "expected", "sizes"), + [ + ( + [("a", "integer"), ("s", "varchar")], + [("1", _PADDING), ("2", _PADDING)], + {"chunksize": 1}, + [(1, _PADDING), (2, _PADDING)], + [1, 1], + ), + ( + [("s", "varchar")], + [(_PADDING,)], + {"quoting": csv.QUOTE_NONE}, + [(f'"{_PADDING}"',)], + [1], + ), + ( + [("x", "integer"), ("x", "varchar")], + [("1", f"007{_PADDING}")], + {}, + [(1, f"007{_PADDING}")], + [1], + ), + ( + [("b", "varbinary"), ("s", "varchar")], + [("00 ff", _PADDING)], + {}, + [(b"\x00\xff", _PADDING)], + [1], + ), + ], + ids=["chunksize", "quoting", "duplicate_names", "converters"], + ) + def test_pyarrow_engine_incompatible_options( + self, tmp_path, columns, rows, read_options, expected, sizes + ): + """engine="pyarrow" reads with the C engine what the PyArrow engine cannot read. + + That is a result in chunks, with another quoting, with columns of the same name, + or with converters. + """ + text = io.StringIO(newline="") + writer = csv.writer(text, quoting=csv.QUOTE_ALL, lineterminator="\n") + writer.writerows([[name for name, _ in columns], *rows]) + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + columns, + text.getvalue(), + engine="pyarrow", + **read_options, + ) + chunks = list(result_set.iter_chunks()) + assert [len(chunk) for chunk in chunks] == sizes + assert [tuple(row) for chunk in chunks for row in chunk.itertuples(index=False)] == expected + def test_read_small_csv_with_pyarrow_engine(self, tmp_path): """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine. @@ -326,7 +391,7 @@ def test_read_small_csv_with_pyarrow_engine(self, tmp_path): assert len(data) < AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [("d", "date"), ("t", "timestamp")] df = _local_result_set( - _BucketFileSystem(tmp_path), types, data, engine="pyarrow" + LocalBucketFileSystem(tmp_path), types, data, engine="pyarrow" ).as_pandas() assert df.to_dict("list") == { "d": [pd.Timestamp("2024-01-01")], @@ -370,7 +435,7 @@ def test_read_csv_with_pyarrow_engine_options(self, tmp_path, kwargs, expected): data = '"v","d"\n' + "".join(f"{line}\n" for line in lines) assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES types = [("v", "varchar"), ("d", "date")] - fs = _BucketFileSystem(tmp_path) + fs = LocalBucketFileSystem(tmp_path) if expected is OperationalError: with pytest.raises(OperationalError): _local_result_set(fs, types, data, engine="pyarrow", **kwargs) @@ -459,7 +524,7 @@ def test_unload_chunks_with_read_options(self, tmp_path, kwargs, sizes, columns) pa.table({"a": [1, 2], "s": [{"x": 1}, {"x": 2}]}), pa.table({"a": [3], "s": [{"x": 3}]}), ] - fs = _BucketFileSystem(tmp_path) + fs = LocalBucketFileSystem(tmp_path) if sizes is None: with pytest.raises(OperationalError): _local_result_set(fs, tables=tables, chunksize=1, **kwargs) @@ -468,29 +533,26 @@ def test_unload_chunks_with_read_options(self, tmp_path, kwargs, sizes, columns) assert [len(chunk) for chunk in chunks] == sizes assert [chunk.columns.tolist() for chunk in chunks] == [columns] * len(sizes) + def test_unload_user_filesystem(self, tmp_path): + """A filesystem given to execute() reads the UNLOAD result instead of PyAthena's.""" + fs = LocalBucketFileSystem(tmp_path / "pyathena") + user_fs = LocalBucketFileSystem(tmp_path / "user") + _write_unload_tables(user_fs, [pa.table({"a": [2]})]) + result_set = _local_result_set(fs, tables=[pa.table({"a": [1]})], filesystem=user_fs) + assert result_set.fetchall() == [(2,)] + @pytest.mark.parametrize( ("option", "value"), - [ - ("filesystem", "user"), - ("storage_options", {}), - ("storage_options", None), - ("filesystem", None), - ], - ids=["filesystem", "storage_options", "storage_options-none", "filesystem-none"], + [("storage_options", {}), ("storage_options", None), ("filesystem", None)], + ids=["storage_options", "storage_options-none", "filesystem-none"], ) def test_unload_filesystem_options(self, tmp_path, option, value): - """filesystem or storage_options given to execute(), even None, replace PyAthena's. + """storage_options, even None, or a None filesystem given to execute() replace PyAthena's. - A given filesystem reads its own files. Otherwise pandas opens the S3 files - itself, so the read is observed. + pandas opens the S3 files itself, so the read is observed. """ - fs = _BucketFileSystem(tmp_path / "pyathena") + fs = LocalBucketFileSystem(tmp_path / "pyathena") tables = [pa.table({"a": [1]})] - if value == "user": - value = _BucketFileSystem(tmp_path / "user") - _write_unload_tables(value, [pa.table({"a": [2]})]) - assert _local_result_set(fs, tables=tables, filesystem=value).fetchall() == [(2,)] - return with patch("pandas.read_parquet", return_value=pd.DataFrame({"a": [2]})) as read_parquet: assert _local_result_set(fs, tables=tables, **{option: value}).fetchall() == [(2,)] # The S3 location, which pandas opens with the given options instead of PyAthena's. @@ -499,19 +561,22 @@ def test_unload_filesystem_options(self, tmp_path, option, value): assert read_parquet.call_args.kwargs.get("filesystem") is None @pytest.mark.parametrize( - ("data", "types", "chunksize"), - [ - pytest.param('"v"\n"a"\n"b"\n', [("v", "varchar")], None, id="text"), - pytest.param('"v"\n"a"\n"b"\n', [("v", "varchar")], 1, id="text-chunks"), - pytest.param('"b"\n"00 ff"\n', [("b", "varbinary")], None, id="binary"), - ], - ) - @pytest.mark.parametrize( - ("kwargs", "prefetch"), + ("data", "types", "chunksize", "kwargs", "prefetch"), [ - pytest.param({}, True, id="default"), - pytest.param({"block_size": 5 * 2**20}, False, id="block_size"), - pytest.param({"cache_type": "bytes"}, False, id="cache_type"), + pytest.param(_TEXT_CSV, [("v", "varchar")], None, {}, True, id="text"), + pytest.param(_TEXT_CSV, [("v", "varchar")], 1, {}, True, id="text-chunks"), + pytest.param('"b"\n"00 ff"\n', [("b", "varbinary")], None, {}, True, id="binary"), + pytest.param( + _TEXT_CSV, + [("v", "varchar")], + None, + {"block_size": 5 * 2**20}, + False, + id="block_size", + ), + pytest.param( + _TEXT_CSV, [("v", "varchar")], None, {"cache_type": "bytes"}, False, id="cache_type" + ), ], ) def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefetch): @@ -519,7 +584,7 @@ def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefet With block_size or cache_type, the file is opened with the filesystem's settings. """ - fs = _BucketFileSystem(tmp_path) + fs = LocalBucketFileSystem(tmp_path) result_set = _local_result_set(fs, types, data, chunksize=chunksize, **kwargs) assert len(result_set.fetchall()) == len(data.splitlines()) - 1 [(options, _)] = fs.opened @@ -873,27 +938,6 @@ def _read_csv_result(data, types, engine="c", chunksize=None, **kwargs): return PandasDataFrameIterator(result, result_set._finish_csv_frame, result_set._csv_stream) -class _BucketFileSystem(DirFileSystem): - """The local files under a directory, addressed by S3 paths, with the files it opened.""" - - def __init__(self, root): - super().__init__(path=str(root), fs=LocalFileSystem(), skip_instance_cache=True) - # The keyword arguments of each open() and the file that it returned. - self.opened = [] - - def _join(self, path): - if isinstance(path, str): - path = path.removeprefix("s3://") - elif isinstance(path, list): - path = [p.removeprefix("s3://") for p in path] - return super()._join(path) - - def open(self, path, *args, **kwargs): - file = super().open(path, *args, **kwargs) - self.opened.append((kwargs, file)) - return file - - def _write_unload_tables(fs, tables): """Write the Parquet files of an UNLOAD result under ``s3://bucket/unload/``. @@ -945,17 +989,7 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * } connection.s3_client.meta.config.max_pool_connections = 2 connection.s3_client.head_object.return_value = {"ContentLength": length} - connection.client.get_query_results.return_value = { - "ResultSet": { - "ResultSetMetadata": { - "ColumnInfo": [ - {"Name": n, "Type": t, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} - for n, t in types - ] - }, - "Rows": [], - } - } + connection.client.get_query_results.return_value = query_results_response(types) with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): return AthenaPandasResultSet( connection=connection, diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 8d3a08412..2b827ab8a 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -19,6 +19,8 @@ from botocore.config import Config from botocore.exceptions import ClientError from botocore.response import StreamingBody +from fsspec.implementations.dirfs import DirFileSystem +from fsspec.implementations.local import LocalFileSystem from jinja2 import Environment, FileSystemLoader from sqlalchemy import types @@ -28,7 +30,7 @@ from pyathena.filesystem.s3_core import S3Core from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution -from pyathena.util import retry_api_call +from pyathena.util import override, retry_api_call _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -637,3 +639,56 @@ def stub_multipart_copy(stubber, fail_list=False, fail_part=False, fail_annotati ) return stubber.add_response("put_object_annotation", {}, put) + + +def query_results_response(columns): + """Build a ``GetQueryResults`` response with the metadata of the given columns. + + Args: + columns: The (name, Athena type) pairs of the result columns. + + Returns: + The response, without rows, as for a result that is read from its result file. + """ + return { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [ + {"Name": name, "Type": type_, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} + for name, type_ in columns + ] + }, + "Rows": [], + } + } + + +class LocalBucketFileSystem(DirFileSystem): + """The local files under a directory, addressed by S3 paths, with the files it opened. + + ``s3://bucket/key`` is the file ``bucket/key`` under the directory. + """ + + def __init__(self, root): + """Initialize the filesystem. + + Args: + root: The directory that holds the files. + """ + super().__init__(path=str(root), fs=LocalFileSystem(), skip_instance_cache=True) + # The keyword arguments of each open() and the file that it returned. + self.opened = [] + + @override + def _join(self, path): + if isinstance(path, str): + path = path.removeprefix("s3://") + elif isinstance(path, list): + path = [p.removeprefix("s3://") for p in path] + return super()._join(path) + + @override + def open(self, path, *args, **kwargs): + file = super().open(path, *args, **kwargs) + self.opened.append((kwargs, file)) + return file From efbdb7c5057c825f25ed26c473665a2e91d26436 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:38:33 +0900 Subject: [PATCH 6/7] Check the pandas read options that the independent review found untested The UNLOAD chunk test checks the dtypes of each dtype_backend and an empty column list, which reads the result whole, and a repeated column is fetched twice. The CSV open-options test checks the documented 32 MiB blocks and two cached blocks and the options that PyAthena's filesystem is created with. The storage-options test reads a binary NULL and an empty binary value, and engine="pyarrow" is checked through the values of numeric-looking, multiline, and NULL varchar values. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 8 +++- tests/pyathena/pandas/test_result_set.py | 52 +++++++++++++++++++++--- 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index ffb6e6fd4..32f073aca 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -357,7 +357,13 @@ def test_result_sets_share_s3_client(self, pandas_cursor): @pytest.mark.parametrize( ("query", "expected"), - [("SELECT * FROM one_row", [(1,)]), ("SELECT X'01' AS value", [(b"\x01",)])], + [ + ("SELECT * FROM one_row", [(1,)]), + ( + "SELECT X'01' AS value, CAST(NULL AS VARBINARY) AS null_value, X'' AS empty_value", + [(b"\x01", None, b"")], + ), + ], ids=["plain", "binary"], ) def test_csv_storage_options(self, pandas_cursor, query, expected): diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 910a132f6..3b59e960d 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -382,6 +382,24 @@ def test_pyarrow_engine_incompatible_options( assert [len(chunk) for chunk in chunks] == sizes assert [tuple(row) for chunk in chunks for row in chunk.itertuples(index=False)] == expected + def test_read_csv_with_pyarrow_engine_values(self, tmp_path): + """engine="pyarrow" keeps the text of varchar values, numeric-looking and multiline.""" + values = ["007", "1e3", "nan", "line1\nline2", None, _PADDING] + data = '"v","n"\n' + "".join( + f'{"" if v is None else chr(34) + v + chr(34)},"{i}"\n' for i, v in enumerate(values) + ) + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("v", "varchar"), ("n", "integer")], + data, + engine="pyarrow", + ) + rows = result_set.fetchall() + assert [(None if pd.isna(v) else v, n) for v, n in rows] == [ + (v, i) for i, v in enumerate(values) + ] + def test_read_small_csv_with_pyarrow_engine(self, tmp_path): """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine. @@ -511,6 +529,7 @@ def test_as_pandas_parquet_chunks_engine(self, tmp_path): pytest.param({}, [1, 1, 1], ["a", "s"], id="default"), pytest.param({"columns": ["a"]}, [1, 1, 1], ["a"], id="columns"), pytest.param({"dtype_backend": "pyarrow"}, [1, 1, 1], ["a", "s"], id="dtype_backend"), + pytest.param({"columns": []}, [3], [], id="no_columns"), # Other options read the result whole, where pandas applies or rejects them. pytest.param({"filters": [("a", "=", 1)]}, [1], ["a", "s"], id="filters"), pytest.param({"columns": ["s.x"]}, [3], ["x"], id="nested_column"), @@ -532,6 +551,22 @@ def test_unload_chunks_with_read_options(self, tmp_path, kwargs, sizes, columns) chunks = list(_local_result_set(fs, tables=tables, chunksize=1, **kwargs).iter_chunks()) assert [len(chunk) for chunk in chunks] == sizes assert [chunk.columns.tolist() for chunk in chunks] == [columns] * len(sizes) + # The chunks have the dtypes of the backend. + backend = kwargs.get("dtype_backend") + assert all( + isinstance(dtype, pd.ArrowDtype) is (backend == "pyarrow") + for chunk in chunks + for dtype in chunk.dtypes + ) + + def test_unload_repeated_columns(self, tmp_path): + """Rows of an UNLOAD result read with a column repeated have the column twice.""" + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + tables=[pa.table({"a": [1, 2], "b": ["x", "y"]})], + columns=["a", "a"], + ) + assert result_set.fetchall() == [(1, 1), (2, 2)] def test_unload_user_filesystem(self, tmp_path): """A filesystem given to execute() reads the UNLOAD result instead of PyAthena's.""" @@ -580,9 +615,10 @@ def test_unload_filesystem_options(self, tmp_path, option, value): ], ) def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefetch): - """A CSV result is read in CSV_BLOCK_SIZE blocks with the background cache. + """A CSV result is read in 32 MiB blocks with the background cache, keeping two blocks. - With block_size or cache_type, the file is opened with the filesystem's settings. + With block_size or cache_type, the file is opened with the filesystem's settings, + which are created with them. """ fs = LocalBucketFileSystem(tmp_path) result_set = _local_result_set(fs, types, data, chunksize=chunksize, **kwargs) @@ -590,10 +626,13 @@ def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefet [(options, _)] = fs.opened if prefetch: assert options["cache_type"] == "background" - assert options["block_size"] == AthenaPandasResultSet.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 + assert fs.created_with["default_block_size"] == kwargs.get("block_size") + assert fs.created_with["default_cache_type"] == kwargs.get("cache_type") @pytest.mark.filterwarnings("error::pandas.errors.DtypeWarning") @pytest.mark.parametrize( @@ -990,8 +1029,8 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * connection.s3_client.meta.config.max_pool_connections = 2 connection.s3_client.head_object.return_value = {"ContentLength": length} connection.client.get_query_results.return_value = query_results_response(types) - with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): - return AthenaPandasResultSet( + with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs) as file_system: + result_set = AthenaPandasResultSet( connection=connection, converter=DefaultPandasTypeConverter(), query_execution=query_execution, @@ -1000,6 +1039,9 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * unload=unload, **kwargs, ) + # The options that PyAthena's filesystem was created with. + fs.created_with = file_system.call_args.kwargs + return result_set def _string_series(values, infer_string): From 102e971229712e53fbcb10d7a665fa7be04540c8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:46:06 +0900 Subject: [PATCH 7/7] Check the pyarrow CSV dispatch with numeric-looking text The engine="pyarrow" values test reads a varchar column of only numeric-looking values, which pandas' own PyArrow reader would read as numbers, and keeps the padding in another column. The CSV open-options test no longer records the arguments of the mocked filesystem. Refs #1244 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_result_set.py | 34 +++++++++++------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 3b59e960d..0f9b5038e 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -382,23 +382,25 @@ def test_pyarrow_engine_incompatible_options( assert [len(chunk) for chunk in chunks] == sizes assert [tuple(row) for chunk in chunks for row in chunk.itertuples(index=False)] == expected - def test_read_csv_with_pyarrow_engine_values(self, tmp_path): + @pytest.mark.parametrize( + "values", + [ + pytest.param(["007", "1e3", "10", None], id="numeric_looking"), + pytest.param(["a", "line1\nline2", "nan", None], id="text"), + ], + ) + def test_read_csv_with_pyarrow_engine_values(self, tmp_path, values): """engine="pyarrow" keeps the text of varchar values, numeric-looking and multiline.""" - values = ["007", "1e3", "nan", "line1\nline2", None, _PADDING] - data = '"v","n"\n' + "".join( - f'{"" if v is None else chr(34) + v + chr(34)},"{i}"\n' for i, v in enumerate(values) + data = '"v","p"\n' + "".join( + f'{"" if v is None else chr(34) + v + chr(34)},"{_PADDING}"\n' for v in values ) - assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES result_set = _local_result_set( LocalBucketFileSystem(tmp_path), - [("v", "varchar"), ("n", "integer")], + [("v", "varchar"), ("p", "varchar")], data, engine="pyarrow", ) - rows = result_set.fetchall() - assert [(None if pd.isna(v) else v, n) for v, n in rows] == [ - (v, i) for i, v in enumerate(values) - ] + assert [None if pd.isna(v) else v for v, _ in result_set.fetchall()] == values def test_read_small_csv_with_pyarrow_engine(self, tmp_path): """A CSV result smaller than PYARROW_MIN_FILE_SIZE_BYTES is read with the C engine. @@ -617,8 +619,7 @@ def test_unload_filesystem_options(self, tmp_path, option, value): def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefetch): """A CSV result is read in 32 MiB blocks with the background cache, keeping two blocks. - With block_size or cache_type, the file is opened with the filesystem's settings, - which are created with them. + With block_size or cache_type, the file is opened with the filesystem's settings. """ fs = LocalBucketFileSystem(tmp_path) result_set = _local_result_set(fs, types, data, chunksize=chunksize, **kwargs) @@ -631,8 +632,6 @@ def test_csv_open_options(self, tmp_path, data, types, chunksize, kwargs, prefet else: assert "cache_type" not in options assert "block_size" not in options - assert fs.created_with["default_block_size"] == kwargs.get("block_size") - assert fs.created_with["default_cache_type"] == kwargs.get("cache_type") @pytest.mark.filterwarnings("error::pandas.errors.DtypeWarning") @pytest.mark.parametrize( @@ -1029,8 +1028,8 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * connection.s3_client.meta.config.max_pool_connections = 2 connection.s3_client.head_object.return_value = {"ContentLength": length} connection.client.get_query_results.return_value = query_results_response(types) - with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs) as file_system: - result_set = AthenaPandasResultSet( + with patch("pyathena.pandas.result_set.S3FileSystem", return_value=fs): + return AthenaPandasResultSet( connection=connection, converter=DefaultPandasTypeConverter(), query_execution=query_execution, @@ -1039,9 +1038,6 @@ def _local_result_set(fs, types=(), data=None, tables=None, name="result.csv", * unload=unload, **kwargs, ) - # The options that PyAthena's filesystem was created with. - fs.created_with = file_system.call_args.kwargs - return result_set def _string_series(values, infer_string):