diff --git a/docs/pandas.md b/docs/pandas.md index 93f36a372..eb0d46ed0 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -556,9 +556,10 @@ Common performance options: - `dtype`: Explicit column data types - `parse_dates`: Columns to parse as dates -With `engine="pyarrow"`, PandasCursor uses the PyArrow engine only when pyarrow is installed, no chunksize is set (explicitly or by `auto_optimize_chunksize`), `quoting` is the default, the result has no columns that need a converter (`boolean`, `decimal`, `varbinary`, `json`, `time with time zone`, and `timestamp with time zone` with the default converter), and the result file is at least `AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES` bytes. +With `engine="pyarrow"`, PandasCursor uses the PyArrow engine only when the result is not a tab-separated `.txt` file, pyarrow is installed, no chunksize is set (explicitly or by `auto_optimize_chunksize`), `quoting` is the default, the result has no columns that need a converter (`boolean`, `decimal`, `varbinary`, `json`, `time with time zone`, and `timestamp with time zone` with the default converter), and the result file is at least `AthenaPandasResultSet.PYARROW_MIN_FILE_SIZE_BYTES` bytes. Otherwise, it falls back to the C engine. -With the PyArrow engine, when `keep_default_na` and `na_values` are the defaults and the only pandas.read_csv() options passed to `execute()` are `dtype` as a mapping and `parse_dates` as a list, PyAthena reads the file with `pyarrow.csv`, allowing newlines in quoted values, and converts it to a DataFrame as `pandas.read_csv(engine="pyarrow")` does, except that columns with a string dtype are read as text, other than in the tab-separated results of DDL statements, and without `future.infer_string`, a column with dtype `str` keeps NULL as a missing value. +With `engine="pyarrow"`, tab-separated `.txt` results from DDL statements such as `SHOW TABLES`, `SHOW COLUMNS`, and `DESCRIBE` use the C engine to preserve leading zeros, exponent notation, and padding in string values. +With the PyArrow engine, when `keep_default_na` and `na_values` are the defaults and the only pandas.read_csv() options passed to `execute()` are `dtype` as a mapping and `parse_dates` as a list, PyAthena reads the file with `pyarrow.csv`, allowing newlines in quoted values, and converts it to a DataFrame as `pandas.read_csv(engine="pyarrow")` does, except that columns with a string dtype are read as text, and without `future.infer_string`, a column with dtype `str` keeps NULL as a missing value. Otherwise, pandas reads the file, and its PyArrow engine raises an error or returns wrong values when a quoted value containing a newline crosses one of pyarrow's read blocks. It also changes the values of a column with a string dtype whose values all look like numbers, such as `"007"` to `"7.0"`, and without `future.infer_string`, turns NULL in a column with dtype `str` into the string `"nan"`. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 3f21f8a5e..f632b9825 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -637,6 +637,10 @@ def _get_csv_engine( # Use PyArrow only when explicitly requested and all compatibility # checks pass; otherwise fall through to the C engine default. if self._engine == "pyarrow": + # Header-less DDL results lose numeric-looking strings and their padding + # when PyArrow infers types before applying the string dtypes. + if self.output_location and self.output_location.endswith(".txt"): + return "c" effective_chunksize = chunksize if chunksize is not None else self._chunksize is_compatible = ( effective_chunksize is None diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 26e701542..145bdd4fc 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1040,6 +1040,7 @@ def test_get_csv_engine_explicit_specification(self): 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 diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 0ed38d817..8d7ceead1 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -81,6 +81,105 @@ def test_as_pandas_single_dataframe(self): class TestAthenaPandasResultSet: + @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( + ("data", "names", "expected_names"), + [ + ("001\n007\n1e3\n" * 10, ["tab_name"], ["001", "007", "1e3"] * 10), + ( + "001 \n007 \n1e3 \n" * 10, + ["field"], + ["001 ", "007 ", "1e3 "] * 10, + ), + ( + ( + "001 \tint \t\n" + "007 \tint \t\n" + "1e3 \tint \t\n" + ) + * 10, + ["col_name", "data_type", "comment"], + ["001 ", "007 ", "1e3 "] * 10, + ), + ], + 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 + ): + """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._fs = MagicMock() + result_set._fs.open.return_value = stream = io.BytesIO(data) + description = [(name, "string", None, None, 0, 0, "UNKNOWN") for name in names] + + with ( + pd.option_context("future.infer_string", infer_string), + patch.object( + AthenaPandasResultSet, + "description", + new_callable=PropertyMock, + return_value=description, + ), + patch.object(result_set, "_get_content_length", return_value=len(data)), + patch("pandas.read_csv", wraps=pd.read_csv) as read_csv, + ): + expected = pd.read_csv(io.BytesIO(data), **result_set._get_csv_read_options("c", None)) + read_csv.reset_mock() + actual = result_set._read_csv() + + 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" + assert stream.closed + @pytest.mark.parametrize( ("execute_kwargs", "path", "filesystem_kwargs"), [