Preserve explicit ARRAY quantifier binds and clarify negation compatibility

Author: laughingman7743

docs/sqlalchemy.md

@@ -1113,8 +1113,9 @@ Comparisons use SQL three-valued logic: NULL elements can produce NULL when no d
 For an empty array, ANY is false and ALL is true.
 A NULL array produces NULL for both.
 SQLAlchemy comparison flipping is preserved.
-For boolean comparisons, SQLAlchemy can turn `~(any_(flags) == True)` into an element-wise `!= True` comparison before dialect compilation.
-To negate the whole match, explicitly group the comparison first: `~(any_(flags) == True).self_group()`.
+SQLAlchemy can rewrite negation into an element-wise comparison before dialect compilation, depending on the version and operand types.
+For example, `~(any_(flags) == True)` can become an element-wise `!= True` comparison; older SQLAlchemy 2.0 releases also rewrite non-boolean comparisons this way.
+To negate the whole match, always explicitly group the comparison first: `~(any_(flags) == True).self_group()`.
 Quantifiers over subqueries retain their usual SQL compilation.
 
 #### Querying ARRAY data

pyathena/sqlalchemy/compiler.py

@@ -318,7 +318,10 @@ def visit_binary(
                 predicate.right = Column(variable, item_type)
                 if isinstance(predicate.left, BindParameter) and (
                     isinstance(item_type, types.ARRAY)
-                    or self._array_type_inspector.array_type(predicate.left.type) is not None
+                    or (
+                        predicate.left.type is aggregate.element.type
+                        and predicate.left.type._type_affinity is not types.ARRAY
+                    )
                 ):
                     predicate.left = predicate.left._with_binary_element_type(item_type)
             if from_linter is not None and operators.is_comparison(binary.operator):

tests/pyathena/sqlalchemy/test_compiler.py

@@ -340,6 +340,15 @@ def test_multidimensional_array_quantifier_bind_type(self):
         items = column("items", AthenaArray(Integer, dimensions=2))
         assert "CAST(ARRAY[1, 2] AS ARRAY(INTEGER)) =" in self._compile_sql(items.any([1, 2]))
 
+    def test_quantifier_preserves_explicit_array_bind(self):
+        items = column("items", AthenaArray(Integer))
+        needle = literal([1, 2], items.type)
+        sql = self._compile_sql(needle == any_(func.array_agg(items)))
+        assert "CAST(ARRAY[1, 2] AS ARRAY(INTEGER)) = _pyathena_element_0" in sql
+        unknown = column("unknown", AthenaArray(types.NullType()))
+        sql = self._compile_sql(needle == any_(unknown))
+        assert "CAST(ARRAY[1, 2] AS ARRAY(INTEGER)) = _pyathena_element_0" in sql
+
     def test_subquery_any_remains_unchanged(self):
         sql = self._compile_sql(any_(select(column("item", Integer)).scalar_subquery()) == 2)
         assert "ANY (SELECT item)" in sql

tests/sqlalchemy/test_suite.py

@@ -182,9 +182,10 @@ def test_decorated_and_variant_arrays(self, connection):
             (1, (1, 2), True, True, (1, 2, 3, 4), True, (1, 2)),
         )
 
-    @pytest.mark.parametrize(
-        ("base_type", "item_type", "values", "needle"),
-        [(String(), String, ["a", "b"], "a"), (Integer(), Integer, [1, 2], 1)],
+    @sa_testing.combinations(
+        (String(), String, ["a", "b"], "a"),
+        (Integer(), Integer, [1, 2], 1),
+        argnames="base_type,item_type,values,needle",
     )
     def test_variant_array_quantifier_scalar_bind(
         self, connection, base_type, item_type, values, needle
@@ -263,6 +264,12 @@ def test_dimensions_zero_indexes_and_bound_index(self, connection):
         eq_(connection.execute(select(array.any([1, 2]))).scalar_one(), True)
         plain = literal([1, 2], types.ARRAY(Integer))
         eq_(connection.execute(select(plain[bindparam("index")]), {"index": 2}).scalar_one(), 2)
+        eq_(
+            connection.execute(
+                select(literal([1, 2], AthenaArray(Integer)) == any_(func.array_agg(plain)))
+            ).scalar_one(),
+            True,
+        )
 
     def test_quantified_comparisons(self, connection):
         cases = [