diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index ea11a419f..77978264a 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -146,6 +146,9 @@ def __init__( import pyarrow as pa self._table = pa.Table.from_pydict({}) + # The fetch methods convert only the values read from a result file. + # GetQueryResults values are already converted. + self._convert_rows = bool(self.output_location) self._batches = iter(self._table.to_batches(arraysize)) def _create_s3_file_system(self): @@ -251,11 +254,15 @@ def _fetch(self) -> None: return else: dict_rows = rows.to_pydict() - column_names = dict_rows.keys() - processed_rows = [ - tuple(self.converters[k](v) for k, v in zip(column_names, row, strict=False)) - for row in zip(*dict_rows.values(), strict=False) - ] + if self._convert_rows: + converters = self.converters + column_names = dict_rows.keys() + processed_rows = [ + tuple(converters[k](v) for k, v in zip(column_names, row, strict=False)) + for row in zip(*dict_rows.values(), strict=False) + ] + else: + processed_rows = list(zip(*dict_rows.values(), strict=False)) self._rows.extend(processed_rows) @override diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index af1abca69..f70e99f65 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -969,5 +969,16 @@ def test_null_vs_empty_string(self, arrow_cursor): indirect=["arrow_cursor"], ) def test_fetch_all_rows(self, arrow_cursor): - arrow_cursor.execute("SELECT 1 AS col") - assert arrow_cursor.fetchall() == [(1,)] + arrow_cursor.execute( + """ + SELECT + 1 AS col + ,CAST('12:34:56' AS TIME) AS col_time + ,X'0102' AS col_varbinary + ,json_parse('{"a": 1}') AS col_json + ,CAST('{"a": 1}' AS JSON) AS col_json_string + """ + ) + assert arrow_cursor.fetchall() == [ + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}') + ]