fix(snowflake): fix new array operations; remove ArrayRemove operation

ibis/backends/snowflake/__init__.py

@@ -89,12 +89,19 @@ def _format_in_memory_table(self, op, ref_op, translator):
         if _NATIVE_ARROW:
             return super()._format_in_memory_table(op, ref_op, translator)
 
-        columns = translator._schema_to_sqlalchemy_columns(ref_op.schema)
-        rows = list(ref_op.data.to_frame().itertuples(index=False))
-        pos_columns = [
-            sa.column(f"${idx}") for idx in range(1, len(ref_op.schema.names) + 1)
-        ]
-        return sa.select(*pos_columns).select_from(sa.values(*columns).data(rows))
+        import ibis
+
+        schema = ref_op.schema
+        selects = (
+            sa.select(
+                *(
+                    translator.translate(ibis.literal(col, typ).op()).label(name)
+                    for col, (name, typ) in zip(row, schema.items())
+                )
+            )
+            for row in ref_op.data.to_frame().itertuples(index=False)
+        )
+        return sa.union_all(*selects)
 
 
 class SnowflakeCompiler(AlchemyCompiler):

ibis/backends/snowflake/registry.py

@@ -332,15 +332,17 @@ def _group_concat(t, op):
                 sa.func.ifnull(arg, sa.func.parse_json("null")), type_=ARRAY
             )
         ),
-        ops.ArrayContains: fixed_arity(sa.func.array_contains, 2),
-        ops.ArrayPosition: fixed_arity(
-            lambda lst, el: sa.func.array_position(lst, el) - 1, 2
+        ops.ArrayContains: fixed_arity(
+            lambda arr, el: sa.func.array_contains(sa.func.to_variant(el), arr), 2
         ),
-        ops.ArrayDistinct: fixed_arity(sa.func.array_distinct, 1),
-        ops.ArrayRemove: fixed_arity(
-            lambda lst, el: sa.func.array_except(lst, sa.func.array_construct(el)),
+        ops.ArrayPosition: fixed_arity(
+            # snowflake is zero-based here, so we don't need to substract 1 from the result
+            lambda lst, el: sa.func.coalesce(
+                sa.func.array_position(sa.func.to_variant(el), lst), -1
+            ),
             2,
         ),
+        ops.ArrayDistinct: fixed_arity(sa.func.array_distinct, 1),
         ops.ArrayUnion: fixed_arity(
             lambda left, right: sa.func.array_distinct(sa.func.array_cat(left, right)),
             2,
@@ -396,6 +398,7 @@ def _group_concat(t, op):
     ops.CumulativeOp,
     ops.NTile,
     # ibis.expr.operations.array
+    ops.ArrayRemove,
     ops.ArrayRepeat,
     ops.ArraySort,
     # ibis.expr.operations.reductions

ibis/backends/tests/test_array.py

@@ -577,6 +577,7 @@ def test_array_position(backend, con):
         "pandas",
         "polars",
         "postgres",
+        "snowflake",
     ],
     raises=com.OperationNotDefinedError,
 )
@@ -629,6 +630,7 @@ def test_array_unique(backend, con):
         "pandas",
         "polars",
         "postgres",
+        "snowflake",
     ],
     raises=com.OperationNotDefinedError,
 )
@@ -659,7 +661,7 @@ def test_array_sort(backend, con):
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.broken(
-    ["trino", "pyspark"],
+    ["snowflake", "trino", "pyspark"],
     raises=AssertionError,
     reason="array_distinct([NULL]) seems to differ from other backends",
 )