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 74a2e099f..490b63c70 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -546,29 +546,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. @@ -729,17 +741,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 @@ -906,14 +917,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 +935,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 +946,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()``. @@ -952,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 = 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 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..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(