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]] 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_cursor.py b/tests/pyathena/pandas/test_cursor.py index 1058386de..32f073aca 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,23 +8,19 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from decimal import Decimal -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import MagicMock, patch import numpy as np import pandas as pd import pytest -from pandas.io.parsers import TextFileReader 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 @@ -33,7 +30,9 @@ CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, RESULT_FILE_VALUES_QUERY, + LocalBucketFileSystem, cached_file_systems, + query_results_response, ) @@ -43,6 +42,68 @@ def _pandas_converter_without_bigint_dtype(): return converter +# 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 +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. + """ + 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 = { + "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 = query_results_response(columns) + connection.s3_client.head_object.return_value = {"ContentLength": len(data)} + 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 + + +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 +128,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 +138,100 @@ 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, + _BINARY_COLUMNS, + _BINARY_ROWS, 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 = [file for _, file in 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("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, _BINARY_COLUMNS, _BINARY_ROWS) 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, 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, + _BINARY_COLUMNS, + _BINARY_ROWS, 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,21 +244,34 @@ 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_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_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 + 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, {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( @@ -211,33 +292,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 +318,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 +330,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 +356,42 @@ 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"), + ("query", "expected"), [ - ("SELECT * FROM one_row", [(1,)], False), - ("SELECT X'01' AS value", [(b"\x01",)], True), + ("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"], ) - @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 +459,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 +474,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 +488,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 +666,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 +783,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 +802,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 +949,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 +986,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 +1038,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 +1180,11 @@ 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"), - [ - (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", - ], - ) - 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 - @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 +1198,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 +1223,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 +1258,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 +1287,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 +1305,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 +1323,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 +1337,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 +1353,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 +1370,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 +1411,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 +1538,73 @@ 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", [{"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( - ("pandas_cursor", "chunked"), + ("cursor_kwargs", "size_options", "read", "expected"), [ - ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, False), - ({"cursor_kwargs": {"auto_optimize_chunksize": True}}, True), + ({}, {}, "iter_chunks", [5]), + ( + {}, + {"AUTO_CHUNK_THRESHOLD_LARGE": 0, "AUTO_CHUNK_SIZE_LARGE": 2}, + "iter_chunks", + [2, 2, 1], + ), + ( + {}, + {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, + "as_pandas", + None, + ), + ( + {"chunksize": 3}, + {"AUTO_CHUNK_THRESHOLD_MEDIUM": 0, "AUTO_CHUNK_SIZE_MEDIUM": 2}, + "iter_chunks", + [3, 2], + ), ], - indirect=["pandas_cursor"], + ids=["small", "large", "medium-as_pandas", "explicit_chunksize"], ) - def test_pandas_cursor_auto_optimize_chunksize_enabled( - self, pandas_cursor, chunked, monkeypatch + def test_auto_optimize_chunksize( + self, tmp_path, monkeypatch, cursor_kwargs, size_options, read, expected ): - """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. 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. 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)") - 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 + for name, value in size_options.items(): + monkeypatch.setattr(AthenaPandasResultSet, name, value) + 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.""" @@ -1863,100 +1620,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 +1897,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 +1928,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] diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index e1f3b47de..0f9b5038e 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -10,7 +10,7 @@ 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 @@ -27,12 +27,21 @@ from pyathena.pandas.result_set import ( AthenaPandasResultSet, PandasDataFrameIterator, - _JSONConverter, _no_trunc_date, _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): @@ -94,6 +103,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 +118,7 @@ def test_iterrows_continuous_row_numbers(self): ], ids=[ "default_index", + "single_chunk", "index_col", "category", "category_null_chunk", @@ -134,13 +145,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 +171,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 +205,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 +284,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 +309,163 @@ 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] + """A DDL result read with engine="pyarrow" keeps numeric-looking names and their padding. - 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() - - 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 + The result file is closed after the read. + """ + assert len(data) >= AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES + types = [(name, "string") for name in names] + 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 + ).as_pandas() + assert all(file.closed for _, file in fs.opened) + assert df.columns.tolist() == names + assert df.iloc[:, 0].tolist() == expected_names @pytest.mark.parametrize( - ("chunksize", "execute_kwargs", "expected"), + ("columns", "rows", "read_options", "expected", "sizes"), [ - (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), + ( + [("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_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_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 @pytest.mark.parametrize( - ("block_size", "cache_type", "expected"), + "values", [ - ( - None, - None, - { - "block_size": 32 * 2**20, - "cache_type": "background", - "cache_options": {"maxblocks": 2}, - }, - ), - (5 * 2**20, None, {}), - (None, "bytes", {}), - (5 * 2**20, "readahead", {}), + pytest.param(["007", "1e3", "10", None], id="numeric_looking"), + pytest.param(["a", "line1\nline2", "nan", None], id="text"), ], ) - 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 + def test_read_csv_with_pyarrow_engine_values(self, tmp_path, values): + """engine="pyarrow" keeps the text of varchar values, numeric-looking and multiline.""" + data = '"v","p"\n' + "".join( + f'{"" if v is None else chr(34) + v + chr(34)},"{_PADDING}"\n' for v in values + ) + result_set = _local_result_set( + LocalBucketFileSystem(tmp_path), + [("v", "varchar"), ("p", "varchar")], + data, + engine="pyarrow", + ) + 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. + + 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")] + df = _local_result_set( + LocalBucketFileSystem(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( - ("execute_kwargs", "path", "filesystem_kwargs"), + ("kwargs", "expected"), [ - ({}, "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({}, ["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", + ), + 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"), + pytest.param( + {"na_values": ("", "NA")}, [None, "null", "x", None, "y", "z", "w"], id="na_values" + ), ], - 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 = LocalBucketFileSystem(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 +483,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 +525,113 @@ 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"), + 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"), + 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 = LocalBucketFileSystem(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) + # 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.""" + 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"), + [("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): + """storage_options, even None, or a None filesystem given to execute() replace PyAthena's. + + pandas opens the S3 files itself, so the read is observed. + """ + fs = LocalBucketFileSystem(tmp_path / "pyathena") + tables = [pa.table({"a": [1]})] + 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", "kwargs", "prefetch"), [ - 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(_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" + ), ], ) - @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" + 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. + """ + 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 + if prefetch: + assert options["cache_type"] == "background" + 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 @pytest.mark.filterwarnings("error::pandas.errors.DtypeWarning") @pytest.mark.parametrize( @@ -625,8 +641,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 +655,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 +664,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) @@ -666,7 +675,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 @@ -676,7 +685,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 +695,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"), @@ -775,7 +745,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 @@ -783,7 +753,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, @@ -796,34 +766,34 @@ 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] - 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 +871,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 +881,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 +896,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 +906,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 +923,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 +942,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()) @@ -1001,19 +969,75 @@ def _read_csv_result(data, types, engine="c", chunksize=None, fs=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) + + +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 = query_results_response(types) + 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): @@ -1190,6 +1214,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 +1277,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}") 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