Fix SQLAlchemy true division semantics

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

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.

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