diff --git a/pyathena/sqlalchemy/compiler.py b/pyathena/sqlalchemy/compiler.py index d6df0a17..6a119eb3 100644 --- a/pyathena/sqlalchemy/compiler.py +++ b/pyathena/sqlalchemy/compiler.py @@ -10,7 +10,7 @@ IdentifierPreparer, SQLCompiler, ) -from sqlalchemy.sql.elements import BindParameter +from sqlalchemy.sql.elements import BindParameter, Cast from sqlalchemy.sql.schema import Column from pyathena.model import ( @@ -19,11 +19,10 @@ AthenaRowFormatSerde, ) from pyathena.sqlalchemy.preparer import AthenaDDLIdentifierPreparer -from pyathena.sqlalchemy.types import AthenaArray, AthenaMap, AthenaStruct +from pyathena.sqlalchemy.types import AthenaArray, AthenaMap, AthenaStruct, get_double_type if TYPE_CHECKING: from sqlalchemy import ( - Cast, CheckConstraint, ForeignKeyConstraint, PrimaryKeyConstraint, @@ -256,6 +255,38 @@ def visit_filter_func(self, fn: Function[Any], **kw: Any) -> str: return f"filter({array_sql}, {lambda_sql})" + def visit_truediv_binary(self, binary, operator, **kw): + """Render true division with explicit Athena numeric coercions.""" + left_type = binary.left.type + right_type = binary.right.type + + if isinstance(left_type, types.Float) or isinstance(right_type, types.Float): + division_type = get_double_type()() + return ( + self.process(Cast(binary.left, division_type), **kw) + + " / " + + self.process(Cast(binary.right, division_type), **kw) + ) + + left_is_numeric = isinstance(left_type, types.Numeric) + right_is_numeric = isinstance(right_type, types.Numeric) + if left_is_numeric or right_is_numeric: + division_type = binary.type + return ( + self.process(Cast(binary.left, division_type), **kw) + + " / " + + self.process(Cast(binary.right, division_type), **kw) + ) + + if isinstance(left_type, types.Integer) and isinstance(right_type, types.Integer): + return ( + self.process(binary.left, **kw) + + " / " + + self.process(Cast(binary.right, get_double_type()()), **kw) + ) + + return super().visit_truediv_binary(binary, operator, **kw) + def visit_cast(self, cast: Cast[Any], **kwargs): if (isinstance(cast.type, types.VARCHAR) and cast.type.length is None) or isinstance( cast.type, types.String @@ -265,6 +296,8 @@ def visit_cast(self, cast: Cast[Any], **kwargs): type_clause = "CHAR" elif isinstance(cast.type, (types.BINARY, types.VARBINARY)): type_clause = "VARBINARY" + elif hasattr(types, "DOUBLE") and isinstance(cast.type, types.DOUBLE): + type_clause = "DOUBLE" elif isinstance(cast.type, (types.FLOAT, types.Float, types.REAL)): # https://docs.aws.amazon.com/athena/latest/ug/data-types.html # In Athena, use float in DDL statements like CREATE TABLE diff --git a/tests/pyathena/sqlalchemy/test_compiler.py b/tests/pyathena/sqlalchemy/test_compiler.py index 887af41b..1fb120f5 100644 --- a/tests/pyathena/sqlalchemy/test_compiler.py +++ b/tests/pyathena/sqlalchemy/test_compiler.py @@ -1,9 +1,21 @@ from unittest.mock import Mock import pytest -from sqlalchemy import Column, Date, Integer, MetaData, String, Table, exc, func, select +from sqlalchemy import ( + Column, + Date, + Float, + Integer, + MetaData, + Numeric, + String, + Table, + exc, + func, + select, +) from sqlalchemy.engine.url import make_url -from sqlalchemy.sql import literal +from sqlalchemy.sql import literal, literal_column from sqlalchemy.sql.ddl import CreateTable from pyathena.sqlalchemy.base import AthenaDialect @@ -229,6 +241,39 @@ def test_visit_char_length_func_existing(self): sql_str = str(compiled) assert "length(" in sql_str + @pytest.mark.parametrize( + ("expression", "expected"), + [ + ( + literal_column("15", type_=Integer()) / literal_column("10", type_=Integer()), + "SELECT 15 / CAST(10 AS DOUBLE) AS anon_1", + ), + ( + literal(15) / literal(10), + "SELECT 15 / CAST(10 AS DOUBLE) AS anon_1", + ), + ( + literal_column("5.52", type_=Numeric(10, 2)) + / literal_column("2.4", type_=Numeric(10, 2)), + "SELECT CAST(5.52 AS DECIMAL(10, 2)) / CAST(2.4 AS DECIMAL(10, 2)) AS anon_1", + ), + ( + literal_column("5.52", type_=Numeric(10, 2)) / literal_column("2", type_=Integer()), + "SELECT CAST(5.52 AS DECIMAL(10, 2)) / CAST(2 AS DECIMAL(10, 2)) AS anon_1", + ), + ( + literal_column("5.52", type_=Float()) / literal_column("2.4", type_=Float()), + "SELECT CAST(5.52 AS DOUBLE) / CAST(2.4 AS DOUBLE) AS anon_1", + ), + ], + ) + def test_visit_truediv_binary(self, expression, expected): + compiled = select(expression).compile( + dialect=self.dialect, compile_kwargs={"literal_binds": True} + ) + + assert str(compiled) == expected + class TestAthenaDDLCompiler: """Compile-only (no AWS) tests for the DDL compiler's S3 Tables support. diff --git a/tests/sqlalchemy/test_suite.py b/tests/sqlalchemy/test_suite.py index 49077eef..10be6405 100644 --- a/tests/sqlalchemy/test_suite.py +++ b/tests/sqlalchemy/test_suite.py @@ -5,7 +5,6 @@ from sqlalchemy.testing.suite import InsertBehaviorTest as _InsertBehaviorTest from sqlalchemy.testing.suite import IntegerTest as _IntegerTest from sqlalchemy.testing.suite import StringTest as _StringTest -from sqlalchemy.testing.suite import TrueDivTest as _TrueDivTest del BinaryTest # noqa: F821 del ComponentReflectionTest # noqa: F821 @@ -48,26 +47,6 @@ def test_no_results_for_non_returning_insert(self, connection, style, executeman pass -class TrueDivTest(_TrueDivTest): - @pytest.mark.skip("Athena returns an integer for operations between integers.") - def test_truediv_integer(self, connection, left, right, expected): - pass - - @pytest.mark.skip("Athena returns an integer for operations between integers.") - def test_truediv_integer_bound(self, connection): - pass - - @pytest.mark.skip("TODO") - def test_truediv_numeric(self, connection, left, right, expected): - # TODO - pass - - @pytest.mark.skip("TODO") - def test_truediv_float(self, connection, left, right, expected): - # TODO: AssertionError: 2.299999908606215 != 2.3 - pass - - class FetchLimitOffsetTest(_FetchLimitOffsetTest): @pytest.mark.skip("Athena does not support expressions in the offset clause.") def test_simple_limit_expr_offset(self, connection):