From e2006b36161dc5e954bb4870a2914eecc89e31ff Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:46:14 +0900 Subject: [PATCH 1/3] Convert the timestamp text of chunked Polars CSV reads in the scan Athena writes up to 12 fractional digits for timestamp columns, and Polars does not parse more than nine into a Datetime. A whole PolarsCursor CSV read catches the resulting ComputeError and reads the data again with those columns as text, truncating them with _to_datetimes(). The chunked reader passed the same Datetime dtypes to scan_csv() without that fallback, so the read raised OperationalError. A retry cannot work for chunks, because a value that fails to parse can be in a later chunk, after the earlier chunks have been returned. The chunked reader now reads the timestamp columns as text and applies _to_datetimes() to the scan_csv() LazyFrame, so the conversion runs in the query plan. _to_datetimes() takes a LazyFrame as well as a DataFrame, selecting the columns with cs.by_name(require_all=False) so that a LazyFrame need not resolve its schema, and takes strict, which the readers set from ignore_errors so that text that is not a timestamp becomes null as Polars reads it. _get_timestamp_dtypes() becomes _get_csv_schema_overrides(), which also returns the schema_overrides with those columns as String. With schema given to execute(), Polars ignores schema_overrides, so the columns are not read as text, as with schema_overrides and with_column_names. The whole read keeps parsing natively, because reading the columns as text raised its peak RSS by about 35-40% in local measurements. Fixes #1191 Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/converter.py | 26 +++--- pyathena/polars/result_set.py | 59 ++++++++---- tests/pyathena/polars/test_converter.py | 20 ++++- tests/pyathena/polars/test_cursor.py | 16 ++++ tests/pyathena/polars/test_result_set.py | 110 ++++++++++++++++++----- 5 files changed, 178 insertions(+), 53 deletions(-) diff --git a/pyathena/polars/converter.py b/pyathena/polars/converter.py index a40e9c7d5..37659f34f 100644 --- a/pyathena/polars/converter.py +++ b/pyathena/polars/converter.py @@ -12,7 +12,7 @@ import logging from collections.abc import Callable from copy import deepcopy -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar from pyathena.converter import ( _TIMESTAMP_TEXT_LENGTHS, @@ -32,6 +32,8 @@ _logger = logging.getLogger(__name__) +_FrameT = TypeVar("_FrameT", "pl.DataFrame", "pl.LazyFrame") + _DEFAULT_POLARS_CONVERTERS: dict[str, Callable[[str | None], Any | None]] = { "date": _to_date, @@ -43,36 +45,38 @@ } -def _to_datetimes(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame: +def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any], strict: bool = True) -> _FrameT: """Convert timestamp text columns to Datetime dtypes, truncating finer fractions. Athena writes up to 12 fractional digits, which Polars does not parse into a Datetime whose time unit holds fewer. Args: - df: The DataFrame with the timestamp text columns, with NULL as null or as - an empty string. + frame: The DataFrame or LazyFrame with the timestamp text columns, with NULL + as null or as an empty string. A column that it does not have, such as + one that the read options leave out, is skipped. dtypes: The Datetime dtypes keyed by column name. + strict: Whether text that is not a timestamp raises an error, rather than + becoming null. Returns: - The DataFrame with the columns converted. + The frame with the columns converted. """ import polars as pl + import polars.selectors as cs exprs = [] for name, dtype in dtypes.items(): - if name not in df.columns: - # Not read, such as a column that the read options leave out. - continue time_unit = (dtype() if isinstance(dtype, type) else dtype).time_unit - text = pl.col(name) + # A selector, not pl.col(), so that a LazyFrame need not resolve its schema. + text = cs.by_name(name, require_all=False).as_expr() exprs.append( pl.when(text != "") .then(text.str.slice(0, _TIMESTAMP_TEXT_LENGTHS[time_unit])) - .str.to_datetime("%Y-%m-%d %H:%M:%S%.f", time_unit=time_unit) + .str.to_datetime("%Y-%m-%d %H:%M:%S%.f", time_unit=time_unit, strict=strict) .cast(dtype) ) - return df.with_columns(exprs) if exprs else df + return frame.with_columns(exprs) if exprs else frame class DefaultPolarsTypeConverter(Converter): diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index b79802547..2b45692a2 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -547,29 +547,41 @@ def _get_converters( for name, d in zip(column_names, description, strict=True) } - def _get_timestamp_dtypes(self, has_header: bool) -> dict[str, Any]: - """Get the Datetime dtypes of the timestamp columns, which are read as text. + def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], dict[str, Any]]: + """Get the Polars data types for reading CSV data, with the timestamp columns as text. + + Athena writes up to 12 fractional digits, which Polars does not parse into a + Datetime whose time unit holds fewer. A chunked read always reads the timestamp + columns as text and converts them with ``_to_datetimes()``, and a whole read + does so when parsing them fails. Args: has_header: Whether the CSV data has a header. Without one, or with ``schema_overrides`` or ``with_column_names`` given to ``execute()``, - Polars does not name the columns by the header, and they are not read - as text. + Polars does not name the columns by the header, and with ``schema``, + Polars ignores ``schema_overrides``, so the columns are not read as text. Returns: - The Datetime dtypes keyed by the header of a CSV file. + Tuple of (schema_overrides, timestamp_dtypes): the data types to read the + columns with, and the Datetime dtypes of the timestamp columns read as text, + both keyed by the header of a CSV file. """ import polars as pl - if not has_header or self._kwargs.keys() & {"schema_overrides", "with_column_names"}: - return {} dtypes = self._csv_dtypes - return { + if not has_header or self._kwargs.keys() & { + "schema", + "schema_overrides", + "with_column_names", + }: + return dtypes, {} + timestamp_dtypes = { name: dtype for name, d in zip(self._get_column_names(), self.description or [], strict=True) if d[1] == "timestamp" and ((dtype := dtypes.get(name)) is pl.Datetime or isinstance(dtype, pl.Datetime)) } + return {**dtypes, **dict.fromkeys(timestamp_dtypes, pl.String)}, timestamp_dtypes def _get_column_names(self) -> list[str]: """Get the names of the result columns in a DataFrame. @@ -731,17 +743,16 @@ def _read_csv(self) -> pl.DataFrame: try: df = pl.read_csv(source, **kwargs) except pl.exceptions.ComputeError: - timestamp_dtypes = self._get_timestamp_dtypes(has_header) + schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) if not timestamp_dtypes: raise - # Athena writes up to 12 fractional digits, which Polars does not parse - # into a Datetime whose time unit holds fewer, so the data is read again - # with the timestamp columns as text. - kwargs["schema_overrides"] = { - **self._csv_dtypes, - **dict.fromkeys(timestamp_dtypes, pl.String), - } - df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes) + # The data is read again with the timestamp columns as text. + kwargs["schema_overrides"] = schema_overrides + df = _to_datetimes( + pl.read_csv(source, **kwargs), + timestamp_dtypes, + strict=not kwargs.get("ignore_errors", False), + ) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df @@ -908,14 +919,16 @@ def _iter_csv_chunks(self) -> Iterator[pl.DataFrame]: raise ProgrammingError("output_location is not available.") separator, has_header, new_columns = self._get_csv_params() + schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) return self._scan_csv_chunks( self.output_location, self.connection, { "separator": separator, "has_header": has_header, - "schema_overrides": self._csv_dtypes, + "schema_overrides": schema_overrides, }, + timestamp_dtypes, new_columns, ) @@ -924,6 +937,7 @@ def _scan_csv_chunks( location: str, connection: Connection[Any], defaults: dict[str, Any], + timestamp_dtypes: dict[str, Any], new_columns: list[str] | None, ) -> Iterator[pl.DataFrame]: """Scan a CSV result file in chunks with Polars. @@ -934,6 +948,8 @@ def _scan_csv_chunks( frozen when the first chunk is read. defaults: The arguments of ``polars.scan_csv()`` that the result set chooses; see ``_read_kwargs()``. + timestamp_dtypes: The Datetime dtypes of the timestamp columns read as + text; see ``_get_csv_schema_overrides()``. new_columns: The names of the first columns, set after reading; see ``_get_csv_params()``. @@ -954,7 +970,12 @@ def _scan_csv_chunks( if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. read_kwargs.pop("new_columns", None) - lazy_df = pl.scan_csv(location, **read_kwargs) + lazy_df = _to_datetimes( + pl.scan_csv(location, **read_kwargs), + timestamp_dtypes, + # As Polars reads a value it cannot parse as null with ignore_errors. + strict=not read_kwargs.get("ignore_errors", False), + ) for batch in lazy_df.collect_batches(chunk_size=self._chunksize): if new_columns: batch.columns = [*new_columns, *batch.columns[len(new_columns) :]] diff --git a/tests/pyathena/polars/test_converter.py b/tests/pyathena/polars/test_converter.py index f025e001d..d103d611f 100644 --- a/tests/pyathena/polars/test_converter.py +++ b/tests/pyathena/polars/test_converter.py @@ -20,6 +20,7 @@ def test_convert_delegates_to_default(self): assert converter.convert("varchar", "hello") == "hello" +@pytest.mark.parametrize("lazy", [False, True], ids=["dataframe", "lazyframe"]) @pytest.mark.parametrize( ("dtype", "microseconds"), [ @@ -28,11 +29,12 @@ def test_convert_delegates_to_default(self): (pl.Datetime("us"), [0, 123000, 123456, 123456, 123456]), ], ) -def test_to_datetimes(dtype, microseconds): +def test_to_datetimes(dtype, microseconds, lazy): """Timestamp text with up to 12 fractional digits is truncated to the time unit. NULL can be null or an empty string, depending on the read options. A column that - the DataFrame does not have, such as one not selected, is skipped. + the frame does not have, such as one not selected, is skipped. A LazyFrame is + converted in its query plan. """ df = pl.DataFrame( { @@ -48,7 +50,10 @@ def test_to_datetimes(dtype, microseconds): ] } ) - result = _to_datetimes(df, {"t": dtype, "missing": dtype}) + result = _to_datetimes(df.lazy() if lazy else df, {"t": dtype, "missing": dtype}) + if lazy: + assert isinstance(result, pl.LazyFrame) + result = result.collect() assert result.schema["t"] == dtype assert result["t"].to_list() == [ *(datetime(2020, 1, 2, 3, 4, 5, us) for us in microseconds), @@ -56,3 +61,12 @@ def test_to_datetimes(dtype, microseconds): None, None, ] + + +def test_to_datetimes_not_strict(): + """Text that is not a timestamp becomes null when the conversion is not strict.""" + df = pl.DataFrame({"t": ["2020-01-02 03:04:05.123456789012", "invalid"]}) + with pytest.raises(pl.exceptions.InvalidOperationError): + _to_datetimes(df, {"t": pl.Datetime}) + result = _to_datetimes(df, {"t": pl.Datetime}, strict=False) + assert result["t"].to_list() == [datetime(2020, 1, 2, 3, 4, 5, 123456), None] diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 2d0cb4631..f0a5443d3 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -178,6 +178,22 @@ def test_temporal_types(self, polars_cursor, unload): assert row == temporal.fetched_row(unload) assert temporal.zones(row) == temporal.zones(temporal.fetched_row(unload)) + @pytest.mark.parametrize("polars_cursor", [{"cursor_kwargs": {"chunksize": 1}}], indirect=True) + def test_temporal_types_chunked(self, polars_cursor): + """Chunks of a CSV result have the temporal types and values of a whole read (#1191). + + Timestamps with more fractional digits than the time unit are truncated. + """ + types = {"timestamp": pl.Datetime("us"), "date": pl.Date()} + columns = temporal.columns() + chunks = list(polars_cursor.execute(temporal.query()).iter_chunks()) + assert len(chunks) == 1 + df = chunks[0] + assert dict(df.schema) == {c.name: types.get(c.athena_type, pl.String()) for c in columns} + assert [df[c.name][0] for c in columns] == [ + c.value if c.athena_type in types else c.text for c in columns + ] + def test_complex(self, polars_cursor): polars_cursor.execute( """ diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index e6ebdc915..3d10bd0d1 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -7,6 +7,7 @@ import inspect import threading +from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from datetime import UTC, date, datetime, time, timedelta from decimal import Decimal @@ -25,6 +26,12 @@ from pyathena.util import RetryConfig _ROWS_BEFORE_FAILURE = 300_000 +# A timestamp column of a CSV result file with 12 fractional digits and NULL, as text and +# as read into a Datetime column with the default time unit. +_TIMESTAMP_TEXTS = pl.Series(["2020-01-02 03:04:05.123456789012", None], dtype=pl.String) +_TRUNCATED_TIMESTAMPS = pl.Series( + [datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us") +) def _chunked_result_set() -> AthenaPolarsResultSet: @@ -42,6 +49,18 @@ def _chunked_result_set() -> AthenaPolarsResultSet: return result_set +def _read_all(result: pl.DataFrame | Iterator[pl.DataFrame]) -> pl.DataFrame: + """Read the result of a whole read or a chunk reader into one DataFrame. + + Args: + result: A DataFrame, or an iterator of DataFrame chunks. + + Returns: + The DataFrame, or the chunks concatenated. + """ + return result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) + + def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose chunks are ``chunks``, without AWS. @@ -411,8 +430,7 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader): patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), ): - result = getattr(result_set, reader)() - df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) + df = _read_all(getattr(result_set, reader)()) assert df.to_dict(as_series=False) == {first: ["1", "2"], second: ["x", "y"]} @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) @@ -443,8 +461,7 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader): patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), ): - result = getattr(result_set, reader)() - df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) + df = _read_all(getattr(result_set, reader)()) assert df.to_dict(as_series=False) == {"z": ["001"], "w": ["x"]} @pytest.mark.parametrize( @@ -521,22 +538,43 @@ def test_csv_storage_options_replace_defaults(self): assert df.to_dict(as_series=False) == {"a": [1]} @pytest.mark.parametrize( - ("kwargs", "expected"), + ("reader", "kwargs", "expected"), [ - ({}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ({"columns": ["v"]}, {}), - ({"new_columns": ["t2", "v2"]}, {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ({"with_column_names": lambda names: [n.upper() for n in names]}, OperationalError), + ("_read_csv", {}, {"t": _TRUNCATED_TIMESTAMPS}), + ("_iter_csv_chunks", {}, {"t": _TRUNCATED_TIMESTAMPS}), + ("_read_csv", {"columns": ["v"]}, {}), + ("_read_csv", {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), + ("_iter_csv_chunks", {"new_columns": ["t2", "v2"]}, {"t2": _TRUNCATED_TIMESTAMPS}), + ("_read_csv", {"schema": {"t": pl.String, "v": pl.String}}, {"t": _TIMESTAMP_TEXTS}), + ( + "_iter_csv_chunks", + {"schema": {"t": pl.String, "v": pl.String}}, + {"t": _TIMESTAMP_TEXTS}, + ), + ( + "_read_csv", + {"with_column_names": lambda names: [n.upper() for n in names]}, + OperationalError, + ), + ( + "_iter_csv_chunks", + {"with_column_names": lambda names: [n.upper() for n in names]}, + OperationalError, + ), ], ) - def test_read_csv_truncates_timestamps(self, kwargs, expected): - """Timestamps that fail to parse are read again as text and truncated. + def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected): + """Timestamps are read as text and truncated to the time unit. + + A whole read does so when parsing them fails; a chunked read always does. With ``with_column_names`` given to execute(), Polars renames the columns as it - reads them, so they are not, and the read fails as before. No AWS calls; the - GetQueryResults rows are mocked. + reads them, so they are not read as text, and the read fails. With ``schema``, + the columns have the types it gives. ``scan_csv()`` does not take ``columns``. """ - result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ + path = tmp_path / "result.csv" + path.write_bytes(b'"t","v"\n"2020-01-02 03:04:05.123456789012","x"\n,"y"\n') + result_set = _chunked_result_set() result_set._query_execution = None result_set._converter = DefaultPolarsTypeConverter() result_set._kwargs = kwargs @@ -544,27 +582,59 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): {"Name": n, "Type": t, "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"} for n, t in (("t", "timestamp"), ("v", "varchar")) ) - data = b'"t","v"\n"2020-01-02 03:04:05.123456789012","x"\n,"y"\n' with ( - patch.object(AthenaPolarsResultSet, "_fetch_all_rows_as_csv", return_value=data), + patch.object( + AthenaPolarsResultSet, + "output_location", + new_callable=PropertyMock, + return_value=str(path), + ), patch.object( AthenaPolarsResultSet, "_csv_storage_options", new_callable=PropertyMock, return_value={}, ), + patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), + patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), ): if expected is OperationalError: with pytest.raises(OperationalError): - result_set._read_csv() + _read_all(getattr(result_set, reader)()) return - df = result_set._read_csv() + df = _read_all(getattr(result_set, reader)()) for name, values in expected.items(): - assert df.schema[name] == pl.Datetime("us") - assert df[name].to_list() == values + assert df[name].equals(values, check_dtypes=True) if not expected: assert df.columns == ["v"] + def test_iter_csv_chunks_ignore_errors(self, tmp_path): + """With ignore_errors, timestamp text that is not a timestamp becomes null in chunks.""" + path = tmp_path / "result.csv" + path.write_bytes(b'"t"\n"2020-01-02 03:04:05.123456789012"\n"invalid"\n') + result_set = _chunked_result_set() + result_set._query_execution = None + result_set._converter = DefaultPolarsTypeConverter() + result_set._kwargs = {"ignore_errors": True} + result_set._metadata = ( + {"Name": "t", "Type": "timestamp", "Precision": 3, "Scale": 0, "Nullable": "UNKNOWN"}, + ) + with ( + patch.object( + AthenaPolarsResultSet, + "output_location", + new_callable=PropertyMock, + return_value=str(path), + ), + patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), + patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), + ): + df = _read_all(result_set._iter_csv_chunks()) + assert df["t"].equals( + pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), + check_dtypes=True, + ) + class TestPolarsDataFrameIterator: @pytest.mark.parametrize( From 7d5584efcf186b7513fe1a68b123f731cb9ec434 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:51:00 +0900 Subject: [PATCH 2/3] Read the timestamps of whole Polars CSV reads as text with ignore_errors With ignore_errors given to execute(), Polars reads a timestamp with more than nine fractional digits as null instead of raising ComputeError, so the whole read never fell back to reading the timestamp columns as text and returned null for such values, while a chunked read now truncates them. With ignore_errors, the whole read now reads those columns as text at once and converts them without strict, so both reads return the truncated values, and null for text that is not a timestamp. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 31 +++++++++++++----------- tests/pyathena/polars/test_result_set.py | 17 ++++++++++--- 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 2b45692a2..8f2273754 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -552,8 +552,9 @@ def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], d Athena writes up to 12 fractional digits, which Polars does not parse into a Datetime whose time unit holds fewer. A chunked read always reads the timestamp - columns as text and converts them with ``_to_datetimes()``, and a whole read - does so when parsing them fails. + columns as text and converts them with ``_to_datetimes()``. A whole read does + so when parsing them fails, and with ``ignore_errors`` given to ``execute()``, + with which Polars reads the timestamps that it does not parse as null. Args: has_header: Whether the CSV data has a header. Without one, or with @@ -740,19 +741,21 @@ def _read_csv(self) -> pl.DataFrame: 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: - schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) - if not timestamp_dtypes: - raise - # The data is read again with the timestamp columns as text. + schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) + strict = not kwargs.get("ignore_errors", False) + df: pl.DataFrame | None = None + if strict or not timestamp_dtypes: + try: + df = pl.read_csv(source, **kwargs) + except pl.exceptions.ComputeError: + if not timestamp_dtypes: + raise + if df is None: + # The timestamp columns are read as text, again after parsing them + # failed, or at once with ignore_errors, with which Polars reads the + # timestamps that it does not parse as null instead of failing. kwargs["schema_overrides"] = schema_overrides - df = _to_datetimes( - pl.read_csv(source, **kwargs), - timestamp_dtypes, - strict=not kwargs.get("ignore_errors", False), - ) + df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes, strict=strict) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 3d10bd0d1..9d9b7c1c3 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -608,8 +608,13 @@ def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected) if not expected: assert df.columns == ["v"] - def test_iter_csv_chunks_ignore_errors(self, tmp_path): - """With ignore_errors, timestamp text that is not a timestamp becomes null in chunks.""" + @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) + def test_read_csv_ignore_errors(self, tmp_path, reader): + """With ignore_errors, timestamp text that is not a timestamp becomes null. + + Timestamps with more fractional digits than the time unit are still truncated, + which Polars itself reads as null with ignore_errors. + """ path = tmp_path / "result.csv" path.write_bytes(b'"t"\n"2020-01-02 03:04:05.123456789012"\n"invalid"\n') result_set = _chunked_result_set() @@ -626,10 +631,16 @@ def test_iter_csv_chunks_ignore_errors(self, tmp_path): new_callable=PropertyMock, return_value=str(path), ), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + return_value={}, + ), patch.object(AthenaPolarsResultSet, "_get_parquet_storage_options", return_value={}), patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), ): - df = _read_all(result_set._iter_csv_chunks()) + df = _read_all(getattr(result_set, reader)()) assert df["t"].equals( pl.Series([datetime(2020, 1, 2, 3, 4, 5, 123456), None], dtype=pl.Datetime("us")), check_dtypes=True, From 3276a3f472bd3996c154f42f2142c0c6d2eb74a2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 12:52:45 +0900 Subject: [PATCH 3/3] Say that Polars parses up to nine fractional digits Polars parses timestamps with up to nine fractional digits into a Datetime of any time unit, truncating them, and fails only on more, so the docstrings no longer say that it fails on more digits than the time unit holds. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/converter.py | 4 ++-- pyathena/polars/result_set.py | 4 ++-- tests/pyathena/polars/test_cursor.py | 2 +- tests/pyathena/polars/test_result_set.py | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pyathena/polars/converter.py b/pyathena/polars/converter.py index 37659f34f..771c6e21b 100644 --- a/pyathena/polars/converter.py +++ b/pyathena/polars/converter.py @@ -48,8 +48,8 @@ def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any], strict: bool = True) -> _FrameT: """Convert timestamp text columns to Datetime dtypes, truncating finer fractions. - Athena writes up to 12 fractional digits, which Polars does not parse into a - Datetime whose time unit holds fewer. + Athena writes up to 12 fractional digits, and Polars does not parse more than + nine into a Datetime, so the text is cut to the digits that the time unit holds. Args: frame: The DataFrame or LazyFrame with the timestamp text columns, with NULL diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 8f2273754..3d29a7acb 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -550,8 +550,8 @@ def _get_converters( def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], dict[str, Any]]: """Get the Polars data types for reading CSV data, with the timestamp columns as text. - Athena writes up to 12 fractional digits, which Polars does not parse into a - Datetime whose time unit holds fewer. A chunked read always reads the timestamp + Athena writes up to 12 fractional digits, and Polars does not parse more than + nine into a Datetime. A chunked read always reads the timestamp columns as text and converts them with ``_to_datetimes()``. A whole read does so when parsing them fails, and with ``ignore_errors`` given to ``execute()``, with which Polars reads the timestamps that it does not parse as null. diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index f0a5443d3..adf7bea6c 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -182,7 +182,7 @@ def test_temporal_types(self, polars_cursor, unload): def test_temporal_types_chunked(self, polars_cursor): """Chunks of a CSV result have the temporal types and values of a whole read (#1191). - Timestamps with more fractional digits than the time unit are truncated. + Timestamps with up to 12 fractional digits are truncated to microseconds. """ types = {"timestamp": pl.Datetime("us"), "date": pl.Date()} columns = temporal.columns() diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 9d9b7c1c3..c54fd4496 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -612,8 +612,8 @@ def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected) def test_read_csv_ignore_errors(self, tmp_path, reader): """With ignore_errors, timestamp text that is not a timestamp becomes null. - Timestamps with more fractional digits than the time unit are still truncated, - which Polars itself reads as null with ignore_errors. + Timestamps with more than nine fractional digits, which Polars itself reads as + null with ignore_errors, are still truncated. """ path = tmp_path / "result.csv" path.write_bytes(b'"t"\n"2020-01-02 03:04:05.123456789012"\n"invalid"\n')