Read the timestamp columns of chunked Polars CSV results as text

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

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) :]]

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

@@ -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: