Keep integer and JSON values exact in PandasCursor results with NULL - #1044
Conversation
| for i in range(df.shape[1]): | ||
| if (values := _unwrap_csv_objects(df.iloc[:, i])) is not None: | ||
| df.isetitem(i, pd.Series(values, index=df.index, dtype=object)) | ||
| index = df.index |
There was a problem hiding this comment.
Self-review round one (implementation behavior) — FINDINGS, repaired in eff8530.
Scope: git diff a73c6db457110ea13eae8f7d1e7670b9095605dd..fd778a21bce684fef14d0a0eeae0f2a228637e62. Covered: CSV converter wrapping in _get_csv_read_options() and _configure_binary_csv_read(), unwrapping in _finish_csv_frame() for whole results, chunks (PandasDataFrameIterator), and indexes; managed _as_pandas_from_api() dtypes; the pyarrow CSV engine (never chosen with converters, so it never sees wrapped values); UNLOAD (keeps _no_trunc_date); sync/async/aio pandas cursors (shared result set); the regression test.
Finding 1: with index_col=["col_int", "col_json"], the json level of the MultiIndex kept the _CSVObject wrappers, because the first version unwrapped only a flat Index. Before this PR, that level held the converted values. Repair: unwrap each level and rebuild the MultiIndex (pyathena/pandas/result_set.py:570).
| pytest.param({}, {}, id="default"), | ||
| pytest.param({}, {"chunksize": 1}, id="chunked"), | ||
| pytest.param({}, {"index_col": "col_json_index"}, id="index"), | ||
| pytest.param({}, {"index_col": ["col_int", "col_json_index"]}, id="multi_index"), |
There was a problem hiding this comment.
Self-review round one — Finding 2: the converter rebuild for resolved column names (_configure_binary_csv_read(), pyathena/pandas/result_set.py:831) was not exercised; reverting it to self._converter.get() still passed. Repair: added multi_index and usecols params (usecols with a varbinary column takes the name-resolution path). Mutation check: reverting that line fails usecols, and dropping the MultiIndex handling fails multi_index. 6 passed with eff8530. reset_index() re-infers dtypes, so the test reads index levels directly.
| if d[1] in self._INTEGER_TYPES: | ||
| dtypes[d[0]] = self._converter.get_dtype(d[1], d[4], d[5]) | ||
| elif d[1] == "json": | ||
| dtypes[d[0]] = object |
There was a problem hiding this comment.
Self-review round one — checked and no change: pd.Series(values, dtype=get_dtype(...)) uses the same types mapping that read_csv(dtype=...) uses. A custom non-nullable integer dtype now raises on NULL on managed storage as on the S3 path (release note in the PR body). Duplicate column names still collapse in _rows_to_columnar(); that is the existing #1032, not changed here. Wrapper detection checks the first value only. The C and Python parsers call converters on every cell, including NA (''), so a json column is either fully wrapped or not converted.
| # and json values stay as decoded, so that NULL does not make them floats. | ||
| dtypes: dict[str, Any] = {} | ||
| for d in description: | ||
| if d[1] in self._INTEGER_TYPES: |
There was a problem hiding this comment.
Self-review round two (claims, callers, operations): FINDINGS. I corrected the PR body; the code is unchanged.
Scope: git diff a73c6db457110ea13eae8f7d1e7670b9095605dd..eff8530ca31909b48703275631794a7147348913.
Claims checked:
- Managed integers and JSON numbers with NULL were
float64, and S3-path JSON numbers too: measured on Athena before and after the change, for 22 types on both storage modes. - An all-integer JSON column was
int64: measured. - With pandas 3, all-string JSON was
str: offlineread_csv()andpd.DataFrame()checks. read_csv()re-infers converter values even withdtype=object: offline check, pandas 3.0.6.- The async and aio cursors share the result set:
pyathena/pandas/async_cursor.pyandpyathena/aio/pandas/cursor.pyimportAthenaPandasResultSet.
Finding (claim): the release note said a non-nullable custom integer dtype now fails on managed storage "as it already did" on the S3 path. Verified offline: managed raises a raw TypeError from pd.Series(...) here, while the S3 path raises OperationalError wrapping ValueError: Integer column has NA values. I corrected the PR body to state both exceptions. The managed path has never wrapped _fetch_all_rows() errors, so I left the exception type as is.
Callers: user converters/dtype kwargs still override PyAthena's (read_csv_kwargs.update(self._kwargs)), so they never see wrapped values. The pyarrow CSV engine is only chosen without converters. Repeated as_pandas() returns the same _df. Chunks are unwrapped one at a time. There are no AWS API changes: the same GetQueryResults calls and CSV reads.
Docs: docs/pandas.md and docs/null_handling.md state no json or integer dtype that this changes, so nothing became obsolete.
Evidence: published head eff8530, tests/pyathena/pandas/ + tests/pyathena/aio/pandas/: 300 passed locally (managed params included). The SQLAlchemy suites have not run yet; CI runs them on Ready.
There was a problem hiding this comment.
Rebase check: rebased a73c6db..05b08d9 onto master ec5323e, giving head cfc3976. git range-diff shows the four commits unchanged apart from test-file context: #1022's test_read_options and #1033's CONVERTED_VALUES_* imports landed at the same spot as test_integer_and_json_with_null, and I kept both sides.
Upstream contracts checked against this patch:
- Convert time zone values to aware times and empty TIME/JSON text to NULL #1033 switches the pandas
jsonconverter to_to_json(''→ None); the wrapper passes its result through unchanged. - Convert time zone values to aware times and empty TIME/JSON text to NULL #1033's new
time with time zone/timestamp with time zoneconverters are not json, so they aren't wrapped. _time_columnsis nowtimeonly, which the test'scol_timecovers.- Let execute() override cursor read options and accept them in the AsyncPandasCursor constructor #1022 only changes how read options reach the result set.
After the rebase: just lint passed, and 12 passed for -k "test_integer_and_json_with_null or test_integer_without_dtype or test_fetch_all_rows or test_read_options". The full pandas + aio pandas suite is running on cfc3976.
There was a problem hiding this comment.
Rebase check 2: rebased cfc3976..e25b9f0 onto master e49d576, giving head 0b77627.
#1054 (#1041) independently made PandasDataFrameIterator.get_chunk() apply the per-chunk function, so the conflict in that hunk resolves to upstream's code, which behaves the same. b7d0729 (was fbeb26a) now holds only the pd.array() change and the test, and I reworded its message. The get_chunk() release-note item moved out of this PR's body to #1054.
#1054's _trunc_date() now returns None for NULL time; the test's col_time has no NULL. Its _to_array parser changes don't touch json columns.
After the rebase: just lint passed, and 62 passed for -k "test_integer_and_json_with_null or test_integer_without_dtype or test_fetch_all_rows or test_complex_as_pandas or binary or get_chunk".
| df = self._reader.get_chunk(size) | ||
| else: | ||
| df = next(self._reader) | ||
| return self._trunc_date(df) |
There was a problem hiding this comment.
Independent review (relayed): Codex CLI 0.160.0, model gpt-6-astra, model_reasoning_effort=high, codex exec -s read-only, session 01a102ac-063d-7353-93a0-b918333a4bb8. Static review only: no edits, tests, or network. Snapshot: detached worktree at eff8530, base a73c6db, given the literal diff and the intended behavior, without PR text or prior findings. The snapshot was verified unchanged afterwards.
Reviewer coverage: managed integer/JSON construction (NULLs, empty results, precision, duplicate names); CSV whole-result and chunk access; index/MultiIndex; binary NULL preservation; converters/dtype/names/usecols overrides; C/Python engines and the PyArrow fallback; async/aio callers; pandas>=3.0; test assertions. Result: FINDINGS (2 × P2).
Finding 1 (P2): PandasDataFrameIterator.get_chunk() (was pyathena/pandas/result_set.py:199) never applied the per-chunk function. With chunksize=1, cursor.as_pandas().get_chunk(1) returned _CSVObject cells for JSON. Skipping the function predates this PR (time columns stayed timestamps there), but leaking the wrappers was introduced by it.
Verified by reading the code. Repaired in fbeb26a: get_chunk() applies the same function as __next__(). That also truncates time columns through get_chunk(), which fixes the pre-existing gap; I'll add it to the release notes. The test's chunked param now reads with get_chunk(1) followed by iteration and checks a time column. Reverting the repair fails [chunked].
| dtypes[d[0]] = object | ||
| return pd.DataFrame( | ||
| { | ||
| name: values if name not in dtypes else pd.array(values, dtype=dtypes[name]) |
There was a problem hiding this comment.
Independent review (relayed): Finding 2 (P2): with managed storage, SELECT 1 AS x, 2 AS x, 3 AS y gives x=[1, 2], y=[3] from _rows_to_columnar(), because the duplicate names collide (existing #1032). Building the columns as Series aligned them on the index, so the result had two rows with NA instead of raising ValueError as at the base.
Verified offline (pandas 3.0.6): dict of lists gives ValueError: All arrays must be of the same length; dict of Series gives a 2-row frame with NaN. Repaired in fbeb26a: typed columns use pd.array(values, dtype=...), which is not aligned and raises the same ValueError. Dtype results are unchanged (Int64, object, float, str; int64 with NULL raises the same TypeError). I didn't add a test, because one would pin the #1032 ValueError, which that issue will replace.
There was a problem hiding this comment.
Independent follow-up (relayed): Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a102b6-373c-7f43-a351-549a16767249. Static review of the delta eff8530..fbeb26a on a detached snapshot at fbeb26a. Coverage: get_chunk()/__next__()/collection and sync/async callers (CSV chunks finished once; non-chunked, UNLOAD, and managed use the no-op function); unequal column lengths raise ValueError again; pd.array() vs pd.Series() dtypes. Result: FINDINGS (1 × P2).
Finding: when a custom converter's types omits an integer type, get_dtype() returns None. pd.array([1, None], dtype=None) infers Int64, where pd.Series() and read_csv() (S3 path) infer float64/NaN. Verified. Repaired in 05b08d9: integer columns without a converter dtype keep their list and pandas' inference. New test_integer_without_dtype (default + managed, converter without a bigint dtype) expects float64 and NaN on both. Recording None again fails [managed].
Declined: the reviewer also suggested a managed duplicate-name test. The ValueError it would assert is the defect tracked in #1032, which will replace that error, so I didn't pin it.
Self-review of the repairs (fbeb26a, 05b08d9), round one: get_chunk's callback is _no_trunc_date except for CSV chunks; explicit dtypes (Int64, object, float, str) are unchanged under pd.array(). Round two: the commit messages' claims (ValueError restored, float64 inference) are checked offline and by the test; the PR body is updated with the get_chunk() time truncation release note.
There was a problem hiding this comment.
Independent follow-up 2 (relayed): Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a102c0-e95f-70a0-a6c4-c4fd0ea6bfdb. Static review of the delta fbeb26a..05b08d9 on a detached snapshot at 05b08d9. Coverage: API/CSV dtype handling, converter propagation, regression-test assertions, helper and fixture state isolation. Result: CLEAN. Omitting None dtypes restores float64/NaN inference, and explicit dtypes and JSON handling are unchanged. The managed case would fail against the previous Int64 result. Each helper call builds its own converter, so removing bigint doesn't leak to other tests.
05b08d9 to
cfc3976
Compare
| self.value = value | ||
|
|
||
| @classmethod | ||
| def wrap(cls, converter: Callable[[str | None], Any]) -> Callable[[str | None], _CSVObject]: |
There was a problem hiding this comment.
Refactor e25b9f0 (maintainer request: gather the CSV wrapping helpers): _convert_csv_object, _unwrap_csv_objects, and the column/index unwrapping that was inline in _finish_csv_frame() become _CSVObject.wrap(), unwrap(), and _unwrap_values(). They are classmethods, not staticmethods, because each one uses cls (constructor, isinstance). The maintainer chose to keep the wrapper over reading json as text and decoding afterwards, or dropping the S3-path JSON fix.
Self-review, round one: a pure move; the closure replaces functools.partial with the same call and exception behavior, and partial is still used for varbinary.
Round two: the commit's "behavior is unchanged" claim was checked by diff inspection and by -k "test_integer_and_json_with_null or test_integer_without_dtype or test_fetch_all_rows or test_complex_as_pandas or binary": 58 passed at e25b9f0. The class docstring's inference claim matches the earlier measurement.
Independent follow-up 3 (relayed): Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a1044b-4454-7940-bb0e-f25bf9e1a1e1. Static review of cfc3976..e25b9f0 on a detached snapshot at e25b9f0. Coverage: C/Python converters, whole/chunked/get_chunk() reads, varbinary wrapping, Index/MultiIndex, docstrings, structure. Result: CLEAN.
With managed query result storage, integer columns with NULL became float64, losing integers above 2**53, because the DataFrame was built from GetQueryResults rows with inferred dtypes. They now get the dtype that the CSV result file reads them with (Int64 by default), and json columns stay object. In the CSV result file, pandas inferred a numeric dtype from the values that the json converter returned, so JSON numbers with NULL became float64 on that path too. The json converter's values are now wrapped during read_csv() and unwrapped into an object column afterwards, for whole results, chunks, and an index column. Closes #1027 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
An index_col list with a json column left the wrapped values in its MultiIndex level. The regression test also covers usecols, which rebuilds the converters in _configure_binary_csv_read(), and a MultiIndex. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Building the managed columns as Series aligned them on their index, so a duplicate column name, whose values _rows_to_columnar() collects into one list, filled extra rows with NA instead of raising ValueError as before. The typed columns are now pandas arrays, which are not aligned. The chunked test reads the first chunk with get_chunk(), which unwraps JSON values and truncates time columns as iteration does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
pd.array() infers Int64 for integers with NULL where the CSV result file and the earlier pd.Series() infer float64, so integer columns without a converter dtype keep their list and pandas' inference. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The converter wrapper and the unwrapping of columns and index levels become _CSVObject.wrap() and _CSVObject.unwrap(), so that the class holds everything about values that pandas.read_csv() must not infer a dtype from. Behavior is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
e25b9f0 to
0b77627
Compare
_CSVObject.unwrap() selected the columns of every frame through iloc, which made results without JSON 24% slower to read in 1,000-row chunks of 200 columns; it now reads df.dtypes first and touches only object columns. Iterating values.array was also several times slower than the object ndarray from to_numpy() (172 ms vs 39 ms per million values). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…each value Wrapping every converted json value in an object made json columns about 30% slower to read, half of it garbage collection. pandas only turns JSON numbers into floats when the values mix numbers and None, so the json converter now returns one shared placeholder in place of None, which keeps the column object, and _finish_csv_frame() puts None back in the columns or index levels whose converter returned it. Reading json columns now costs within noise of master. json columns without NULL keep the dtype that pandas infers, as on master, and the managed path makes json columns object only when they have NULL, to match. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| return df | ||
|
|
||
|
|
||
| class _JSONConverter: |
There was a problem hiding this comment.
Design change (maintainer decision): _CSVObject → _JSONConverter (NULL placeholder), 88cd9cd.
Measured: wrapping every json value cost about 30% on json columns, half of it GC (+35% with GC on, +18% off). Fable 5.1 and Codex were consulted read-only on the alternatives:
| Design | Measured cost on json columns |
|---|---|
| Decode raw strings after reading | +8–32% |
| Wrap only scalars | +11% on objects, +18% on unrelated object columns (full scans) |
| Per-column value buffers (Codex) | +15–25% |
| NULL placeholder (Fable) | about noise |
Both reviewers agree pandas has no option that stops inference on converter output. The maintainer chose the placeholder.
Self-review, round one (git diff e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1..88cd9cde004bf73a7d3ec05aa7e461ac4a702b1f): CLEAN.
- Converter instances live per result set, created in
_get_csv_converter()._csv_converterstakes the finalread_csv_kwargs["converters"]: after_configure_binary_csv_read()(which rebuilds converters on name resolution), and after userconverterskwargs, which replace ours, so nothing is restored then. - The flag is reset in
restore(), which_finish_csv_frame()runs right after each frame or chunk (__next__andget_chunk()). - JSON
nullliterals give None from_to_jsonand are handled like NULL. usecolsthat drops a json column never calls its converter.- The python engine never converts
index_colcolumns (pre-existing), so their flag stays unset. - Mutations: returning None instead of the placeholder fails 4 params; skipping the restore fails 5; skipping the index restore fails
indexandmulti_index; making every managed json columnobjectfailsmanaged.
Round two (claims): the commit message's "within noise of master" holds for the combined benchmarks (−1.1%, +3.3%, +2.3%), but single-column cases measured up to +13.8% (numbers without NULL) and +16.4% (objects with NULL in 50k chunks). The PR body states per-value costs instead (0.06 µs per call, 0.04 µs per restored value).
New user-visible effect, now in the release notes: large low_memory reads with sparse NULL emit DtypeWarning with exact values. Reproduced offline: 600,000 rows, one NULL at the end. Before, the column was silently float64.
Evidence: 317 passed on tests/pyathena/pandas/ + tests/pyathena/aio/pandas/ locally; CI test/run (3.14) passed on 88cd9cd.
| if d[1] in self._INTEGER_TYPES: | ||
| if (dtype := self._converter.get_dtype(d[1], d[4], d[5])) is not None: | ||
| dtypes[d[0]] = dtype | ||
| elif d[1] == "json" and None in columnar[d[0]]: |
There was a problem hiding this comment.
Independent review (relayed): Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a104a0-9478-7a22-87d7-b77f10273eda. Full static review of e49d576..88cd9cd, run because the design changed, on a detached snapshot at 88cd9cd. The prompt gave the intended behavior and accepted limits, without PR text or prior findings.
Coverage:
- managed dtypes, SQL NULL and JSON
null, and empty results; - column and index restore;
- iteration,
get_chunk(), and iteratoras_pandas(); converters/dtype/names/usecolsoverrides, engines,low_memory, and varbinary;- async/aio, UNLOAD, converter lifetime, performance structure, and tests.
Result: one P2, pre-existing. On managed storage, SELECT 1 AS x, 2 AS x puts both values in one column, through _rows_to_columnar() (pyathena/result_set.py:764). That is #1032, unchanged by this PR; my prompt overstated a duplicate-name guarantee. No new regressions found.
WHAT
PandasCursor(and the async and aio pandas cursors, which shareAthenaPandasResultSet) keeps integer and JSON values exact when a column has NULL._as_pandas_from_api()):tinyint,smallint,integer, andbigintcolumns get the dtype that the pandas converter gives the CSV result file (Int64by default; pandas infers it when the converter'stypeshas none).jsoncolumns with NULL areobject. Before, pandas inferredfloat64from theGetQueryResultsvalues when a column had NULL, so integers above 2**53 lost precision and NULL becameNaN. The typed columns are pandas arrays, which are not aligned on an index.read_csv()): pandas infers a dtype from the values that a converter returns, so JSON numbers with NULL becamefloat64. json columns now use_JSONConverter, which returns one shared placeholder (_JSONConverter.NULL) in place of None, so the column staysobject._finish_csv_frame()puts None back in the columns or index levels whose converter returned it, for the whole result and for each chunk (iteration andget_chunk()).Cost: one extra Python call per json value (about 0.06 µs per value), plus a restore pass (about 0.04 µs per value) for frames or chunks whose json column has NULL. Results without json columns don't change. An earlier revision wrapped every json value in an object, which made json columns about 30% slower to read; this design replaced it.
Release notes (behavior changes)
objectdtype and keep their decoded values, withNonefor NULL. Before, JSON numbers with NULL werefloat64and NULL wasNaN. JSONnullliterals count as NULL.Int64(or the dtype that a custom converter'stypesgives them), with<NA>for NULL, as with an S3 result file. Without a dtype intypes, pandas infers it as before. A custom converter whosetypesmaps an integer type to a non-nullable dtype such asint64now fails on NULL on managed storage:pandas.Series()raisesTypeError. With an S3 result file, the same converter already failed withOperationalErrorwrapping pandas'ValueError: Integer column has NA values.pandas.read_csv()call can now emit pandas'DtypeWarning(mixed types) when its json numbers have sparse NULL; the values are exact. This applies when the read is unchunked or chunks are larger than pandas' internal parser block (the largest power of two below 2**20 divided by the column count, for example 524,288 rows for one column and 65,536 for ten). Before, the column was silentlyfloat64. Passlow_memory=Falsetoexecute()to avoid it. Verified offline: 600,000 rows with one NULL at the end giveDtypeWarningand exact values, and withlow_memory=Falseno warning.Unchanged limits
These behave as on master:
int64orstr).float64.NaN.index_colcolumns.WHY
Closes #1027.
The maintainer chose this scope on the issue: integers and JSON only on the managed path, and the S3-path JSON precision fix in the same PR. For the S3 path, the maintainer chose the NULL placeholder design after Fable 5.1 and Codex compared the alternatives. Wrapping every value, decoding raw strings after reading, and per-column buffers each cost 15–30% on json columns.
The remaining managed-path differences for
date,array,map, androware filed as #1042, and custom converters on the managed path are #1028.TEST
Tested commit: 88cd9cd (base e49d576).
just lint: passed.TestPandasCursor.test_integer_and_json_with_nullchecks default, chunked (get_chunk()then iteration),index_coland MultiIndexindex_colon the json column with NULL,usecols, and managed.test_integer_without_dtypechecks default and managed with a converter without abigintdtype. Both: 8 passed.index,multi_index;object:managed;test_integer_without_dtype[managed].uv run --env-file .env pytest -n 2 tests/pyathena/pandas/ tests/pyathena/aio/pandas/: 317 passed (includes the managed work group params).test / run (3.14)passed on 88cd9cd.read_csv()infersfloat64from converter values even withdtype=object; a shared placeholder keeps the columnobject.Not run locally: the SQLAlchemy suites and the rest of
just test pyathena. CI skips the SQLAlchemy jobs by path filter, because this PR doesn't change SQLAlchemy modules.🤖 Generated with Claude Code