Support ORM ARRAY updates and preserve typed assignment expressions

Author: laughingman7743

docs/sqlalchemy.md

@@ -1132,6 +1132,7 @@ with engine.begin() as conn:
 ```
 
 Element assignment beyond the end extends the array and fills intervening positions with NULL.
+NULL padding uses Athena's `repeat()` function and follows its size limits.
 A NULL destination array is treated as empty for partial updates.
 Assigning `None` to an element stores NULL.
 Nested indices rebuild the corresponding inner arrays, including missing inner arrays.

pyathena/sqlalchemy/_array_update.py

@@ -32,6 +32,10 @@ def process(value):
     def literal_processor(self, dialect):
         return lambda value: _literal_complex(value, self.item_type, dialect)
 
+    def bind_expression(self, bindvalue):
+        expression = self.item_type.bind_expression(bindvalue)
+        return bindvalue if expression is None else expression
+
 
 class _IndexType(types.TypeDecorator[int]):
     impl = types.Integer
@@ -62,7 +66,9 @@ def __init__(self, column, path, value, value_type):
         self.type = column.type
         self.value_type = value_type
         self.value = (
-            value._with_binary_element_type(_AssignmentType(value_type))
+            value._with_binary_element_type(
+                _AssignmentType(value_type if value.type._isnull else value.type)
+            )
             if isinstance(value, BindParameter)
             else value
         )
@@ -89,7 +95,10 @@ def rewrite_array_update(statement):
             base = base.left
         name = base if isinstance(base, str) else getattr(base, "key", None)
         if path:
-            if not isinstance(base, Column) or base.table is not statement.table:
+            if (
+                not isinstance(base, Column)
+                or base.table._deannotate() is not statement.table._deannotate()
+            ):
                 raise exc.CompileError("ARRAY updates require a column of the target table")
             if name in seen:
                 raise exc.CompileError("Only one assignment per ARRAY column is supported")
@@ -119,6 +128,7 @@ def _index_sql(compiler, index: ColumnElement[Any], **kw):
     if (
         isinstance(index, BindParameter)
         and not index.required
+        and index.callable is None
         and (type(index.value) is not int or index.value <= 0)
     ):
         raise exc.CompileError("ARRAY write indices must be positive integers after normalization")
@@ -139,7 +149,12 @@ def compile_array_update(compiler, expression, **kw):
     final_slice = isinstance(expression.path[-1], Slice)
     if final_slice and (
         isinstance(value, Null)
-        or (isinstance(value, BindParameter) and not value.required and value.value is None)
+        or (
+            isinstance(value, BindParameter)
+            and not value.required
+            and value.callable is None
+            and value.value is None
+        )
     ):
         raise exc.CompileError("An ARRAY slice assignment requires a non-NULL array")
     rhs = compiler.process(value, **kw)
@@ -189,15 +204,17 @@ def rebuild(array, array_type, path):
                 f"concat({prefix}, {padding}, {rhs}, {suffix})", bound.step, array_type, **kw
             )
         index = _index_sql(compiler, bound, **kw)
-        variable = compiler._array_lambda_name()
         previous = f"element_at({array}, {index})"
         replacement = (
             rebuild(previous, _array_item_type(array_type), path[1:]) if len(path) > 1 else rhs
         )
-        return (
-            f"transform(sequence(1, greatest(cardinality({array}), {index})), "
-            f"{variable} -> IF({variable} = {index}, {replacement}, "
-            f"element_at({array}, {variable})))"
+        prefix = f"slice({array}, 1, least({index} - 1, cardinality({array})))"
+        element_type = compiler._complex_dml_type(_array_item_type(array_type))
+        padding = (
+            f"repeat(CAST(NULL AS {element_type}), "
+            f"CAST(greatest({index} - 1 - cardinality({array}), 0) AS INTEGER))"
         )
+        suffix = f"slice({array}, {index} + 1, greatest(cardinality({array}) - {index}, 0))"
+        return f"concat({prefix}, {padding}, ARRAY[{replacement}], {suffix})"
 
     return rebuild(compiler.process(expression.column, **kw), expression.type, expression.path)

tests/pyathena/sqlalchemy/test_array_update.py

@@ -1,6 +1,7 @@
 import pytest
-from sqlalchemy import Column, Integer, MetaData, Table, bindparam, func, types
+from sqlalchemy import Column, Integer, MetaData, Table, bindparam, func, types, update
 from sqlalchemy import exc as sa_exc
+from sqlalchemy.orm import declarative_base
 
 from pyathena.formatter import DefaultParameterFormatter
 from pyathena.sqlalchemy.base import AthenaDialect
@@ -79,15 +80,16 @@ def test_nested_and_zero_indexed_update():
     table = array_table(AthenaArray(Integer, dimensions=2, zero_indexes=True))
     statement = table.update().values({table.c["items"][0][2]: 7})
     sql = str(statement.compile(dialect=AthenaDialect(), compile_kwargs={"literal_binds": True}))
-    assert sql.count("transform(sequence(") == 2
+    assert "ARRAY[concat(" in sql
+    assert "sequence(" not in sql
     assert "IF(1 > 0, 1," in sql
     assert "IF(3 > 0, 3," in sql
 
 
 def test_generic_array_partial_update():
     table = array_table(types.ARRAY(Integer))
     sql = str(table.update().values({table.c["items"][1]: 2}).compile(dialect=AthenaDialect()))
-    assert "SET items=transform(" in sql
+    assert "SET items=concat(" in sql
 
 
 def test_write_index_expression_keeps_its_argument_types():
@@ -121,3 +123,48 @@ def test_binary_element_assignment_uses_native_hex_parameter():
         for name, value in compiled.params.items()
     }
     assert "FROM_HEX('00ff')" in DefaultParameterFormatter().format(str(compiled), params)
+
+
+class PrefixString(types.TypeDecorator):
+    impl = types.String
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect):
+        return f"prefix:{value}"
+
+    def bind_expression(self, bindvalue):
+        return func.upper(bindvalue)
+
+
+def test_explicit_assignment_type_and_callable_bindings():
+    table = array_table(AthenaArray(types.String))
+    stmt = table.update().values(
+        {
+            table.c["items"][bindparam("index", callable_=lambda: 1)]: bindparam(
+                "value", type_=PrefixString(), callable_=lambda: "a"
+            )
+        }
+    )
+    compiled = stmt.compile(dialect=AthenaDialect())
+    assert compiled._bind_processors["value"]("a") == "prefix:a"
+    assert "upper(%(value)s)" in str(compiled)
+    assert compiled.params["index"] == 1
+    assert compiled.params["value"] == "a"
+    table.update().values(
+        {table.c["items"][1:2]: bindparam("values", callable_=lambda: ["a"])}
+    ).compile(dialect=AthenaDialect())
+
+
+def test_orm_partial_update_and_renamed_attribute_conflicts():
+    base = declarative_base()
+
+    class Model(base):
+        __tablename__ = "arrays"
+        id = Column(Integer, primary_key=True)
+        values = Column("stored", AthenaArray(Integer), key="db_key")
+
+    sql = str(update(Model).values({Model.values[1]: 2}).compile(dialect=AthenaDialect()))
+    assert "UPDATE arrays SET stored=concat(" in sql
+    for whole in (Model.values, "values"):
+        with pytest.raises(sa_exc.CompileError, match="one assignment"):
+            update(Model).values({Model.values[1]: 2, whole: []}).compile(dialect=AthenaDialect())

tests/sqlalchemy/test_suite.py

@@ -23,10 +23,12 @@
     select,
     text,
     types,
+    update,
 )
 from sqlalchemy import Table as SATable
 from sqlalchemy import exc as sa_exc
 from sqlalchemy import testing as sa_testing
+from sqlalchemy.orm import Session, registry
 from sqlalchemy.sql.elements import quoted_name
 from sqlalchemy.testing import eq_, fixtures
 from sqlalchemy.testing.schema import Column, Table
@@ -275,6 +277,39 @@ def test_expression_values_and_indices(self, connection, metadata):
         connection.execute(table.update().values({items[1:2]: items[2:3].concat([4])}))
         eq_(connection.execute(select(table)).one(), (2, [2, 9, 4, 9], [b"\x00\xff"]))
 
+    def test_orm_and_long_array_update(self, connection, metadata):
+        table = Table(
+            "array_orm_updates",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("items", AthenaArray(Integer)),
+        )
+        table.create(connection)
+        connection.execute(
+            table.insert().values(
+                id=1,
+                items=func.concat(func.sequence(1, 10000), literal([10001], AthenaArray(Integer))),
+            )
+        )
+        mapping = registry()
+
+        class Record:
+            pass
+
+        mapping.map_imperatively(Record, table)
+        try:
+            with Session(bind=connection) as session:
+                session.execute(update(Record).where(Record.id == 1).values({Record.items[1]: 99}))
+                session.flush()
+            row = connection.execute(select(table.c["items"])).scalar_one()
+            eq_((len(row), row[0], row[-1]), (10001, 99, 10001))
+            connection.execute(
+                table.update().values({table.c["items"][2]: select(literal(77)).scalar_subquery()})
+            )
+            eq_(connection.execute(select(table.c["items"])).scalar_one()[:3], [99, 77, 3])
+        finally:
+            mapping.dispose()
+
     def test_null_slice_binding_rejected(self, connection, metadata):
         table = Table("array_null_slice", metadata, Column("items", types.ARRAY(Integer)))
         table.create(connection)