Fix ARRAY variant quantifier binds and document expression limits

Author: laughingman7743

docs/sqlalchemy.md

@@ -1104,14 +1104,17 @@ Slice stops are inclusive, following SQLAlchemy's SQL array convention.
 Omitted boundaries mean the beginning or end of the array; explicit boundaries are clipped to the array.
 A reversed slice returns an empty array, and slicing a NULL array returns NULL.
 Only `step=None` and `step=1` are supported.
+Slice SQL can evaluate the array and boundary expressions more than once; use deterministic expressions rather than volatile functions such as `random()` in slices.
 Use `AthenaArray` for open-ended slices with `zero_indexes=True`; SQLAlchemy's generic `ARRAY` comparator requires explicit bounds in that mode.
 Indexing a multidimensional array retains the remaining dimensions, while slicing retains its array type.
 
 `any_(array)` and `all_(array)`, and the legacy `array.any(value)` and `array.all(value)` methods, compile to Athena's `any_match` and `all_match` functions.
 Comparisons use SQL three-valued logic: NULL elements can produce NULL when no decisive true or false result exists.
 For an empty array, ANY is false and ALL is true.
 A NULL array produces NULL for both.
-SQLAlchemy comparison flipping and negation are preserved.
+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()`.
 Quantifiers over subqueries retain their usual SQL compilation.
 
 #### Querying ARRAY data

pyathena/sqlalchemy/compiler.py

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

tests/pyathena/sqlalchemy/test_array.py

@@ -277,16 +277,22 @@ def test_decorated_array_index_and_slice(self):
 
 class TestArrayTypeInspector:
     @pytest.mark.parametrize(
-        "type_", [TupleArray(), String().with_variant(TupleArray(), "awsathena")]
+        ("type_", "value", "expected"),
+        [
+            (TupleArray(), 2, "2"),
+            (String().with_variant(TupleArray(), "awsathena"), 2, "2"),
+            (Integer().with_variant(TupleArray(), "awsathena"), 2, "2"),
+            (String().with_variant(AthenaArray(String), "awsathena"), "a", "'a'"),
+        ],
     )
-    def test_decorated_and_variant_array_quantifiers(self, type_):
+    def test_decorated_and_variant_array_quantifiers(self, type_, value, expected):
         items = column("items", type_)
-        statement = select(any_(items) == 2, all_(items) > 0)
+        statement = select(any_(items) == value, all_(items) > value)
         sql = str(
             statement.compile(dialect=AthenaDialect(), compile_kwargs={"literal_binds": True})
         )
-        assert "any_match((items), _pyathena_element_0 -> 2 = _pyathena_element_0)" in sql
-        assert "all_match((items), _pyathena_element_1 -> 0 < _pyathena_element_1)" in sql
+        assert f"any_match((items), _pyathena_element_0 -> {expected} = _pyathena_element_0)" in sql
+        assert f"all_match((items), _pyathena_element_1 -> {expected} < _pyathena_element_1)" in sql
 
     @pytest.mark.parametrize(
         ("type_", "ddl", "dml"),

tests/pyathena/sqlalchemy/test_compiler.py

@@ -362,6 +362,13 @@ def test_quantifier_boolean_left_operand(self):
         assert "any_match" in self._compile_sql(any_(flags) == True)  # noqa: E712
         assert "all_match" in self._compile_sql(all_(flags) != False)  # noqa: E712
         assert "IS DISTINCT FROM" in self._compile_sql(any_(flags).is_distinct_from(True))
+        comparison = any_(flags) == True  # noqa: E712
+        assert self._compile_sql(~comparison) == (
+            "any_match((flags), _pyathena_element_0 -> _pyathena_element_0 != true)"
+        )
+        assert self._compile_sql(~comparison.self_group()) == (
+            "NOT (any_match((flags), _pyathena_element_0 -> _pyathena_element_0 = true))"
+        )
 
     def test_quantifier_join_linter_tracks_original_tables(self):
         left = table("left_table", column("value", Integer))

tests/sqlalchemy/test_suite.py

@@ -182,6 +182,18 @@ 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)],
+    )
+    def test_variant_array_quantifier_scalar_bind(
+        self, connection, base_type, item_type, values, needle
+    ):
+        array = literal(values, base_type.with_variant(AthenaArray(item_type), "awsathena"))
+        statement = select(any_(array) == needle, all_(array) == needle)
+        eq_(tuple(connection.execute(statement).one()), (True, False))
+        eq_(tuple(connection.execute(statement).one()), (True, False))
+
     def test_cached_steps_and_boolean_quantifiers(self, connection):
         inferred = func.array_agg(func.length(literal("abc")))
         eq_(connection.execute(select(inferred[1:2:1])).scalar_one(), [3])
@@ -191,9 +203,14 @@ def test_cached_steps_and_boolean_quantifiers(self, connection):
             connection.execute(select(array[1:2:2])).all()
         assert isinstance(error.value.orig, ValueError)
         flags = literal([True, False], AthenaArray(types.Boolean))
+        comparison = any_(flags) == True  # noqa: E712
         eq_(
-            tuple(connection.execute(select(any_(flags) == True, all_(flags) == True)).one()),  # noqa: E712
-            (True, False),
+            tuple(
+                connection.execute(
+                    select(comparison, all_(flags) == True, ~comparison, ~comparison.self_group())  # noqa: E712
+                ).one()
+            ),
+            (True, False, True, False),
         )
 
     def test_index_slice_and_concat(self, connection):