# Copyright 2019 The PyAthena authors
#
# Licensed under the MIT License.
# See LICENSE or https://opensource.org/licenses/MIT.
#
# SPDX-License-Identifier: MIT
"""Type converters for Polars cursor results."""
from __future__ import annotations
import logging
from collections.abc import Callable
from copy import deepcopy
from typing import TYPE_CHECKING, Any
from pyathena.converter import (
_TIMESTAMP_TEXT_LENGTHS,
Converter,
_to_binary,
_to_date,
_to_datetime_with_tz,
_to_default,
_to_json,
_to_time,
_to_time_with_tz,
)
from pyathena.util import override
if TYPE_CHECKING:
import polars as pl
_logger = logging.getLogger(__name__)
_DEFAULT_POLARS_CONVERTERS: dict[str, Callable[[str | None], Any | None]] = {
"date": _to_date,
"time": _to_time,
"time with time zone": _to_time_with_tz,
"timestamp with time zone": _to_datetime_with_tz,
"varbinary": _to_binary,
"json": _to_json,
}
def _to_datetimes(df: pl.DataFrame, dtypes: dict[str, Any]) -> pl.DataFrame:
"""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.
dtypes: The Datetime dtypes keyed by column name.
Returns:
The DataFrame with the columns converted.
"""
import polars as pl
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)
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
[docs]
class DefaultPolarsTypeConverter(Converter):
"""Optimized type converter for Polars DataFrame results.
This converter is specifically designed for the PolarsCursor and provides
optimized type conversion for Polars DataFrames.
The converter focuses on:
- Mapping Athena types to Polars dtypes when reading CSV results,
including ``decimal`` as ``pl.Decimal(precision, scale)``
- Converting date, time, and varbinary values, and parsing JSON values,
in rows returned by fetchone(), fetchmany(), and fetchall()
- Keeping array, map, and row values as strings
Example:
>>> from pyathena.polars.converter import DefaultPolarsTypeConverter
>>> converter = DefaultPolarsTypeConverter()
>>>
>>> # Used automatically by PolarsCursor
>>> cursor = connection.cursor(PolarsCursor)
>>> # converter is applied automatically to results
Note:
This converter is used by default in PolarsCursor.
Most users don't need to instantiate it directly.
"""
[docs]
def __init__(self) -> None:
"""Initialize the converter with the default Polars conversion functions and dtypes."""
super().__init__(
mappings=deepcopy(_DEFAULT_POLARS_CONVERTERS),
default=_to_default,
types=self._dtypes,
)
@property
def _dtypes(self) -> dict[str, Any]:
import polars as pl
if not hasattr(self, "__dtypes"):
self.__dtypes = {
"boolean": pl.Boolean,
"tinyint": pl.Int8,
"smallint": pl.Int16,
"integer": pl.Int32,
"bigint": pl.Int64,
"float": pl.Float32,
"real": pl.Float64,
"double": pl.Float64,
"char": pl.String,
"varchar": pl.String,
"string": pl.String,
"timestamp": pl.Datetime,
"date": pl.Date,
"time": pl.String,
"time with time zone": pl.String,
"timestamp with time zone": pl.String,
"varbinary": pl.String,
"array": pl.String,
"map": pl.String,
"row": pl.String,
"decimal": pl.Decimal,
"json": pl.String,
}
return self.__dtypes
[docs]
@override
def get_dtype(self, type_: str, precision: int = 0, scale: int = 0) -> Any:
"""Get the Polars data type for a given Athena type.
Args:
type_: The Athena data type name.
precision: The precision for decimal types.
scale: The scale for decimal types.
Returns:
The Polars data type.
"""
import polars as pl
if type_ == "decimal":
return pl.Decimal(precision=precision, scale=scale)
return self._types.get(type_)
[docs]
@override
def convert(self, type_: str, value: str | None, type_hint: str | None = None) -> Any | None:
converter = self.get(type_)
return converter(value)
[docs]
class DefaultPolarsUnloadTypeConverter(Converter):
"""Type converter for Polars UNLOAD operations.
This converter is designed for use with UNLOAD queries that write
results directly to Parquet files in S3. Since UNLOAD operations
bypass the normal conversion process and write data in native
Parquet format, this converter has minimal functionality.
Note:
Used automatically when PolarsCursor is configured with unload=True.
UNLOAD results are read directly as Polars DataFrames from Parquet files.
"""
[docs]
def __init__(self) -> None:
"""Initialize the converter with no type mappings."""
super().__init__(
mappings={},
default=_to_default,
)
[docs]
@override
def convert(self, type_: str, value: str | None, type_hint: str | None = None) -> Any | None:
converter = self.get(type_)
return converter(value)