Convert the timestamp text of chunked Polars CSV reads in the scan

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.
+    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:
-        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):

pyathena/polars/result_set.py

@@ -547,29 +547,42 @@ 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, 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.
 
         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.
@@ -728,20 +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:
-                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)
+            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=strict)
             if new_columns:
                 df.columns = [*new_columns, *df.columns[len(new_columns) :]]
             return df
@@ -908,14 +922,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 +940,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 +951,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 +973,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) :]]

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,11 +50,23 @@ 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),
         datetime(1, 1, 1),
         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]

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 up to 12 fractional digits are truncated to microseconds.
+        """
+        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(
             """