From 5ea497cfcb051c8fb148553326e60bb4661a4a4d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 20:07:31 +0900 Subject: [PATCH 1/9] Keep Athena json results decoded by the converter in SQLAlchemy JSON PyAthena's converters already decode results of the Athena json type, and SQLAlchemy's JSON result processor decoded them again, raising TypeError for JSON objects and arrays. The dialect now maps JSON to AthenaJSON, which returns json results unchanged and decodes JSON text of other Athena types, such as varchar, as before. A CAST of text AS JSON is a JSON string scalar in Athena, so it now comes back as a str instead of being decoded a second time. Closes #934 Co-Authored-By: Claude Opus 5.5 --- docs/sqlalchemy.md | 47 +++++++------- pyathena/sqlalchemy/base.py | 2 + pyathena/sqlalchemy/types.py | 33 +++++++++- tests/pyathena/sqlalchemy/test_base.py | 81 ++++++++----------------- tests/pyathena/sqlalchemy/test_types.py | 41 +++++++++++++ 5 files changed, 120 insertions(+), 84 deletions(-) diff --git a/docs/sqlalchemy.md b/docs/sqlalchemy.md index 6e285f19c..48be0c389 100644 --- a/docs/sqlalchemy.md +++ b/docs/sqlalchemy.md @@ -1365,60 +1365,55 @@ events = Table('events', metadata, #### Querying JSON data -When querying JSON data, PyAthena automatically parses JSON strings into Python dictionaries: +PyAthena's converters decode results of Athena's `json` type, and the `JSON` type returns those values unchanged. +A JSON object becomes a `dict`, an array a `list`, and a string scalar a `str`: ```python from sqlalchemy import select, literal_column from sqlalchemy.sql import type_coerce from sqlalchemy.types import JSON -# Query with explicit type coercion result = connection.execute( select( type_coerce( - literal_column('CAST(\'{"name": "test", "value": 123}\' AS JSON)'), + literal_column('json_parse(\'{"name": "test", "items": [1, 2, 3]}\')'), JSON ).label("json_col") ) ).fetchone() -# Result is automatically parsed as a dictionary -print(result.json_col) # {"name": "test", "value": 123} +print(result.json_col) # {'items': [1, 2, 3], 'name': 'test'} print(type(result.json_col)) # ``` -#### Important limitations - -Athena's JSON type support has specific limitations: - -- **JSON objects and arrays are supported** - `CAST('...' AS JSON)` accepts an object or a top-level array such as `[1, 2, 3]` -- **Arrays within objects are supported** - JSON objects can contain arrays as property values -- **DML only** - JSON type is supported for SELECT queries but not in CREATE TABLE statements; compiling `CREATE TABLE` with a `JSON` column raises `CompileError` +`json_parse()` parses text into a JSON value, while `CAST('...' AS JSON)` returns the text as a JSON string scalar: ```python -# Supported: JSON object with nested array result = connection.execute( select( - type_coerce( - literal_column('CAST(\'{"items": [1, 2, 3]}\' AS JSON)'), - JSON - ).label("json_col") + type_coerce(literal_column("json_parse('[1, 2, 3]')"), JSON).label("parsed"), + type_coerce(literal_column("CAST('[1, 2, 3]' AS JSON)"), JSON).label("cast"), ) ).fetchone() -print(result.json_col) # {"items": [1, 2, 3]} -# Supported: Top-level array +print(result.parsed) # [1, 2, 3] +print(result.cast) # '[1, 2, 3]' +``` + +The `JSON` type decodes JSON text in columns of other types, such as a `varchar` column, with `json.loads()` or the dialect's `json_deserializer`: + +```python result = connection.execute( - select( - type_coerce( - literal_column("CAST('[1, 2, 3]' AS JSON)"), - JSON - ).label("json_col") - ) + select(type_coerce(literal_column("'{\"a\": 1}'"), JSON).label("json_col")) ).fetchone() -print(result.json_col) # [1, 2, 3] + +print(result.json_col) # {'a': 1} ``` +#### Important limitations + +- **DML only** - JSON type is supported for SELECT queries but not in CREATE TABLE statements; compiling `CREATE TABLE` with a `JSON` column raises `CompileError` + #### Best practices 1. **Use with SELECT queries** - JSON type works best for querying existing data diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index d620fce6e..f408fdf06 100644 --- a/pyathena/sqlalchemy/base.py +++ b/pyathena/sqlalchemy/base.py @@ -38,6 +38,7 @@ AthenaArray, AthenaBinary, AthenaDate, + AthenaJSON, AthenaMap, AthenaStruct, AthenaTimestamp, @@ -212,6 +213,7 @@ class AthenaDialect(DefaultDialect): types.ARRAY: AthenaArray, types.Date: AthenaDate, types.DateTime: AthenaTimestamp, + types.JSON: AthenaJSON, } ischema_names: dict[str, type[Any]] = ischema_names diff --git a/pyathena/sqlalchemy/types.py b/pyathena/sqlalchemy/types.py index 07c8e69d8..c1bb26afb 100644 --- a/pyathena/sqlalchemy/types.py +++ b/pyathena/sqlalchemy/types.py @@ -6,7 +6,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from sqlalchemy import types from sqlalchemy.sql import sqltypes @@ -19,7 +19,7 @@ if TYPE_CHECKING: from sqlalchemy import Dialect - from sqlalchemy.sql.type_api import _LiteralProcessorType + from sqlalchemy.sql.type_api import _LiteralProcessorType, _ResultProcessorType __all__ = [ "ARRAY", @@ -29,6 +29,7 @@ "AthenaArray", "AthenaBinary", "AthenaDate", + "AthenaJSON", "AthenaMap", "AthenaStruct", "AthenaTimestamp", @@ -47,6 +48,34 @@ def process(value: bytes) -> str: return process +class AthenaJSON(types.JSON): + """SQLAlchemy JSON type that keeps the values PyAthena has decoded. + + PyAthena's converters decode results of the Athena ``json`` type, so this + type returns them unchanged, and a JSON string scalar stays a ``str``. + Results of other Athena types, such as JSON text in a ``varchar`` column, + are decoded with the dialect's JSON deserializer. + """ + + @override + def result_processor( + self, dialect: Dialect, coltype: object + ) -> _ResultProcessorType[Any] | None: + """Return a processor decoding JSON text of Athena types other than ``json``. + + Args: + dialect: The dialect fetching the value. + coltype: The Athena type name from the cursor description. + + Returns: + The processor, or None for the Athena ``json`` type. + """ + if coltype == "json": + return None + processor: _ResultProcessorType[Any] | None = super().result_processor(dialect, coltype) + return processor + + class Tinyint(sqltypes.Integer): """SQLAlchemy type for Athena TINYINT (8-bit signed integer). diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 9e0536bc7..eb6055d46 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -891,67 +891,36 @@ def test_basic_query(self, engine): assert rows[0].number_of_rows == 1 assert len(rows[0]) == 1 - def test_json_type_with_cast(self, engine): - """Test JSON type support with CAST operation in SELECT query.""" - engine, conn = engine - # Note: Athena JSON type support has limitations - # - JSON objects are supported - # - Direct CAST of JSON arrays is not supported - # - JSON is primarily used with DML operations, not DDL - - # Test 1: Simple JSON object with type_coerce for proper type handling - result = conn.execute( - select( - type_coerce( - literal_column('CAST(\'{"name": "test", "value": 123}\' AS JSON)'), - types.JSON, - ).label("json_col") - ) - ).fetchone() - assert result.json_col == {"name": "test", "value": 123} - assert isinstance(result.json_col, dict) - - # Test 2: Nested JSON object with arrays inside - # (Arrays are supported as part of JSON objects, just not as top-level CAST) - nested_json_str = '{"user": {"id": 1, "name": "Alice"}, "scores": [95, 87, 92]}' - result = conn.execute( - select( - type_coerce(literal_column(f"CAST('{nested_json_str}' AS JSON)"), types.JSON).label( - "nested_json" - ) - ) - ).fetchone() - assert result.nested_json == { - "user": {"id": 1, "name": "Alice"}, - "scores": [95, 87, 92], + def test_json_type(self, engine): + engine, conn = engine + nested = '{"user": {"id": 1, "name": "Alice"}, "scores": [95, 87, 92]}' + scalars = '{"str": "value", "num": 42, "bool": true, "nil": null}' + columns = { + "obj": f"json_parse('{nested}')", + "arr": "json_parse('[1, 2, 3]')", + "scalars": f"json_parse('{scalars}')", + # Athena returns a cast of text as a JSON string scalar. + "cast_text": "CAST('{\"a\": 1}' AS JSON)", + "missing": "CAST(NULL AS JSON)", + # JSON text in a varchar column is decoded. + "text_obj": "'{\"a\": 1}'", } - assert result.nested_json["user"]["name"] == "Alice" - assert result.nested_json["scores"][0] == 95 - assert isinstance(result.nested_json["scores"], list) - - # Test 3: JSON with null value result = conn.execute( select( - type_coerce(literal_column("CAST('{\"key\": null}' AS JSON)"), types.JSON).label( - "json_with_null" + *( + type_coerce(literal_column(expr), types.JSON).label(name) + for name, expr in columns.items() ) ) - ).fetchone() - assert result.json_with_null == {"key": None} - assert result.json_with_null["key"] is None - - # Test 4: JSON with various types - result = conn.execute( - select( - type_coerce( - literal_column( - 'CAST(\'{"str": "value", "num": 42, "bool": true, "nil": null}\' AS JSON)' - ), - types.JSON, - ).label("json_types") - ) - ).fetchone() - assert result.json_types == {"str": "value", "num": 42, "bool": True, "nil": None} + ).one() + assert result._asdict() == { + "obj": {"user": {"id": 1, "name": "Alice"}, "scores": [95, 87, 92]}, + "arr": [1, 2, 3], + "scalars": {"str": "value", "num": 42, "bool": True, "nil": None}, + "cast_text": '{"a": 1}', + "missing": None, + "text_obj": {"a": 1}, + } def test_select_nested_struct_query(self, engine): """Test SELECT query with nested STRUCT (ROW) types (Issue #627).""" diff --git a/tests/pyathena/sqlalchemy/test_types.py b/tests/pyathena/sqlalchemy/test_types.py index 665a16cbc..e144bc5ea 100644 --- a/tests/pyathena/sqlalchemy/test_types.py +++ b/tests/pyathena/sqlalchemy/test_types.py @@ -1,7 +1,48 @@ +import json + +import pytest from sqlalchemy import types from pyathena.sqlalchemy.base import ischema_names +from pyathena.sqlalchemy.rest import AthenaRestDialect def test_double_column_type(): assert ischema_names["double"] is types.DOUBLE + + +class _JSONText(types.TypeDecorator): + impl = types.JSON + cache_ok = True + + +class TestAthenaJSON: + @staticmethod + def _process(type_, value, coltype, dialect=None): + dialect = dialect or AthenaRestDialect() + processor = type_.dialect_impl(dialect)._cached_result_processor(dialect, coltype) + return processor(value) if processor else value + + @pytest.mark.parametrize("type_", [types.JSON(), _JSONText()]) + @pytest.mark.parametrize( + "value", + [{"a": 1}, [1, 2, 3], '{"a": 1}', "text", 1, True, None], + ) + def test_keeps_converted_json_results(self, type_, value): + assert self._process(type_, value, "json") == value + + @pytest.mark.parametrize("type_", [types.JSON(), _JSONText()]) + @pytest.mark.parametrize( + ("value", "expected"), + [('{"a": 1}', {"a": 1}), ("[1, 2, 3]", [1, 2, 3]), ('"text"', "text"), (None, None)], + ) + def test_decodes_json_text_of_other_types(self, type_, value, expected): + assert self._process(type_, value, "varchar") == expected + + def test_decodes_json_text_with_dialect_deserializer(self): + dialect = AthenaRestDialect(json_deserializer=lambda v: ("custom", json.loads(v))) + assert self._process(types.JSON(), '{"a": 1}', "varchar", dialect) == ( + "custom", + {"a": 1}, + ) + assert self._process(types.JSON(), {"a": 1}, "json", dialect) == {"a": 1} From 428287b580625a2e62fa7f6dbf86e054e592e81a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 20:08:42 +0900 Subject: [PATCH 2/9] Use the public result_processor in the AthenaJSON tests Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/sqlalchemy/test_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pyathena/sqlalchemy/test_types.py b/tests/pyathena/sqlalchemy/test_types.py index e144bc5ea..a5be6b0d3 100644 --- a/tests/pyathena/sqlalchemy/test_types.py +++ b/tests/pyathena/sqlalchemy/test_types.py @@ -20,7 +20,7 @@ class TestAthenaJSON: @staticmethod def _process(type_, value, coltype, dialect=None): dialect = dialect or AthenaRestDialect() - processor = type_.dialect_impl(dialect)._cached_result_processor(dialect, coltype) + processor = type_.dialect_impl(dialect).result_processor(dialect, coltype) return processor(value) if processor else value @pytest.mark.parametrize("type_", [types.JSON(), _JSONText()]) From b175da964eefa4f1365ee2c30e9bc1a00d1709d4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 21:35:45 +0900 Subject: [PATCH 3/9] Qualify the JSON converter statements for custom converters Co-Authored-By: Claude Opus 5.5 --- docs/sqlalchemy.md | 5 +++-- pyathena/sqlalchemy/types.py | 10 ++++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/sqlalchemy.md b/docs/sqlalchemy.md index 48be0c389..66b8b40e1 100644 --- a/docs/sqlalchemy.md +++ b/docs/sqlalchemy.md @@ -1365,8 +1365,9 @@ events = Table('events', metadata, #### Querying JSON data -PyAthena's converters decode results of Athena's `json` type, and the `JSON` type returns those values unchanged. -A JSON object becomes a `dict`, an array a `list`, and a string scalar a `str`: +PyAthena's default converters decode results of Athena's `json` type, and the `JSON` type returns those values unchanged. +A JSON object becomes a `dict`, an array a `list`, and a string scalar a `str`. +With a custom converter that does not decode `json` results, the `JSON` type returns their text. ```python from sqlalchemy import select, literal_column diff --git a/pyathena/sqlalchemy/types.py b/pyathena/sqlalchemy/types.py index c1bb26afb..65b65d540 100644 --- a/pyathena/sqlalchemy/types.py +++ b/pyathena/sqlalchemy/types.py @@ -51,10 +51,12 @@ def process(value: bytes) -> str: class AthenaJSON(types.JSON): """SQLAlchemy JSON type that keeps the values PyAthena has decoded. - PyAthena's converters decode results of the Athena ``json`` type, so this - type returns them unchanged, and a JSON string scalar stays a ``str``. - Results of other Athena types, such as JSON text in a ``varchar`` column, - are decoded with the dialect's JSON deserializer. + PyAthena's default converters decode results of the Athena ``json`` type, + so this type returns them unchanged, and a JSON string scalar stays a + ``str``. With a custom converter that does not decode ``json`` results, + this type returns their text. Results of other Athena types, such as JSON + text in a ``varchar`` column, are decoded with the dialect's JSON + deserializer. """ @override From c5d9830ae76e9df0d1024031c2eca90ed6e2df4f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 21:42:26 +0900 Subject: [PATCH 4/9] State the JSON text pass-through by what the cursor returns Co-Authored-By: Claude Opus 5.5 --- docs/sqlalchemy.md | 2 +- pyathena/sqlalchemy/types.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/sqlalchemy.md b/docs/sqlalchemy.md index 66b8b40e1..b6a287a25 100644 --- a/docs/sqlalchemy.md +++ b/docs/sqlalchemy.md @@ -1367,7 +1367,7 @@ events = Table('events', metadata, PyAthena's default converters decode results of Athena's `json` type, and the `JSON` type returns those values unchanged. A JSON object becomes a `dict`, an array a `list`, and a string scalar a `str`. -With a custom converter that does not decode `json` results, the `JSON` type returns their text. +If the cursor returns the text of a `json` result instead, for example through a custom converter that does not decode it, the `JSON` type returns that text unchanged. ```python from sqlalchemy import select, literal_column diff --git a/pyathena/sqlalchemy/types.py b/pyathena/sqlalchemy/types.py index 65b65d540..74872d11f 100644 --- a/pyathena/sqlalchemy/types.py +++ b/pyathena/sqlalchemy/types.py @@ -53,8 +53,9 @@ class AthenaJSON(types.JSON): PyAthena's default converters decode results of the Athena ``json`` type, so this type returns them unchanged, and a JSON string scalar stays a - ``str``. With a custom converter that does not decode ``json`` results, - this type returns their text. Results of other Athena types, such as JSON + ``str``. If the cursor returns the text of a ``json`` result instead, for + example through a custom converter that does not decode it, this type + returns that text unchanged. Results of other Athena types, such as JSON text in a ``varchar`` column, are decoded with the dialect's JSON deserializer. """ From a96a81edfdf60a25dde22636a38ed29d06cb021e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 23:33:46 +0900 Subject: [PATCH 5/9] Fix json results of the pandas and S3FS cursors - _to_json returns None for an empty string, which CSV results use for SQL NULL, so PandasCursor and ArrowCursor no longer raise JSONDecodeError on a NULL json value. - S3FSCursor converts GetQueryResults rows with its own converter, as it does for the CSV result file, instead of DefaultTypeConverter. - The ArrowCursor, PandasCursor, and PolarsCursor fallback tests cover NULL json values and JSON string scalars (the ArrowCursor fallback conversion itself was fixed in #1023). Co-Authored-By: Claude Opus 5.5 --- pyathena/converter.py | 11 +++++++++- pyathena/s3fs/result_set.py | 5 +++-- tests/pyathena/arrow/test_cursor.py | 3 ++- tests/pyathena/pandas/test_cursor.py | 10 +++++++-- tests/pyathena/polars/test_cursor.py | 10 +++++++-- tests/pyathena/s3fs/test_cursor.py | 31 ++++++++++++++++++++++++++++ tests/pyathena/test_converter.py | 17 +++++++++++++++ 7 files changed, 79 insertions(+), 8 deletions(-) diff --git a/pyathena/converter.py b/pyathena/converter.py index 3ea1577a2..d8b3405b0 100644 --- a/pyathena/converter.py +++ b/pyathena/converter.py @@ -119,7 +119,16 @@ def _to_binary(varchar_value: str | None) -> bytes | None: def _to_json(varchar_value: str | None) -> Any | None: - if varchar_value is None: + """Convert an Athena JSON value to the decoded Python value. + + Args: + varchar_value: The value as JSON text, or None. An empty string, which + CSV results use for SQL NULL, is also treated as NULL. + + Returns: + The decoded value, or None for SQL NULL. + """ + if not varchar_value: return None return json.loads(varchar_value) diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 74dff6c7b..0bc268fc9 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -120,8 +120,9 @@ def __init__( if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: self._init_csv_reader() elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: - # Managed query result storage: no output_location, use API - rows = self._fetch_all_rows() + # Managed query result storage: no output_location, use API. + # The converter reads text values, as from the CSV result file. + rows = self._fetch_all_rows(self._converter) self._rows.extend(rows) # If CSV reader was not initialized (e.g., CTAS, DDL), diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index f70e99f65..ae113d803 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -977,8 +977,9 @@ def test_fetch_all_rows(self, arrow_cursor): ,X'0102' AS col_varbinary ,json_parse('{"a": 1}') AS col_json ,CAST('{"a": 1}' AS JSON) AS col_json_string + ,CAST(NULL AS JSON) AS col_json_null """ ) assert arrow_cursor.fetchall() == [ - (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}') + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}', None) ] diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 3d0f50e5e..afa1ff3ed 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1590,5 +1590,11 @@ def test_pandas_cursor_iter_chunks_consistency(self, pandas_cursor): indirect=["pandas_cursor"], ) def test_fetch_all_rows(self, pandas_cursor): - pandas_cursor.execute("SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time") - assert pandas_cursor.fetchall() == [(1, datetime(2017, 1, 1, 12, 34, 56).time())] + pandas_cursor.execute( + "SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time, X'0001' AS col_binary, " + "json_parse('{\"a\": 1}') AS col_json, CAST('[1, 2]' AS JSON) AS col_json_string, " + "CAST(NULL AS JSON) AS col_json_null" + ) + assert pandas_cursor.fetchall() == [ + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None) + ] diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index c9e30da04..681542bdb 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -751,5 +751,11 @@ def test_null_vs_empty_string(self, polars_cursor): indirect=["polars_cursor"], ) def test_fetch_all_rows(self, polars_cursor): - polars_cursor.execute("SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time") - assert polars_cursor.fetchall() == [(1, datetime(2017, 1, 1, 12, 34, 56).time())] + polars_cursor.execute( + "SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time, X'0001' AS col_binary, " + "json_parse('{\"a\": 1}') AS col_json, CAST('[1, 2]' AS JSON) AS col_json_string, " + "CAST(NULL AS JSON) AS col_json_null" + ) + assert polars_cursor.fetchall() == [ + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None) + ] diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 777303f4e..8ff23115d 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -8,7 +8,9 @@ import pytest +from pyathena.converter import _to_default from pyathena.error import DatabaseError, ProgrammingError +from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.cursor import S3FSCursor from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet @@ -17,6 +19,12 @@ from tests.pyathena.util import cached_file_systems +def _json_text_converter(): + converter = DefaultS3FSTypeConverter() + converter.set("json", _to_default) + return converter + + class TestS3FSCursor: def test_fetchone(self, s3fs_cursor): s3fs_cursor.execute("SELECT * FROM one_row") @@ -528,3 +536,26 @@ def test_quoted_string_with_comma(self, csv_reader_class): def test_fetch_all_rows(self, s3fs_cursor): s3fs_cursor.execute("SELECT 1 AS col") assert s3fs_cursor.fetchall() == [(1,)] + + @pytest.mark.parametrize( + "s3fs_cursor", + [ + pytest.param({"converter": _json_text_converter()}, id="default"), + pytest.param( + { + "work_group": ENV.managed_work_group, + "s3_staging_dir": "", + "converter": _json_text_converter(), + }, + id="managed", + marks=pytest.mark.skipif( + not ENV.managed_work_group, + reason="AWS_ATHENA_MANAGED_WORKGROUP not set", + ), + ), + ], + indirect=["s3fs_cursor"], + ) + def test_fetch_all_rows_custom_converter(self, s3fs_cursor): + s3fs_cursor.execute("SELECT 1 AS col, json_parse('{\"a\": 1}') AS col_json") + assert s3fs_cursor.fetchall() == [(1, '{"a":1}')] diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index 3dd102498..81322d04f 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -8,6 +8,7 @@ _to_array, _to_datetime, _to_datetime_with_tz, + _to_json, _to_map, _to_struct, ) @@ -325,6 +326,22 @@ def test_to_array_invalid_formats(input_value): assert _to_array(input_value) is None +@pytest.mark.parametrize( + ("input_value", "expected"), + [ + (None, None), + ("", None), + ('{"a":1}', {"a": 1}), + ("[1,2]", [1, 2]), + ('""', ""), + ('"[1, 2]"', "[1, 2]"), + ("null", None), + ], +) +def test_to_json(input_value, expected): + assert _to_json(input_value) == expected + + class TestDefaultTypeConverter: @pytest.mark.parametrize( ("input_value", "expected"), From 90537de87ff745e970fa0232c10ce5f4d7a85e59 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 22:30:26 +0900 Subject: [PATCH 6/9] Keep json as text in the Arrow and Polars GetQueryResults tables With managed query result storage, the Arrow table and Polars DataFrame were built from decoded json values, so pyarrow and Polars inferred struct or list columns and failed when rows held different JSON shapes (for example an object and an array). Build them with json values kept as text, as in the CSV result file, and decode only the json columns in the fetch methods. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/result_set.py | 22 +++++++++++++-------- pyathena/polars/result_set.py | 9 ++++++--- pyathena/result_set.py | 29 +++++++++++++++++++++++++++- tests/pyathena/arrow/test_cursor.py | 8 +++++++- tests/pyathena/pandas/test_cursor.py | 7 +++++-- tests/pyathena/polars/test_cursor.py | 8 ++++++-- 6 files changed, 66 insertions(+), 17 deletions(-) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 77978264a..50a58dba2 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -12,7 +12,7 @@ from pyathena import OperationalError from pyathena.arrow.util import to_column_info -from pyathena.converter import Converter +from pyathena.converter import Converter, _to_default from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet @@ -146,8 +146,8 @@ 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. + # The fetch methods convert the values read from a result file. GetQueryResults + # values are already converted, except json values, which stay text. self._convert_rows = bool(self.output_location) self._batches = iter(self._table.to_batches(arraysize)) @@ -254,11 +254,16 @@ def _fetch(self) -> None: return else: dict_rows = rows.to_pydict() - if self._convert_rows: - converters = self.converters + converters = ( + self.converters if self._convert_rows else self._json_converters(self.converters) + ) + if converters: column_names = dict_rows.keys() processed_rows = [ - tuple(converters[k](v) for k, v in zip(column_names, row, strict=False)) + tuple( + converters.get(k, _to_default)(v) + for k, v in zip(column_names, row, strict=False) + ) for row in zip(*dict_rows.values(), strict=False) ] else: @@ -381,11 +386,12 @@ def _as_arrow_from_api(self, converter: Converter | None = None) -> Table: Args: converter: Type converter for result values. Defaults to - ``DefaultTypeConverter`` if not specified. + ``DefaultTypeConverter`` with json values kept as text, as in + the CSV result file. """ import pyarrow as pa - rows = self._fetch_all_rows(converter) + rows = self._fetch_all_rows(converter or self._json_text_converter()) if not rows: return pa.Table.from_pydict({}) description = self.description if self.description else [] diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 039850751..35d267165 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -262,7 +262,7 @@ def __init__( # Note: _as_polars() and _create_dataframe_iterator() update _metadata for unload # queries, so the converters and column names must be read AFTER them. self._df: pl.DataFrame | None = None - # Converters for the rows of self._df. GetQueryResults values are already converted. + # Converters for the rows of self._df. self._df_converters: dict[str, Callable[[str | None], Any | None]] = {} if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: if self._chunksize is None: @@ -272,6 +272,8 @@ def __init__( self._df_iter = self._create_dataframe_iterator() elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: self._df = self._as_polars_from_api() + # GetQueryResults values are already converted, except json values kept as text. + self._df_converters = self._json_converters(self.converters) else: self._df = pl.DataFrame() if self._df is not None: @@ -539,11 +541,12 @@ def _as_polars_from_api(self, converter: Converter | None = None) -> pl.DataFram Args: converter: Type converter for result values. Defaults to - ``DefaultTypeConverter`` if not specified. + ``DefaultTypeConverter`` with json values kept as text, as in + the CSV result file. """ import polars as pl - rows = self._fetch_all_rows(converter) + rows = self._fetch_all_rows(converter or self._json_text_converter()) if not rows: return pl.DataFrame() description = self.description if self.description else [] diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 410c471b3..6341f88b1 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -12,12 +12,14 @@ ) from pyathena.common import BaseCursor, CursorIterator -from pyathena.converter import Converter, DefaultTypeConverter +from pyathena.converter import Converter, DefaultTypeConverter, _to_default from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig, override, parse_output_location, retry_api_call if TYPE_CHECKING: + from collections.abc import Callable + from pyathena.connection import Connection _logger = logging.getLogger(__name__) @@ -679,6 +681,31 @@ def _is_first_row_column_labels(self, rows: list[dict[str, Any]]) -> bool: return False return True + @staticmethod + def _json_text_converter() -> DefaultTypeConverter: + """Return a ``DefaultTypeConverter`` that keeps json values as text. + + Returns: + The converter. + """ + converter = DefaultTypeConverter() + converter.set("json", _to_default) + return converter + + def _json_converters( + self, converters: dict[str, Callable[[str | None], Any | None]] + ) -> dict[str, Callable[[str | None], Any | None]]: + """Select the converters of the json columns. + + Args: + converters: The converters keyed by column name. + + Returns: + The converters of the columns whose Athena type is json. + """ + description = self.description if self.description else [] + return {d[0]: converters[d[0]] for d in description if d[1] == "json"} + def _fetch_all_rows( self, converter: Converter | None = None, diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index ae113d803..df58ada7d 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -978,8 +978,14 @@ def test_fetch_all_rows(self, arrow_cursor): ,json_parse('{"a": 1}') AS col_json ,CAST('{"a": 1}' AS JSON) AS col_json_string ,CAST(NULL AS JSON) AS col_json_null + UNION ALL + SELECT + 2, CAST('12:34:56' AS TIME), X'0102', json_parse('[1, "x"]'), json_parse('"s"'), NULL + ORDER BY col """ ) assert arrow_cursor.fetchall() == [ - (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}', None) + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", {"a": 1}, '{"a": 1}', None), + (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x01\x02", [1, "x"], "s", None), ] + assert arrow_cursor.as_arrow().schema.field("col_json").type == pa.string() diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index afa1ff3ed..db5c55c70 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1593,8 +1593,11 @@ def test_fetch_all_rows(self, pandas_cursor): pandas_cursor.execute( "SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time, X'0001' AS col_binary, " "json_parse('{\"a\": 1}') AS col_json, CAST('[1, 2]' AS JSON) AS col_json_string, " - "CAST(NULL AS JSON) AS col_json_null" + "CAST(NULL AS JSON) AS col_json_null " + "UNION ALL SELECT 2, CAST('12:34:56' AS TIME), X'0001', json_parse('[1, \"x\"]'), " + "json_parse('\"s\"'), NULL ORDER BY col" ) assert pandas_cursor.fetchall() == [ - (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None) + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None), + (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", [1, "x"], "s", None), ] diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 681542bdb..c8727f03d 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -754,8 +754,12 @@ def test_fetch_all_rows(self, polars_cursor): polars_cursor.execute( "SELECT 1 AS col, CAST('12:34:56' AS TIME) AS col_time, X'0001' AS col_binary, " "json_parse('{\"a\": 1}') AS col_json, CAST('[1, 2]' AS JSON) AS col_json_string, " - "CAST(NULL AS JSON) AS col_json_null" + "CAST(NULL AS JSON) AS col_json_null " + "UNION ALL SELECT 2, CAST('12:34:56' AS TIME), X'0001', json_parse('[1, \"x\"]'), " + "json_parse('\"s\"'), NULL ORDER BY col" ) assert polars_cursor.fetchall() == [ - (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None) + (1, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", {"a": 1}, "[1, 2]", None), + (2, datetime(2017, 1, 1, 12, 34, 56).time(), b"\x00\x01", [1, "x"], "s", None), ] + assert polars_cursor.as_polars()["col_json"].dtype == pl.String From fe4d744e2de885ca2894c0991427e943c9d54022 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 22:50:01 +0900 Subject: [PATCH 7/9] Treat empty CSV fields as NULL only in the CSV json converters Returning None for an empty string from _to_json also changed nested values: the typed parser passes an empty JSON string element of array(json) or map(..., json) as "" to it, so [""] became [None]. Keep _to_json as it was and give the pandas and Arrow converters, which read CSV result files, a _csv_to_json that maps the empty field to None. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/converter.py | 4 ++-- pyathena/converter.py | 8 +++++++- pyathena/pandas/converter.py | 4 ++-- tests/pyathena/test_converter.py | 17 ++++++++++++++--- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/pyathena/arrow/converter.py b/pyathena/arrow/converter.py index ea6482557..88943efba 100644 --- a/pyathena/arrow/converter.py +++ b/pyathena/arrow/converter.py @@ -9,11 +9,11 @@ from pyathena.converter import ( Converter, + _csv_to_json, _to_binary, _to_date, _to_decimal, _to_default, - _to_json, _to_time, ) from pyathena.util import override @@ -26,7 +26,7 @@ "time": _to_time, "decimal": _to_decimal, "varbinary": _to_binary, - "json": _to_json, + "json": _csv_to_json, } diff --git a/pyathena/converter.py b/pyathena/converter.py index d8b3405b0..b6831d01b 100644 --- a/pyathena/converter.py +++ b/pyathena/converter.py @@ -119,7 +119,13 @@ def _to_binary(varchar_value: str | None) -> bytes | None: def _to_json(varchar_value: str | None) -> Any | None: - """Convert an Athena JSON value to the decoded Python value. + if varchar_value is None: + return None + return json.loads(varchar_value) + + +def _csv_to_json(varchar_value: str | None) -> Any | None: + """Convert an Athena JSON value read from a CSV result file. Args: varchar_value: The value as JSON text, or None. An empty string, which diff --git a/pyathena/pandas/converter.py b/pyathena/pandas/converter.py index a80975871..a3b207e90 100644 --- a/pyathena/pandas/converter.py +++ b/pyathena/pandas/converter.py @@ -9,11 +9,11 @@ from pyathena.converter import ( Converter, + _csv_to_json, _to_binary, _to_boolean, _to_decimal, _to_default, - _to_json, ) from pyathena.util import override @@ -24,7 +24,7 @@ "boolean": _to_boolean, "decimal": _to_decimal, "varbinary": _to_binary, - "json": _to_json, + "json": _csv_to_json, } diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index 81322d04f..5d7ad4012 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -5,10 +5,10 @@ from pyathena.converter import ( DefaultTypeConverter, + _csv_to_json, _to_array, _to_datetime, _to_datetime_with_tz, - _to_json, _to_map, _to_struct, ) @@ -338,8 +338,19 @@ def test_to_array_invalid_formats(input_value): ("null", None), ], ) -def test_to_json(input_value, expected): - assert _to_json(input_value) == expected +def test_csv_to_json(input_value, expected): + assert _csv_to_json(input_value) == expected + + +@pytest.mark.parametrize( + ("value", "type_hint", "expected"), + [ + ('[""]', "array(json)", [""]), + ('{"k": ""}', "map(varchar,json)", {"k": ""}), + ], +) +def test_nested_json_empty_string(value, type_hint, expected): + assert DefaultTypeConverter().convert(type_hint.split("(")[0], value, type_hint) == expected class TestDefaultTypeConverter: From d5eac4c78ef6fb7c2f009137d9dd537087598487 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 23:34:51 +0900 Subject: [PATCH 8/9] Define the json-text converter in the converter module The result sets only choose converters; the converter that keeps json values as text is a converter definition, so it moves from AthenaResultSet to pyathena.converter. Co-Authored-By: Claude Opus 5.5 --- pyathena/arrow/result_set.py | 4 ++-- pyathena/converter.py | 11 +++++++++++ pyathena/polars/result_set.py | 4 ++-- pyathena/result_set.py | 13 +------------ 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 50a58dba2..abdc66abb 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -12,7 +12,7 @@ from pyathena import OperationalError from pyathena.arrow.util import to_column_info -from pyathena.converter import Converter, _to_default +from pyathena.converter import Converter, _json_text_converter, _to_default from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet @@ -391,7 +391,7 @@ def _as_arrow_from_api(self, converter: Converter | None = None) -> Table: """ import pyarrow as pa - rows = self._fetch_all_rows(converter or self._json_text_converter()) + rows = self._fetch_all_rows(converter or _json_text_converter()) if not rows: return pa.Table.from_pydict({}) description = self.description if self.description else [] diff --git a/pyathena/converter.py b/pyathena/converter.py index b6831d01b..eff987421 100644 --- a/pyathena/converter.py +++ b/pyathena/converter.py @@ -742,3 +742,14 @@ def _parse_type_hint(self, type_hint: str) -> TypeNode: if normalized not in self._parsed_hints: self._parsed_hints[normalized] = self._parser.parse(normalized) return self._parsed_hints[normalized] + + +def _json_text_converter() -> DefaultTypeConverter: + """Return a ``DefaultTypeConverter`` that keeps json values as text. + + Returns: + The converter. + """ + converter = DefaultTypeConverter() + converter.set("json", _to_default) + return converter diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 35d267165..34dbebe40 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -20,7 +20,7 @@ ) from pyathena import OperationalError -from pyathena.converter import Converter +from pyathena.converter import Converter, _json_text_converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.polars.util import to_column_info @@ -546,7 +546,7 @@ def _as_polars_from_api(self, converter: Converter | None = None) -> pl.DataFram """ import polars as pl - rows = self._fetch_all_rows(converter or self._json_text_converter()) + rows = self._fetch_all_rows(converter or _json_text_converter()) if not rows: return pl.DataFrame() description = self.description if self.description else [] diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 6341f88b1..65842761e 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -12,7 +12,7 @@ ) from pyathena.common import BaseCursor, CursorIterator -from pyathena.converter import Converter, DefaultTypeConverter, _to_default +from pyathena.converter import Converter, DefaultTypeConverter from pyathena.error import DataError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig, override, parse_output_location, retry_api_call @@ -681,17 +681,6 @@ def _is_first_row_column_labels(self, rows: list[dict[str, Any]]) -> bool: return False return True - @staticmethod - def _json_text_converter() -> DefaultTypeConverter: - """Return a ``DefaultTypeConverter`` that keeps json values as text. - - Returns: - The converter. - """ - converter = DefaultTypeConverter() - converter.set("json", _to_default) - return converter - def _json_converters( self, converters: dict[str, Callable[[str | None], Any | None]] ) -> dict[str, Callable[[str | None], Any | None]]: From df476e484c3c0c8208288c54d27761cecf6f2662 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 00:36:18 +0900 Subject: [PATCH 9/9] Name the S3FS test converter after what it builds The production _json_text_converter() returns a DefaultTypeConverter for the GetQueryResults fallback; the S3FS test helper builds a DefaultS3FSTypeConverter, so give it its own name. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/s3fs/test_cursor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 8ff23115d..1d12d7041 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -19,7 +19,7 @@ from tests.pyathena.util import cached_file_systems -def _json_text_converter(): +def _s3fs_converter_with_json_text(): converter = DefaultS3FSTypeConverter() converter.set("json", _to_default) return converter @@ -540,12 +540,12 @@ def test_fetch_all_rows(self, s3fs_cursor): @pytest.mark.parametrize( "s3fs_cursor", [ - pytest.param({"converter": _json_text_converter()}, id="default"), + pytest.param({"converter": _s3fs_converter_with_json_text()}, id="default"), pytest.param( { "work_group": ENV.managed_work_group, "s3_staging_dir": "", - "converter": _json_text_converter(), + "converter": _s3fs_converter_with_json_text(), }, id="managed", marks=pytest.mark.skipif(