From 0c6a0b2de52fe1c3d8d03b79f01b2b3bf52c1ab9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 14 Oct 2026 21:56:44 +0900 Subject: [PATCH 1/4] Read CSV results through fsspec for Polars 2.0 Polars 2.0 reads a path in read_csv() with its native object store and rejects the fsspec storage options that PyAthena passed, so every CSV result read without chunksize failed. The result set now opens the file with fsspec.open() and the same options, and passes its contents to Polars, which keeps PyAthena's S3 filesystem and its read settings with Polars 1.x and 2.0. Polars 2.0 also names headerless columns from column_0 and requires new_columns to name every column, so the tests that pass those options through no longer depend on the Polars version. The dev lock tests Polars 2 on Python 3.14, which pull requests run, and keeps the newest Polars 1.x release on older versions. Closes #1168 Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 17 ++++-- pyproject.toml | 4 ++ tests/pyathena/polars/test_cursor.py | 7 ++- tests/pyathena/polars/test_result_set.py | 40 +++++++++++-- uv.lock | 73 ++++++++++++++++++++++-- 5 files changed, 125 insertions(+), 16 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 48964cc56..0438219d8 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -23,6 +23,8 @@ cast, ) +import fsspec + from pyathena import OperationalError from pyathena.arrow.reader import ParquetChunkReader from pyathena.converter import Converter @@ -346,10 +348,10 @@ def _read_kwargs( @property def _csv_storage_options(self) -> dict[str, Any]: - """Get storage options for Polars CSV reading via fsspec. + """Get storage options for reading a CSV result file via fsspec. - Polars read_csv uses fsspec for cloud storage access, which works - with PyAthena's registered S3FileSystem. + ``_read_csv()`` opens the file with ``fsspec.open()``, which uses + PyAthena's registered S3FileSystem, and passes its contents to Polars. Returns: Dictionary with fsspec-compatible options for S3 access. @@ -592,17 +594,22 @@ def _read_csv(self) -> pl.DataFrame: if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. kwargs.pop("new_columns", None) + # Polars 2.0 reads a path with its native object store, which does not take + # fsspec options, so the result file is read through fsspec here. + storage_options = kwargs.pop("storage_options") source: str | bytes if self.output_location: - source = location = self.output_location + location = self.output_location else: source = self._fetch_all_rows_as_csv() if not source: return pl.DataFrame() location = "the GetQueryResults rows" - del kwargs["storage_options"] try: + if self.output_location: + with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file: + source = file.read() try: df = pl.read_csv(source, **kwargs) except pl.exceptions.ComputeError: diff --git a/pyproject.toml b/pyproject.toml index 1e74fd0f4..2813446d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,10 @@ dev = [ "pandas>=3.0.0", "pyarrow>=22.0.0", "polars>=1.39.0", + # Pull requests test the newest Python with Polars 2, and the older Pythons keep + # the newest Polars 1.x release. + "polars>=2.0.0; python_version >= '3.14'", + "polars<2.0.0; python_version < '3.14'", "Jinja2>=3.1.0", "mypy>=0.900", "pytest>=3.5", diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index c922d8368..d0b512281 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -922,9 +922,12 @@ def test_new_columns_renaming_first_columns(self, polars_cursor, chunksize): assert polars_cursor.fetchall() == [("001", 2)] def test_new_columns_with_schema_overrides(self, polars_cursor): - """schema_overrides given with new_columns are keyed by the new names, as in Polars.""" + """schema_overrides given with new_columns are keyed by the new names, as in Polars. + + Polars 2.0 requires new_columns to name every column of the file. + """ polars_cursor.execute( - "SELECT '001' AS x, 2 AS y", new_columns=["z"], schema_overrides={"z": pl.String} + "SELECT '001' AS x, 2 AS y", new_columns=["z", "w"], schema_overrides={"z": pl.String} ) assert polars_cursor.fetchall() == [("001", 2)] diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 5007a8184..b5019d2ee 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -117,11 +117,13 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader): """Read arguments given to execute() replace the ones the result set chooses.""" path = tmp_path / "result.csv" path.write_text("1;x\n2;y\n") + # Polars names headerless columns from column_1 before 2.0 and from column_0 since. + first, second = pl.read_csv(b"1;x\n", separator=";", has_header=False).columns result_set = _chunked_result_set() result_set._kwargs = { "separator": ";", "has_header": False, - "schema_overrides": {"column_1": pl.Utf8}, + "schema_overrides": {first: pl.Utf8}, } with ( patch.object( @@ -152,7 +154,7 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader): ): result = getattr(result_set, reader)() df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) - assert df.to_dict(as_series=False) == {"column_1": ["1", "2"], "column_2": ["x", "y"]} + assert df.to_dict(as_series=False) == {first: ["1", "2"], second: ["x", "y"]} @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): @@ -160,7 +162,8 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): path = tmp_path / "result.txt" path.write_text("001\tx\n") result_set = _chunked_result_set() - result_set._kwargs = {"new_columns": ["z"], "schema_overrides": {"z": pl.Utf8}} + # Polars 2.0 requires new_columns to name every column of the file. + result_set._kwargs = {"new_columns": ["z", "w"], "schema_overrides": {"z": pl.Utf8}} with ( patch.object( AthenaPolarsResultSet, @@ -188,12 +191,11 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): ): result = getattr(result_set, reader)() df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) - assert df.to_dict(as_series=False) == {"z": ["001"], "b": ["x"]} + assert df.to_dict(as_series=False) == {"z": ["001"], "w": ["x"]} @pytest.mark.parametrize( ("reader", "function"), [ - ("_read_csv", "read_csv"), ("_iter_csv_chunks", "scan_csv"), ("_read_parquet", "read_parquet"), ("_iter_parquet_chunks", "scan_parquet"), @@ -237,6 +239,34 @@ def test_storage_options_replace_defaults(self, reader, function): list(result) assert read.call_args.kwargs["storage_options"] == {"anon": True} + def test_csv_storage_options_replace_defaults(self): + """storage_options given to execute() open a CSV result through fsspec, not Polars.""" + result_set = _chunked_result_set() + result_set._kwargs = {"storage_options": {"anon": True}} + with ( + patch.object( + AthenaPolarsResultSet, + "output_location", + new_callable=PropertyMock, + return_value="s3://bucket/result.csv", + ), + patch.object( + AthenaPolarsResultSet, "_csv_dtypes", new_callable=PropertyMock, return_value={} + ), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + side_effect=AssertionError("replaced storage options were computed"), + ), + patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), + patch("pyathena.polars.result_set.fsspec.open") as open_, + ): + open_.return_value.__enter__.return_value.read.return_value = b"a\n1\n" + df = result_set._read_csv() + open_.assert_called_once_with("s3://bucket/result.csv", "rb", anon=True) + assert df.to_dict(as_series=False) == {"a": [1]} + @pytest.mark.parametrize( ("kwargs", "expected"), [ diff --git a/uv.lock b/uv.lock index d172a5dbe..ef5c2bdd4 100644 --- a/uv.lock +++ b/uv.lock @@ -827,18 +827,57 @@ wheels = [ name = "polars" version = "1.44.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] dependencies = [ - { name = "polars-runtime-32" }, + { name = "polars-runtime-32", version = "1.44.2", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/a4/15/e8541eefc22fbc7ca89bcb5112298a153729f73cfbc0cf6a668e509f975c/polars-1.44.2.tar.gz", hash = "sha256:86c8e26b6c2de8c8d344bb910b74dfc47b118ac3fe0f19b44909467990a0b281", size = 766272, upload-time = "2026-09-09T07:42:08.859Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/51/6d/3014112c7f717d1253223faa13b6db3ac3a64ed00ab2a3bc1b942bc9cdd4/polars-1.44.2-py3-none-any.whl", hash = "sha256:1bb331f17a40d9d931101533dcd33637b66edc61eb377b07020dac16a0f0377b", size = 865768, upload-time = "2026-09-09T07:40:12.053Z" }, ] +[[package]] +name = "polars" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "polars-runtime-32", version = "2.0.0", source = { registry = "https://pypi.org/simple" } }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8e/e9/001f371ec6a1bb54893f599ceebd56e6144fed4091f09f09fec0021a9276/polars-2.0.0.tar.gz", hash = "sha256:62da109e27a19a9d36657ee25dc035c9d3f87e7bd610526fe467dc37ea7dc115", size = 778215, upload-time = "2026-10-06T11:51:29.679Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/09/cc33bbd5463749c116b62c204d88bed6c02a6cb901eac7adab0d38651b07/polars-2.0.0-py3-none-any.whl", hash = "sha256:35d62f3541b7a6d4c360a2e2f07fccc0c2bcbd33b0ea51c83a25417a47a3f3ad", size = 876611, upload-time = "2026-10-06T11:44:04.327Z" }, +] + [[package]] name = "polars-runtime-32" version = "1.44.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] sdist = { url = "https://files.pythonhosted.org/packages/d4/a1/a7eace6587b56f22cf2a21ab4d5e695db372dc23fd96accb68b1ec12660b/polars_runtime_32-1.44.2.tar.gz", hash = "sha256:b84842f7d621aaca7a52e165e19a24f89db45f8aa13744941430218419a14a67", size = 3172205, upload-time = "2026-09-09T07:42:10.851Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/99/5b/a5215f82c3dd443dc5d6911b0d3e937f97056e0ef7753f7e123422481a18/polars_runtime_32-1.44.2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:1fd536720668ba203a16a20b08cd6b23057e407a0279cf36b2f35f879d6e3208", size = 48334014, upload-time = "2026-09-09T07:40:16.43Z" }, @@ -851,6 +890,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/a0/d0dd0d2ec95fa328dd47055905fae53ba3cd79f11c8973326ebe75a49e4c/polars_runtime_32-1.44.2-cp310-abi3-win_arm64.whl", hash = "sha256:d51040d3ab40157f6db3c62be59cab5b80fb3c8d158924769c4982a1c8eef730", size = 44303345, upload-time = "2026-09-09T07:40:47.081Z" }, ] +[[package]] +name = "polars-runtime-32" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +sdist = { url = "https://files.pythonhosted.org/packages/34/ad/dbb6f6d7070867951532bcfe5e6a648d8777b416b18cddabc07030404e8c/polars_runtime_32-2.0.0.tar.gz", hash = "sha256:b5f9afcc742b4a67eabd2c680ff0f12eb02ede9b4bf807bffabd6dbb9a58d5c7", size = 3591339, upload-time = "2026-10-06T11:51:31.076Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/88/d35dec6c8928dfbaa1cccf9b626a1067da906e792c92d9f994ca825ab2b5/polars_runtime_32-2.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ffb7ac6cf4e8c4a652df1951e3c3840c7c23a033603d5a9efd422fa8dd699d82", size = 52494314, upload-time = "2026-10-06T11:44:07.768Z" }, + { url = "https://files.pythonhosted.org/packages/5f/fd/2237bf53ffaff47cdf1edc6c10587a7a6444d4951150eeb08d84f3493ff8/polars_runtime_32-2.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:7012d8a0201bd95638545ce8f256c0efe2c5cab0f806eb043021dddde5a9498b", size = 47930083, upload-time = "2026-10-06T11:44:11.592Z" }, + { url = "https://files.pythonhosted.org/packages/0d/0d/85e3ed90417996fc09770be91b39979074fe2978fc15b431bf8a9459760d/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8b85bb42e6009acc9629afcc70a83473fd468694d6a30ffb0ab376c8dd1a0a17", size = 50417889, upload-time = "2026-10-06T11:50:20.774Z" }, + { url = "https://files.pythonhosted.org/packages/83/88/e9fecfd49159da92f54ff2445883577a0f1bc195da53ecc9535c458d55dd/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d6ac584ea2b38913784db943879412380d92e28ab9cb88e20a77ba71ba3f911", size = 54475036, upload-time = "2026-10-06T11:50:24.411Z" }, + { url = "https://files.pythonhosted.org/packages/48/ad/b2abf732697b21467aaaeaac0f3bf7eee0d89c59ce8125f1ed41b28a2d97/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:a6bf5e260e0a6f00d0f9181438fe9e45776df8c66cee9cba16e3675cc3888488", size = 50579474, upload-time = "2026-10-06T11:50:28.377Z" }, + { url = "https://files.pythonhosted.org/packages/7f/05/304deee59a95865e1b5e9ec7b066069b49093b81b768f473d9d3b165c686/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:55c26eef325b6840584d91aac232e9cf3ac19e1b904594b9b54131be1edeab4d", size = 54413293, upload-time = "2026-10-06T11:50:31.828Z" }, + { url = "https://files.pythonhosted.org/packages/61/59/8c9fd7199f7c4eb1b64e640306a946a2e4a46337b3bbb33b840972c7d84b/polars_runtime_32-2.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:7da1caf3c7b4f397fb213c984013a0c755557619a2d511899a1ff74392484078", size = 54229989, upload-time = "2026-10-06T11:50:35.206Z" }, + { url = "https://files.pythonhosted.org/packages/e2/93/43608026f38aa6ed4d22da8597706a61682ee403caef0021ce8e6dc73227/polars_runtime_32-2.0.0-cp310-abi3-win_arm64.whl", hash = "sha256:c30ba698c8904048df4a9bc3d6c5033cc2d0a7cbb0e13f4fd2de5a1947b61994", size = 48730655, upload-time = "2026-10-06T11:50:38.756Z" }, +] + [[package]] name = "psutil" version = "7.2.2" @@ -944,7 +1004,8 @@ pandas = [ { name = "pandas" }, ] polars = [ - { name = "polars" }, + { name = "polars", version = "1.44.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "polars", version = "2.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, ] sqlalchemy = [ { name = "sqlalchemy" }, @@ -958,7 +1019,8 @@ dev = [ { name = "mypy" }, { name = "myst-parser" }, { name = "pandas" }, - { name = "polars" }, + { name = "polars", version = "1.44.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "polars", version = "2.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "pyarrow" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -997,6 +1059,8 @@ dev = [ { name = "myst-parser" }, { name = "pandas", specifier = ">=3.0.0" }, { name = "polars", specifier = ">=1.39.0" }, + { name = "polars", marker = "python_full_version < '3.14'", specifier = "<2.0.0" }, + { name = "polars", marker = "python_full_version >= '3.14'", specifier = ">=2.0.0" }, { name = "pyarrow", specifier = ">=22.0.0" }, { name = "pytest", specifier = ">=3.5" }, { name = "pytest-asyncio" }, @@ -1017,7 +1081,8 @@ version = "0.1.0" source = { virtual = "benchmarks" } dependencies = [ { name = "awswrangler" }, - { name = "polars" }, + { name = "polars", version = "1.44.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "polars", version = "2.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "psutil" }, { name = "pyathena", extra = ["arrow", "pandas", "polars", "sqlalchemy"] }, ] From 6ec3e0087590eec57e94b9a0c188038e7e45bb68 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 14 Oct 2026 21:58:39 +0900 Subject: [PATCH 2/4] Document where Polars storage_options go and narrow the new_columns notes Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 1 + tests/pyathena/polars/test_cursor.py | 2 +- tests/pyathena/polars/test_result_set.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/polars.md b/docs/polars.md index de8e3b69d..28295b61a 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -397,6 +397,7 @@ df = cursor.execute("SELECT * FROM many_rows").as_polars() They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option, which use `"bytes"` and 5 MiB blocks by default, and the Parquet files of UNLOAD results read with the chunksize option when pyarrow is installed and `execute()` is given no Polars read options. Of the latter, a file of up to 32 MiB is read at once, and a larger one is read with the given `cache_type`, or without it with the `"none"` cache, which requests only the ranges that pyarrow reads. With `storage_options` given to `execute()`, the given options replace PyAthena's S3 settings, and `block_size` and `cache_type` do not apply. +For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to Polars, which uses its native S3 access. (async-polars-cursor)= diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index d0b512281..521e1e70e 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -924,7 +924,7 @@ def test_new_columns_renaming_first_columns(self, polars_cursor, chunksize): def test_new_columns_with_schema_overrides(self, polars_cursor): """schema_overrides given with new_columns are keyed by the new names, as in Polars. - Polars 2.0 requires new_columns to name every column of the file. + Polars 2.0 rejects new_columns that do not name every column of the file. """ polars_cursor.execute( "SELECT '001' AS x, 2 AS y", new_columns=["z", "w"], schema_overrides={"z": pl.String} diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index b5019d2ee..4e80ce321 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -162,7 +162,7 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): path = tmp_path / "result.txt" path.write_text("001\tx\n") result_set = _chunked_result_set() - # Polars 2.0 requires new_columns to name every column of the file. + # Polars 2.0 rejects new_columns that do not name every column of the file. result_set._kwargs = {"new_columns": ["z", "w"], "schema_overrides": {"z": pl.Utf8}} with ( patch.object( From a473e9cd816c3b58d2d782080b49d90b2af3458c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 14 Oct 2026 22:13:10 +0900 Subject: [PATCH 3/4] Describe Polars map columns of UNLOAD results as map Polars 2.0 reads a Parquet map as a Map column, which get_athena_type() did not know, so UNLOAD results described map columns as string. Nested types are now matched by their outer type, which also describes a struct containing a list as row instead of array. The docs now say that Polars reads with fsspec when use_pyarrow=True is given. Co-Authored-By: Claude Opus 5.5 --- docs/polars.md | 2 +- pyathena/polars/util.py | 14 +++++++++---- tests/pyathena/polars/test_cursor.py | 23 +++++++++++++++++++++ tests/pyathena/polars/test_util.py | 31 ++++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 5 deletions(-) create mode 100644 tests/pyathena/polars/test_util.py diff --git a/docs/polars.md b/docs/polars.md index 28295b61a..b72f86dc2 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -397,7 +397,7 @@ df = cursor.execute("SELECT * FROM many_rows").as_polars() They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option, which use `"bytes"` and 5 MiB blocks by default, and the Parquet files of UNLOAD results read with the chunksize option when pyarrow is installed and `execute()` is given no Polars read options. Of the latter, a file of up to 32 MiB is read at once, and a larger one is read with the given `cache_type`, or without it with the `"none"` cache, which requests only the ranges that pyarrow reads. With `storage_options` given to `execute()`, the given options replace PyAthena's S3 settings, and `block_size` and `cache_type` do not apply. -For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to Polars, which uses its native S3 access. +For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to the Polars read function, which uses its native S3 access by default and fsspec with `use_pyarrow=True`. (async-polars-cursor)= diff --git a/pyathena/polars/util.py b/pyathena/polars/util.py index dc7f8680e..98976a087 100644 --- a/pyathena/polars/util.py +++ b/pyathena/polars/util.py @@ -104,12 +104,18 @@ def get_athena_type(dtype: Any) -> tuple[str, int, int]: if base_dtype == polars_type or dtype == polars_type: return athena_info - # Handle parameterized types that didn't match above - dtype_str = str(dtype).lower() - if "list" in dtype_str: + # Nested types are matched by their outer type, not by the types they contain. + if base_dtype == pl.List: return ("array", 0, 0) - if "struct" in dtype_str: + if base_dtype == pl.Struct: return ("row", 0, 0) + # Polars 2.0 reads a Parquet map as Map; Polars 1.x reads it as a list of structs. + map_dtype = getattr(pl, "Map", None) + if map_dtype is not None and base_dtype == map_dtype: + return ("map", 0, 0) + + # Handle parameterized types that didn't match above + dtype_str = str(dtype).lower() if "decimal" in dtype_str: # Extract precision and scale from Decimal type if available if hasattr(dtype, "precision") and hasattr(dtype, "scale"): diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 521e1e70e..4a050c579 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -768,6 +768,29 @@ def test_description_with_chunksize_unload(self, polars_cursor, query, expected) assert description == ([("a", "integer"), ("b", "varchar")] if expected else []) assert polars_cursor.fetchall() == expected + @pytest.mark.parametrize( + "polars_cursor", + [ + {"cursor_kwargs": {"unload": True}}, + {"cursor_kwargs": {"unload": True, "chunksize": 5}}, + ], + indirect=["polars_cursor"], + ) + def test_nested_type_description_unload(self, polars_cursor): + """UNLOAD results describe nested Parquet columns by their outer type.""" + polars_cursor.execute( + """ + SELECT + ARRAY[1, 2] AS a + ,CAST(ROW(ARRAY[1]) AS ROW(x ARRAY(INTEGER))) AS r + ,MAP(ARRAY['k'], ARRAY[1]) AS m + """ + ) + # Polars 1.x reads a Parquet map as a list of key-value structs. + map_type = "map" if hasattr(pl, "Map") else "array" + description = [d[:2] for d in polars_cursor.description] + assert description == [("a", "array"), ("r", "row"), ("m", map_type)] + @pytest.mark.parametrize( "polars_cursor", [ diff --git a/tests/pyathena/polars/test_util.py b/tests/pyathena/polars/test_util.py new file mode 100644 index 000000000..800387ccc --- /dev/null +++ b/tests/pyathena/polars/test_util.py @@ -0,0 +1,31 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +import polars as pl +import pytest + +from pyathena.polars.util import get_athena_type + + +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + (pl.List(pl.Int32), ("array", 0, 0)), + (pl.List(pl.Struct({"x": pl.Int32})), ("array", 0, 0)), + (pl.Struct({"x": pl.Int32}), ("row", 0, 0)), + (pl.Struct({"x": pl.List(pl.Int32)}), ("row", 0, 0)), + ], +) +def test_get_athena_type_nested(dtype, expected): + assert get_athena_type(dtype) == expected + + +@pytest.mark.skipif(not hasattr(pl, "Map"), reason="Polars 1.x has no Map type") +def test_get_athena_type_map(): + assert get_athena_type(pl.Map(pl.String, pl.Int32)) == ("map", 0, 0) + assert get_athena_type(pl.Map(pl.String, pl.List(pl.Int32))) == ("map", 0, 0) + assert get_athena_type(pl.List(pl.Map(pl.String, pl.Int32))) == ("array", 0, 0) From 22de45412887431ca90c180b46b243abc5ef39c7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 14 Oct 2026 22:31:37 +0900 Subject: [PATCH 4/4] Expect Polars 2.0's map conversion in ArrowCursor.as_polars() Polars 2.0 converts an Arrow map to a Map column whose values are dicts; Polars 1.x converts it to a list of key-value structs. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/arrow/test_cursor.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 54352fa9c..162a146a3 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -626,6 +626,13 @@ def test_complex_unload_as_polars(self, arrow_cursor): ).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, @@ -641,7 +648,7 @@ def test_complex_unload_as_polars(self, arrow_cursor): pl.Date, pl.Binary, pl.List(pl.Int32), - pl.List(pl.Struct([pl.Field("key", pl.Int32), pl.Field("value", pl.Int32)])), + map_dtype, pl.Struct([pl.Field("a", pl.Int32), pl.Field("b", pl.Int32)]), pl.Decimal(precision=10, scale=1), ) @@ -661,7 +668,7 @@ def test_complex_unload_as_polars(self, arrow_cursor): "col_date": datetime(2017, 1, 2).date(), "col_binary": b"123", "col_array": [1, 2], - "col_map": [{"key": 1, "value": 2}, {"key": 3, "value": 4}], + "col_map": map_value, "col_struct": {"a": 1, "b": 2}, "col_decimal": Decimal("0.1"), }