Skip to content

Remove and rewrite unnecessary Polars, Arrow, and S3FS cursor tests #1252

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 8 commits into from
Oct 18, 2026
Merged

Remove and rewrite unnecessary Polars, Arrow, and S3FS cursor tests #1252

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ff8012b
Remove and rewrite redundant Polars cursor tests
laughingman7743 Oct 17, 2026
9529eb2
Remove and rewrite redundant Arrow cursor tests
laughingman7743 Oct 17, 2026
d894953
Remove and rewrite redundant S3FS cursor tests
laughingman7743 Oct 17, 2026
5a8b810
Check that storage options reach chunked CSV reads and that fetches stop
laughingman7743 Oct 17, 2026
6b8f9ea
Read the local Polars results through GetQueryResults and as_polars()
laughingman7743 Oct 17, 2026
738069e
Check the read options and conversions that the removed tests protected
laughingman7743 Oct 17, 2026
558f368
Test the cleanup and routing that the follow-up review found untested
laughingman7743 Oct 17, 2026
5b64fc7
Read the S3FS test metadata through GetQueryResults
laughingman7743 Oct 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 20 additions & 223 deletions tests/pyathena/arrow/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from unittest.mock import MagicMock, patch

import pandas as pd
import polars as pl
import pyarrow as pa
import pytest

Expand Down Expand Up @@ -566,186 +565,11 @@ def test_complex_unload_as_arrow(self, arrow_cursor):
)
]

@pytest.mark.parametrize(
"arrow_cursor",
[{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}],
indirect=["arrow_cursor"],
)
def test_as_polars(self, arrow_cursor):

@laughingman7743 laughingman7743 Oct 17, 2026

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 CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a149cf-467f-78a3-9798-57123cbadabe. Static review of 713d57bde5248f9a77b763b44eb4c6945ceb4361..6b8f9ea66f66a752bba7cb832ab6e9b5a357a6d4 in a clean detached snapshot, with the diff, the docs/testing.md "Choose what to test" section, and no PR text or author findings; no edits, builds, tests, network, or GitHub access. Covered: all 11 changed files, the removed tests at the base, the cursor and result-set code, remaining tests under tests/pyathena/, the three changed helpers, live-query expectations, IDs, marks, fixture scopes, and cleanup. Result: FINDINGS (7).

Finding 1 (P2): with test_many_as_polars removed, only single-row results exercise AthenaArrowResultSet.as_polars(); pl.from_arrow(self._table.slice(0, 1)) at pyathena/arrow/result_set.py:448 would pass. The removed complex tests were the only multi-column inputs.

@laughingman7743 laughingman7743 Oct 17, 2026

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.

Repaired in 738069e: test_as_polars reads VALUES (1, 'a'), (2, 'b') (two rows, two columns, one query) and compares to_dicts(). A truncating as_polars() now fails it. Not yet run against Athena.

df = arrow_cursor.execute("SELECT * FROM one_row").as_polars()
assert df.height == 1
assert df.width == 1
assert df.to_dicts() == [{"number_of_rows": 1}]

@pytest.mark.parametrize(
"arrow_cursor",
[{"cursor_kwargs": {"unload": False}}, {"cursor_kwargs": {"unload": True}}],
indirect=["arrow_cursor"],
)
def test_many_as_polars(self, arrow_cursor):
df = arrow_cursor.execute("SELECT * FROM many_rows").as_polars()
assert df.height == 10000
assert df.width == 1
assert df.to_dicts() == [{"a": i} for i in range(10000)]

def test_complex_as_polars(self, arrow_cursor):
df = arrow_cursor.execute(
"""
SELECT
col_boolean
,col_tinyint
,col_smallint
,col_int
,col_bigint
,col_float
,col_double
,col_string
,col_varchar
,col_timestamp
,CAST(col_timestamp AS time) AS col_time
,col_date
,col_binary
,col_array
,CAST(col_array AS json) AS col_array_json
,col_map
,CAST(col_map AS json) AS col_map_json
,col_struct
,col_decimal
FROM one_row_complex
"""
).as_polars()
assert df.height == 1
assert df.width == 19
dtypes = tuple(df.dtypes)
assert dtypes == (
pl.Boolean,
pl.Int8,
pl.Int16,
pl.Int32,
pl.Int64,
pl.Float32,
pl.Float64,
pl.String,
pl.String,
pl.Datetime("us"),
pl.String,
pl.Date,
pl.String,
pl.String,
pl.String,
pl.String,
pl.String,
pl.String,
pl.String,
)
rows = df.to_dicts()
assert rows == [
{
"col_boolean": True,
"col_tinyint": 127,
"col_smallint": 32767,
"col_int": 2147483647,
"col_bigint": 9223372036854775807,
"col_float": 0.5,
"col_double": 0.25,
"col_string": "a string",
"col_varchar": "varchar",
"col_timestamp": datetime(2017, 1, 1, 0, 0, 0),
"col_time": "00:00:00.000",
"col_date": datetime(2017, 1, 2).date(),
"col_binary": "31 32 33",
"col_array": "[1, 2]",
"col_array_json": "[1,2]",
"col_map": "{1=2, 3=4}",
"col_map_json": '{"1":2,"3":4}',
"col_struct": "{a=1, b=2}",
"col_decimal": "0.1",
}
]

@pytest.mark.parametrize(
"arrow_cursor",
[
{
"cursor_kwargs": {"unload": True},
},
],
indirect=["arrow_cursor"],
)
def test_complex_unload_as_polars(self, arrow_cursor):
# NOT_SUPPORTED: Unsupported Hive type: time
# NOT_SUPPORTED: Unsupported Hive type: json
df = arrow_cursor.execute(
"""
SELECT
col_boolean
,col_tinyint
,col_smallint
,col_int
,col_bigint
,col_float
,col_double
,col_string
,col_varchar
,col_timestamp
,col_date
,col_binary
,col_array
,col_map
,col_struct
,col_decimal
FROM one_row_complex
"""
"SELECT * FROM (VALUES (1, 'a'), (2, 'b')) AS t(k, s) ORDER BY k"
).as_polars()
assert df.height == 1
assert df.width == 16
# Polars 2.0 converts an Arrow map to Map; Polars 1.x to a list of key-value structs.
if hasattr(pl, "Map"):
map_dtype = pl.Map(pl.Int32, pl.Int32)
map_value = {1: 2, 3: 4}
else:
map_dtype = pl.List(pl.Struct([pl.Field("key", pl.Int32), pl.Field("value", pl.Int32)]))
map_value = [{"key": 1, "value": 2}, {"key": 3, "value": 4}]
dtypes = tuple(df.dtypes)
assert dtypes == (
pl.Boolean,
pl.Int8,
pl.Int16,
pl.Int32,
pl.Int64,
pl.Float32,
pl.Float64,
pl.String,
pl.String,
pl.Datetime("us"),
pl.Date,
pl.Binary,
pl.List(pl.Int32),
map_dtype,
pl.Struct([pl.Field("a", pl.Int32), pl.Field("b", pl.Int32)]),
pl.Decimal(precision=10, scale=1),
)
rows = df.to_dicts()
assert rows == [
{
"col_boolean": True,
"col_tinyint": 127,
"col_smallint": 32767,
"col_int": 2147483647,
"col_bigint": 9223372036854775807,
"col_float": 0.5,
"col_double": 0.25,
"col_string": "a string",
"col_varchar": "varchar",
"col_timestamp": datetime(2017, 1, 1, 0, 0, 0),
"col_date": datetime(2017, 1, 2).date(),
"col_binary": b"123",
"col_array": [1, 2],
"col_map": map_value,
"col_struct": {"a": 1, "b": 2},
"col_decimal": Decimal("0.1"),
}
]
assert df.to_dicts() == [{"k": 1, "s": "a"}, {"k": 2, "s": "b"}]

def test_cancel(self, arrow_cursor):
def cancel(c):
Expand Down Expand Up @@ -980,49 +804,6 @@ def test_callback(query_id: str):
assert callback_results[0] == arrow_cursor.query_id
assert arrow_cursor.query_id is not None

@pytest.mark.parametrize(
"arrow_cursor",
[
{
"cursor_kwargs": {
"connect_timeout": 10,
"request_timeout": 30,
}
}
],
indirect=["arrow_cursor"],
)
def test_timeout_parameters(self, arrow_cursor):
"""Test that timeout parameters are correctly passed to ArrowCursor and result set."""
# Verify timeout parameters are set on cursor
assert arrow_cursor._connect_timeout == 10
assert arrow_cursor._request_timeout == 30

# Execute a simple query to create a result set
arrow_cursor.execute("SELECT 1")

# Verify timeout parameters are passed to result set
assert arrow_cursor.result_set._connect_timeout == 10
assert arrow_cursor.result_set._request_timeout == 30

@pytest.mark.parametrize(
"arrow_cursor",
[{"cursor_kwargs": {"connect_timeout": 5.5, "request_timeout": 15.5}}],
indirect=["arrow_cursor"],
)
def test_timeout_parameters_float(self, arrow_cursor):
"""Test that timeout parameters accept float values."""
# Verify float timeout parameters are set on cursor
assert arrow_cursor._connect_timeout == 5.5
assert arrow_cursor._request_timeout == 15.5

# Execute a simple query to create a result set
arrow_cursor.execute("SELECT 1")

# Verify float timeout parameters are passed to result set
assert arrow_cursor.result_set._connect_timeout == 5.5
assert arrow_cursor.result_set._request_timeout == 15.5

@pytest.mark.parametrize(
"arrow_cursor",
[
Expand Down Expand Up @@ -1189,10 +970,13 @@ def test_managed_results_match_result_file(self):
indirect=["arrow_cursor"],
)
def test_duplicate_column_names(self, arrow_cursor):
"""Columns with the same name keep their own types, as do their timestamps and NULLs."""
arrow_cursor.execute(
"SELECT 1 AS x, 2 AS x, CAST('01:02:03' AS TIME) AS x, 'a' AS y, "
"json_parse('[1]') AS j, 'b' AS j, CAST('12:34:56' AS TIME) AS t, "
"CAST('01:02:03' AS TIME) AS t"
"CAST('01:02:03' AS TIME) AS t, "
"CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP(9)) AS x, "
"CAST(NULL AS TIMESTAMP) AS x"

@laughingman7743 laughingman7743 Oct 17, 2026

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.

Round 1 FINDINGS, repaired in 5a8b810: the merge of test_read_csv_timestamps_with_the_same_name had dropped its NULL timestamp in a duplicate-name column. This column restores it in the same query. Not yet run against Athena.

)
assert arrow_cursor.fetchall() == [
(
Expand All @@ -1204,9 +988,22 @@ def test_duplicate_column_names(self, arrow_cursor):
"b",
datetime(2017, 1, 1, 12, 34, 56).time(),
datetime(2017, 1, 1, 1, 2, 3).time(),
datetime(2020, 1, 2, 3, 4, 5, 123456),
None,
)
]
assert arrow_cursor.as_arrow().column_names == ["x", "x", "x", "y", "j", "j", "t", "t"]
assert arrow_cursor.as_arrow().column_names == [
"x",
"x",
"x",
"y",
"j",
"j",
"t",
"t",
"x",
"x",
]

def test_duplicate_column_names_with_newline(self, arrow_cursor):
"""Columns with the same name keep their own types, with a newline in the name."""
Expand Down
Loading
Loading