Skip to content

Derive the one_row_complex expectations of the pandas, Arrow, and Polars cursors - #876

Closed
laughingman7743 wants to merge 2 commits into
test/848-derived-expectationsfrom
test/848-derived-expectations-dataframes
Closed

laughingman7743 wants to merge 2 commits into
test/848-derived-expectationsfrom
test/848-derived-expectations-dataframes

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Sep 28, 2026 •

Copy link
Copy Markdown
Member

WHAT

tests/pyathena/expected.py gains the value rules and type tables of the pandas, Arrow, and Polars cursors. With them, the complex-row tests of those cursors derive their expectations from the ONE_ROW_COMPLEX definition.

Value rules (keyed by base type):

Rules Used by Array, map, and struct values
PANDAS_VALUES PandasCursor rows and DataFrames Athena's text rendering ([1, 2], {1=2, 3=4}, {a=1, b=2}); timestamps and dates as pd.Timestamp
ARROW_VALUES ArrowCursor rows text rendering
ARROW_TABLE_VALUES ArrowCursor as_arrow() and as_polars() text rendering; time, binary (31 32 33), JSON, and decimal as the CSV reader's strings; dates as datetimes
UNLOAD_VALUES pandas and Arrow cursors with unload=True lists, key-value tuples, and dicts; timestamps as pd.Timestamp
UNLOAD_POLARS_VALUES ArrowCursor as_polars() with unload=True maps as {"key": ..., "value": ...} lists
PYTHON_VALUES (existing) PolarsCursor rows not selected: exclude_types=POLARS_EXCLUDED_TYPES leaves out binary and complex columns, as the tests did before

Types and helpers on ExpectedResult:

  • pandas_types(unload): the pandas dtype.type of each column.
  • arrow_schema(unload): the ArrowCursor schema. For UNLOAD, the types of nested columns are built from the nested Athena types.
  • polars_types(): the PolarsCursor schema.
  • description(unload=True): every column NULLABLE, and varchar without its length.
  • with_array_lists(row): turns the NumPy arrays in the array columns of pandas UNLOAD rows into lists.

Other changes:

  • polars_type(arrow_type) gives the Polars type of an Arrow type.
  • STRING_TYPE moves from pandas/test_cursor.py into expected.py.
  • Table.data_file in tables.py now rejects a row whose values change in the columns' Arrow types, for example a timestamp more precise than the type, or a struct missing some of its fields. The expectations are derived from the rows as written, so such a row would otherwise produce a false failure.

Rewritten tests (14):

  • pandas/test_cursor.py: test_complex, test_complex_unload_pyarrow, test_complex_as_pandas, test_complex_unload_as_pandas_pyarrow
  • arrow/test_cursor.py: test_complex, test_complex_unload, test_complex_as_arrow, test_complex_unload_as_arrow, test_complex_as_polars, test_complex_unload_as_polars
  • polars/test_cursor.py: test_complex, test_complex_unload, test_complex_as_polars, test_complex_unload_as_polars

Assertion changes:

  • The CAST columns follow the table columns, as in Derive the one_row_complex expectations of the Python-object cursors #875. Every value, dtype, schema field, and description tuple is still asserted.
  • The shape assertions become column-name assertions plus row-list equality. Together they cover the width and the row count.
  • The pandas dtype and ArrowCursor as_polars() dtype checks now also compare the column names, in order.
  • test_complex_unload_as_polars of PolarsCursor now also asserts the column names.

WHY

Part of #848 (step 2), for #834. Stacked on #875.

This covers the remaining tests that spelled out one_row_complex by hand. After this PR, every test that asserts a whole one_row_complex row, schema, dtype list, or description derives it from the definition. A column added only through the definition follows in the next PR.

TEST

Tested commit: 0afa831. At the maintainer's request, this PR was rebuilt as one commit on the simplified #875 (renamed helpers, no nested-type modeling) after the independent reviews below. 0afa831 fixes the refactor review's finding on nulls in pandas UNLOAD arrays.

  • just lint passed.

  • Offline, a script outside the repository compared the derived expectations with the removed literals, after moving the CAST columns last. It covered 14 items and found no difference in any of them:

    • pandas CSV rows
    • UNLOAD rows
    • Arrow CSV rows
    • Arrow table rows
    • UNLOAD Polars dicts
    • PolarsCursor rows
    • pandas CSV and UNLOAD dtypes
    • Arrow CSV and UNLOAD schemas
    • Arrow CSV and UNLOAD Polars dtypes
    • PolarsCursor schema
    • the UNLOAD varchar description

    Every row value has the same Python type as the removed literal.

  • Against AWS, from this worktree, pytest -n 4 tests/pyathena/pandas/test_cursor.py tests/pyathena/arrow/test_cursor.py tests/pyathena/polars/test_cursor.py -k test_complex gave 20 passed, on each pushed revision including the rebuilt 9e350b7 and 68270a4 (the pandas tests again on 0afa831). That is all 14 tests with their parameter sets.

  • Not run locally: the full suite. AWS CI runs the PyAthena suite when the PR is Ready. The changes job skips the SQLAlchemy compliance suites and the Spark tests, because no SQLAlchemy or Spark path changes here.

🤖 Generated with Claude Code

@laughingman7743
laughingman7743 force-pushed the test/848-derived-expectations-dataframes branch 2 times, most recently from 02fc024 to 9242a15 Compare September 28, 2026 10:41
}


def arrow_type(athena_type: str, unload: bool = False) -> pa.DataType:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one (implementation behavior): CLEAN

Base 4e205a3 (#875 head, the stacked base), head 9242a15. I reviewed the full diff of the 4 files; the first draft was ab0652f and the later heads are rebases.

  • Assertion strength: every removed row, dtype, Arrow schema, Polars dtype, and description literal has a derived counterpart.
    • An offline comparison of 14 items against the removed literals, with the CAST columns moved last, found no difference.
    • Every row value has the same Python type.
    • The shape assertions are replaced by name lists, schema equality, and whole-row-list equality, so the width and row count are still asserted.
  • UNLOAD:
    • description(unload=True) gives NULLABLE and treats varchar as unbounded string. decimal(10,1) keeps its parameters.
    • comparable only converts np.ndarray. The pandas UNLOAD map value is already a list of tuples, as the removed literal asserted.
  • Types:
    • The nested UNLOAD Arrow types are built recursively from the nested Athena types (array_element, entries field names as before).
    • polars_type converts Arrow types explicitly and does not call pl.from_arrow.
    • The CSV path types come from explicit family tables.
  • Non-circularity: expected.py still imports no PyAthena converter.
  • Scope: the PolarsCursor tests keep the 12 scalar columns through exclude=COMPLEX_FAMILIES. A new binary or complex column is excluded there automatically, as before, when those columns were hand-picked.
  • Test evidence: -k test_complex over the three files gives 20 passed on AWS at 9242a15.

return _ARROW_UNLOAD_TYPES[name]


def polars_type(arrow: pa.DataType) -> pl.DataType:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims, callers, AWS operations): FINDINGS (description only, fixed)

Base 4e205a3, head 9242a15.

Claims checked:

  • "20 passed = all 14 tests with their parameter sets": pandas 3+2+3+2, arrow 6, polars 4.
  • "Every test that asserts a whole one_row_complex row ... derives it": the remaining one_row_complex uses assert single columns or type-hinted subsets, such as test_complex_with_type_hints.
  • CI scope: no path under the sqla or spark filters changed, so AWS CI runs the PyAthena suite without Spark.

Finding: the description said the old PolarsCursor schema check was "a dict comparison ... which ignores order". That is false. pl.Schema subclasses OrderedDict, and I measured pl.DataFrame({'a':[1],'b':[2]}).schema == {'b': pl.Int64, 'a': pl.Int64} to be False. The claim is removed. The PolarsCursor check is now described as a name-and-type comparison, equivalent to the old one.

@laughingman7743
laughingman7743 force-pushed the test/848-derived-expectations-dataframes branch from 9242a15 to 590af43 Compare September 28, 2026 10:50
Comment thread tests/pyathena/expected.py Outdated
Returns:
The row as a tuple.
"""
return tuple(list(v) if isinstance(v, np.ndarray) else v for v in row)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review (relayed Codex result): FINDINGS (1 × P3)

  • Reviewer: Codex CLI 0.157.1 (codex exec -s read-only), model gpt-6-astra, session 01a0e79c-9591-76b0-a841-ce4c0578c308.
  • Scope: base 4e205a3, head 9242a15. The review ran on a detached snapshot without .env and without the PR's framing. Afterwards the snapshot was unchanged and still at the head. It was a static review only.
  • Covered, as reported: all four changed files, comparing the following against the removed literals and ONE_ROW_COMPLEX:
    • CSV and UNLOAD rows
    • pandas dtypes
    • Arrow schemas
    • Polars dtypes
    • descriptions
    • dimensions
    • helper independence from the converters

P3 — Map container-type coverage weakened. comparable() normalizes every NumPy array cell; previously only col_array was normalized. If pandas UNLOAD returned col_map as a one-dimensional object ndarray of tuples, both rewritten pandas tests would accept it as [(1, 2), (3, 4)]. The removed assertions would reject that. Restrict normalization to Athena array columns.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repair in 590af43: verified. The module-level comparable is replaced by Selection.comparable(row), which converts a NumPy array to a list only in columns whose Athena type family is array. A map column returned as an ndarray would therefore stay an ndarray and fail the list-of-tuples comparison, as the removed assertions would. Both pandas UNLOAD tests use the new method. The offline comparison still shows 14/14 equal, and the 20 complex-row tests pass on AWS at 590af43. The PR was also rebased onto the latest #875 head (756bdd1). That base change touches only the PYTHON nested-value rules, which these tests use for the PolarsCursor scalars alone.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up (relayed Codex result): CLEAN

  • Reviewer: Codex CLI 0.157.1, model gpt-6-astra, session 01a0e7a3-e872-72b3-8da9-3c7a1586f2f2.
  • Scope: a read-only, detached snapshot at 590af43. The review covered the range-diff of the rebase onto 756bdd1 and the repair diff 13785fb..590af43. It was static only.
  • Reported:
    • The P3 is resolved: only Athena array columns normalize ndarrays, so an ndarray map cell stays unnormalized and fails the comparison.
    • The range-diff shows the rebase left the original patch unchanged, and the new base's rendering and nesting rules remain compatible with the selected fixture types.
    • No new defects were found.

@laughingman7743
laughingman7743 force-pushed the test/848-derived-expectations-dataframes branch from 590af43 to 5f16afb Compare September 28, 2026 10:54
Comment thread tests/pyathena/expected.py Outdated
_TEXT: Representation = dict.fromkeys(("array", "map", "struct"), _athena_text)


def _native(value: Any, athena_type: str) -> Any:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review of a repair (both perspectives): FINDINGS (1, fixed by the author)

Base and range: the patch is now on the #875 head c0b2e3c. The repair commit is 5f16afb.

Round one (behavior)

  • The Derive the one_row_complex expectations of the Python-object cursors #875 follow-up 8 found that PYTHON struct rules ignored declared fields that a value omits. The same flaw exists in this PR's _NATIVE rule: dict(v) omits a declared field that Parquet stores as null. A STRUCT<a: int, b: int> value {"a": 1} would be expected as {"a": 1}, while pyarrow returns {"a": 1, "b": None}.
  • The repair adds _native. It walks the declared fields through _struct_items (from the new base), handles arrays of structs, and returns None for a null value. It applies the same _check_nesting limits as the PYTHON rules. UNLOAD_POLARS maps go through _polars_map, which has the same checks.
  • Checked offline against a pyarrow round trip of the same values: a missing struct field, an array of structs containing a null element, a map, an array with a null, and a null map all match.
  • The 14-item comparison against the removed literals is still 14/14 equal, and the 20 complex-row tests pass on AWS at 5f16afb.

Round two (claims)

  • The description's UNLOAD row now states the declared-fields and nesting behavior.
  • The tested commit is updated.
  • No other claim changes.

Comment thread tests/pyathena/expected.py Outdated
name = family(athena_type)
if name == "struct":
return {k: x for k, x, _ in _struct_items(value, athena_type)}
if name == "array" and family(_type_arguments(athena_type)[0]) == "struct":

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 2 (relayed Codex result): FINDINGS (2 × P2)

  • Reviewer: Codex CLI 0.157.1, model gpt-6-astra, session 01a0e7a7-a3b0-7080-af8d-b39c24244ad8.
  • Scope: a read-only, detached snapshot at 5f16afb. The review covered the range-diff of the rebase onto c0b2e3c and the repair 95cf6b0..5f16afb. It was a static review only.
  • Confirmed: both rebased commits are patch-identical. Missing fields, nulls, and arrays of structs are handled correctly, and the one_row_complex expectations are unchanged.

P2 — Arrays of maps bypass Polars conversion. _check_nesting accepts ARRAY<MAP<int,int>>, but _native only recurses for structs. For [[(1, 2)]], UNLOAD_POLARS returns [[(1, 2)]], while converting the Arrow column to Polars gives [[{"key": 1, "value": 2}]].

P2 — Nested scalar values retain precision discarded by Arrow serialization. For STRUCT<t: timestamp> with Arrow type pa.struct([("t", pa.timestamp("ms"))]), a value with microsecond=123456 passes the nesting check. Table.data_file stores 123000, while _native keeps 123456. Normalize the supported scalar values, or reject values the rules cannot preserve.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repair in 3202ab5: both findings verified.

  1. _native now takes a polars flag. A map becomes {"key", "value"} dicts at any position the nesting check allows, including maps nested in arrays. UNLOAD_POLARS uses the flag for arrays, maps, and structs, which replaces _polars_map. Offline, [[(1, 2)]] for ARRAY<MAP<int, int>> gives [[{"key": 1, "value": 2}]], equal to pl.from_arrow of the same Arrow column.
  2. This problem is not limited to nested values: a top-level timestamp column with microseconds would drift the same way. So the fix is in the table definition, not in each rule. Table.data_file now compares table.to_pylist() with the rows as given, and raises ValueError naming any column whose values change in their Arrow type. That covers extra timestamp precision and also a struct that omits declared fields. Every table in TABLES passes. Offline, a microsecond timestamp in a timestamp("ms") column and {"a": 1} in STRUCT<a: int, b: int> are both rejected.

The 14-item comparison against the removed literals is still 14/14 equal. The 20 complex-row tests pass on AWS at 3202ab5.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 3 (relayed Codex result): CLEAN

  • Reviewer: Codex CLI 0.157.1, model gpt-6-astra, session 01a0e7af-5a58-7b62-813e-eb5b416746f7.
  • Scope: a read-only review of a detached snapshot at 3202ab5, covering the range 5f16afb..3202ab5. It was a static review only, and the snapshot was unchanged afterwards.
  • Reported:
    • Finding 1 is resolved: Polars maps convert to key/value dicts at every permitted nesting position, including inside arrays.
    • Finding 2 is resolved: the Arrow round-trip comparison raises ValueError for precision loss and for missing struct fields.
    • The existing TABLES pass the check, and the current one_row_complex expectations are unchanged.
    • No new defects.

@laughingman7743
laughingman7743 marked this pull request as ready for review September 28, 2026 11:07
@laughingman7743
laughingman7743 force-pushed the test/848-derived-expectations-dataframes branch 2 times, most recently from 9e350b7 to 7845405 Compare September 28, 2026 15:01
…ars cursors

The complex-row tests of PandasCursor, ArrowCursor, and PolarsCursor spelled
out the one_row_complex rows, descriptions, pandas dtypes, Arrow schemas,
and Polars dtypes by hand. They now derive them from the table definition
with the value rules and type tables in tests/pyathena/expected.py: Athena's
text rendering for the CSV paths, native values for the UNLOAD paths, and
explicit pandas, Arrow, and Polars type tables. The PolarsCursor tests keep
leaving out the binary and complex columns, now by type, and also assert
the column names.

Table.data_file rejects a row whose values change in the columns' Arrow
types, such as a timestamp more precise than the type or a struct without
all of its fields, because the expectations are derived from the rows as
written.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
pandas reads a null in an integer array from Parquet as NaN, so
with_array_lists turns NaN in array columns into None, as the derived
expectation has it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
for row in self.table.rows
]

def with_array_lists(self, row: Any) -> tuple[Any, ...]:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review of the refactor (relayed Codex result): FINDINGS (1 × P2), repaired in 0afa831

Reviewer: Codex CLI 0.157.1, model gpt-6-astra, session 01a0e88f-6cdd-75d3-8815-7a5fdf10fe03. Read-only static review of a detached snapshot at 68270a4, covering the #875 repair 0b4d138..36988f4 and confirming that this PR's commit on top is a pure rebase of 0b4d138..7845405, with identical additions and deletions.

The reviewer reported that both refactor findings (nested nulls, struct field order) are resolved in the parsing and CSV rules, and that the current one_row_complex values and types are unchanged.

P2 — pandas UNLOAD still mishandles nullable integer arrays. For an ARRAY<int> containing [1, None], the pandas/PyArrow Parquet path returns a NumPy array containing [1.0, NaN]. with_array_lists() only converts the container and keeps NaN, while UNLOAD_VALUES expects [1, None].

Repair (verified): with_array_lists now converts NaN to None inside array columns. 1.0 == 1 already holds in the comparison. Offline, np.array([1.0, nan]) for an ARRAY<int> row [1, None] compares equal to the derived expectation. The pandas complex-row tests pass on AWS at 0afa831 (10 passed). #877 and #878 were rebased onto it.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 2 of the refactor (relayed Codex result): CLEAN

  • Reviewer: Codex CLI 0.157.1, model gpt-6-astra, session 01a0e894-7809-7ee2-9f7a-ed6fa11b9c06.
  • Scope: a read-only review of a detached snapshot at 0afa831, covering 68270a4..0afa831. It was static only.
  • Result:
    • The finding is resolved: [1.0, NaN] becomes [1.0, None], which matches [1, None]. Integers, strings, and existing None elements are unchanged.
    • Full-row comparisons, array length and order checks, and the existing metadata assertions keep their strength.
    • No new defect was found.

@laughingman7743

Copy link
Copy Markdown
Member Author

Closing without merging, by the maintainer's decision. Deriving the expected results from the table definitions required tests/pyathena/expected.py to re-model how each cursor returns every type: Athena's text rendering, JSON parsing, and the pandas, Arrow, and Polars types. The module grew large enough to need tests of its own, which is more complexity than the one-place column addition is worth. The hand-written expectations stay. See #848 for the summary. The branch is kept for reference.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant