From bc95db9a6441caf4e112754b16c0631106583d16 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 00:17:59 +0900 Subject: [PATCH 1/4] Read the timestamp columns of chunked Polars CSV results as text Athena writes up to 12 fractional digits for timestamp columns. A whole PolarsCursor CSV read caught the ComputeError from parsing them into a Datetime with a shorter time unit and read 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 the rows of the chunks already returned would be read again. Both readers now always read the timestamp columns as text and convert them with _to_datetimes(), the chunked reader for each chunk. Reading 5 million rows locally, this was faster than parsing them as Datetime with Polars 1.39.0, 1.44.2, and 2.0.0, for whole and chunked reads, and returned the same values and dtypes, including with time zones. _get_timestamp_dtypes() becomes _get_csv_schema_overrides(), which also returns the schema_overrides with those columns as String. Fixes #1191 Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 45 +++++++------- tests/pyathena/polars/test_cursor.py | 14 +++++ tests/pyathena/polars/test_result_set.py | 74 ++++++++++++++++++------ 3 files changed, 93 insertions(+), 40 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 74a2e099f..7ff9195ee 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -546,8 +546,12 @@ 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, so the timestamp columns are read as text + and converted with ``_to_datetimes()``. Args: has_header: Whether the CSV data has a header. Without one, or with @@ -556,19 +560,22 @@ def _get_timestamp_dtypes(self, has_header: bool) -> dict[str, Any]: 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_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. @@ -701,11 +708,12 @@ def _read_csv(self) -> pl.DataFrame: return pl.DataFrame() separator, has_header, new_columns = self._get_csv_params() + schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) kwargs = self._read_kwargs( lambda: self._csv_storage_options, separator=separator, has_header=has_header, - schema_overrides=self._csv_dtypes, + schema_overrides=schema_overrides, ) if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. @@ -726,20 +734,7 @@ 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: - timestamp_dtypes = self._get_timestamp_dtypes(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) + df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df @@ -906,14 +901,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, ) @@ -922,6 +919,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. @@ -932,6 +930,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,6 +954,7 @@ def _scan_csv_chunks( read_kwargs.pop("new_columns", None) lazy_df = pl.scan_csv(location, **read_kwargs) for batch in lazy_df.collect_batches(chunk_size=self._chunksize): + batch = _to_datetimes(batch, timestamp_dtypes) if new_columns: batch.columns = [*new_columns, *batch.columns[len(new_columns) :]] yield batch diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index c8183f07b..de8313456 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -638,6 +638,20 @@ def test_iter_chunks_data_consistency(self): # Should have multiple chunks assert len(chunked_dfs) > 1 + def test_iter_chunks_result_file_values(self, polars_cursor): + """Chunks of a CSV result file have the types and values of the whole result.""" + whole = polars_cursor.execute(RESULT_FILE_VALUES_QUERY).as_polars() + chunks = list(polars_cursor.execute(RESULT_FILE_VALUES_QUERY, chunksize=1).iter_chunks()) + assert all(chunk.height <= 1 for chunk in chunks) + chunked = pl.concat(chunks) + assert chunked.schema == whole.schema + assert chunked.equals(whole) + # Fractions finer than microseconds are truncated. + assert chunked["col_timestamp_12"].to_list() == [ + datetime(2020, 1, 2, 3, 4, 5, 123456), + None, + ] + def test_iter_chunks_chunk_sizes(self): """Test that chunks have correct sizes.""" with contextlib.closing(connect(schema_name=ENV.schema)) as conn: diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index e6ebdc915..f81324214 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 @@ -42,6 +43,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 +424,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 +455,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 +532,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": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), + ("_iter_csv_chunks", {}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), + ("_read_csv", {"columns": ["v"]}, {}), + ( + "_read_csv", + {"new_columns": ["t2", "v2"]}, + {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, + ), + ( + "_iter_csv_chunks", + {"new_columns": ["t2", "v2"]}, + {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, + ), + ( + "_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. 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. ``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,21 +576,27 @@ 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 From f5023e827cd62d7cf1a639d000173493045690b8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 00:25:42 +0900 Subject: [PATCH 2/4] Keep the timestamp columns of a Polars schema given to execute() Polars ignores schema_overrides when it is given a schema, so with a schema given to execute() the timestamp columns have the types of that schema, not String. The readers converted them with _to_datetimes() anyway, turning a String column into a Datetime one and failing on a Datetime column. As with schema_overrides, the columns are now not read as text when execute() is given a schema. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 10 +++++--- tests/pyathena/polars/test_result_set.py | 29 +++++++++++++----------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 7ff9195ee..926ada3a5 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -556,8 +556,8 @@ def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], d 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: Tuple of (schema_overrides, timestamp_dtypes): the data types to read the @@ -567,7 +567,11 @@ def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], d import polars as pl dtypes = self._csv_dtypes - if not has_header or self._kwargs.keys() & {"schema_overrides", "with_column_names"}: + if not has_header or self._kwargs.keys() & { + "schema", + "schema_overrides", + "with_column_names", + }: return dtypes, {} timestamp_dtypes = { name: dtype diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index f81324214..157a6522d 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -26,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: @@ -534,18 +540,16 @@ def test_csv_storage_options_replace_defaults(self): @pytest.mark.parametrize( ("reader", "kwargs", "expected"), [ - ("_read_csv", {}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ("_iter_csv_chunks", {}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), + ("_read_csv", {}, {"t": _TRUNCATED_TIMESTAMPS}), + ("_iter_csv_chunks", {}, {"t": _TRUNCATED_TIMESTAMPS}), ("_read_csv", {"columns": ["v"]}, {}), - ( - "_read_csv", - {"new_columns": ["t2", "v2"]}, - {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, - ), + ("_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", - {"new_columns": ["t2", "v2"]}, - {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, + {"schema": {"t": pl.String, "v": pl.String}}, + {"t": _TIMESTAMP_TEXTS}, ), ( "_read_csv", @@ -563,8 +567,8 @@ def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected) """Timestamps are read as text and truncated to the time unit. With ``with_column_names`` given to execute(), Polars renames the columns as it - reads them, so they are not read as text, and the read fails. ``scan_csv()`` - does not take ``columns``. + 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``. """ path = tmp_path / "result.csv" path.write_bytes(b'"t","v"\n"2020-01-02 03:04:05.123456789012","x"\n,"y"\n') @@ -598,8 +602,7 @@ def test_read_csv_truncates_timestamps(self, tmp_path, reader, kwargs, expected) return 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"] From 80424b2eb029ce5a731dd41137b8c6b68155f3c9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 08:21:16 +0900 Subject: [PATCH 3/4] Keep the whole Polars CSV read and convert chunks in the query plan Reading the timestamp columns of a whole read as text was faster than parsing them, but held the text and the converted column at once: reading 5 million rows locally, the peak RSS of a whole read grew by about 35-40% with Polars 1.39.0 and 2.0.0. The whole read again parses them natively and reads the data again with them as text only when parsing fails, as before #1191. The chunked read converts the text columns in the scan's query plan instead of converting each batch. _to_datetimes() takes a LazyFrame as well as a DataFrame and selects the columns with cs.by_name(require_all=False), so it skips a column the frame does not have without resolving the schema of a LazyFrame. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/converter.py | 22 ++++++++++++---------- pyathena/polars/result_set.py | 21 ++++++++++++++------- tests/pyathena/polars/test_converter.py | 11 ++++++++--- tests/pyathena/polars/test_result_set.py | 2 ++ 4 files changed, 36 insertions(+), 20 deletions(-) diff --git a/pyathena/polars/converter.py b/pyathena/polars/converter.py index a40e9c7d5..59e92058f 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,36 @@ } -def _to_datetimes(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame: +def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any]) -> _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. 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) .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 926ada3a5..b99577ce3 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -550,8 +550,9 @@ def _get_csv_schema_overrides(self, has_header: bool) -> tuple[dict[str, Any], d """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, so the timestamp columns are read as text - and converted with ``_to_datetimes()``. + 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 @@ -712,12 +713,11 @@ def _read_csv(self) -> pl.DataFrame: return pl.DataFrame() separator, has_header, new_columns = self._get_csv_params() - schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) kwargs = self._read_kwargs( lambda: self._csv_storage_options, separator=separator, has_header=has_header, - schema_overrides=schema_overrides, + schema_overrides=self._csv_dtypes, ) if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. @@ -738,7 +738,15 @@ 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() - df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes) + 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. + kwargs["schema_overrides"] = schema_overrides + df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df @@ -956,9 +964,8 @@ 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) for batch in lazy_df.collect_batches(chunk_size=self._chunksize): - batch = _to_datetimes(batch, timestamp_dtypes) if new_columns: batch.columns = [*new_columns, *batch.columns[len(new_columns) :]] yield batch diff --git a/tests/pyathena/polars/test_converter.py b/tests/pyathena/polars/test_converter.py index f025e001d..8dbb07648 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), diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 157a6522d..0ab2feb20 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -566,6 +566,8 @@ def test_csv_storage_options_replace_defaults(self): 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 read as text, and the read fails. With ``schema``, the columns have the types it gives. ``scan_csv()`` does not take ``columns``. From eee03465f89838bcb6c128009f8065916b7e3ff2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 08:36:45 +0900 Subject: [PATCH 4/4] Read invalid timestamp text in chunks as null with ignore_errors With ignore_errors given to execute(), Polars reads a value it cannot parse as null. The chunked CSV read now reads the timestamp columns as text, so text that is not a timestamp reached the strict conversion of _to_datetimes() and raised OperationalError instead. _to_datetimes() takes strict, which the readers set from ignore_errors. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/converter.py | 6 ++++-- pyathena/polars/result_set.py | 13 ++++++++++-- tests/pyathena/polars/test_converter.py | 9 ++++++++ tests/pyathena/polars/test_result_set.py | 27 ++++++++++++++++++++++++ 4 files changed, 51 insertions(+), 4 deletions(-) diff --git a/pyathena/polars/converter.py b/pyathena/polars/converter.py index 59e92058f..37659f34f 100644 --- a/pyathena/polars/converter.py +++ b/pyathena/polars/converter.py @@ -45,7 +45,7 @@ } -def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any]) -> _FrameT: +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 @@ -56,6 +56,8 @@ def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any]) -> _FrameT: 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 frame with the columns converted. @@ -71,7 +73,7 @@ def _to_datetimes(frame: _FrameT, dtypes: dict[str, Any]) -> _FrameT: 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 frame.with_columns(exprs) if exprs else frame diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index b99577ce3..490b63c70 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -746,7 +746,11 @@ def _read_csv(self) -> pl.DataFrame: raise # 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) + 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 @@ -964,7 +968,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 = _to_datetimes(pl.scan_csv(location, **read_kwargs), timestamp_dtypes) + 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 8dbb07648..d103d611f 100644 --- a/tests/pyathena/polars/test_converter.py +++ b/tests/pyathena/polars/test_converter.py @@ -61,3 +61,12 @@ def test_to_datetimes(dtype, microseconds, lazy): 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_result_set.py b/tests/pyathena/polars/test_result_set.py index 0ab2feb20..3d10bd0d1 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -608,6 +608,33 @@ 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.""" + 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(