test(arrays): add array repeat test on column expression (#9182)

ibis/backends/snowflake/__init__.py

@@ -64,6 +64,13 @@
     return Object.assign(...keys.map((key, j) => ({[key]: arrays[j][i]})));
 })""",
     },
+    "ibis_udfs.public.array_repeat": {
+        # Integer inputs are not allowed because JavaScript only supports
+        # doubles
+        "inputs": {"value": "ARRAY", "count": "DOUBLE"},
+        "returns": "ARRAY",
+        "source": """return Array(count).fill(value).flat();""",
+    },
 }
 
 

ibis/backends/snowflake/compiler.py

@@ -305,12 +305,7 @@ def visit_TypeOf(self, op, *, arg):
         return self.f.typeof(self.f.to_variant(arg))
 
     def visit_ArrayRepeat(self, op, *, arg, times):
-        return self.f.array_flatten(
-            self.f.transform(
-                self.f.array_generate_range(0, times),
-                sge.Lambda(this=arg, expressions=[sg.to_identifier("__arg__")]),
-            )
-        )
+        return self.f.udf.array_repeat(arg, times)
 
     def visit_ArrayUnion(self, op, *, left, right):
         return self.f.array_distinct(self.f.array_cat(left, right))

ibis/backends/tests/test_array.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+from collections import Counter
 from datetime import datetime
 
 import numpy as np
@@ -117,6 +118,16 @@ def test_array_repeat(con):
     assert np.array_equal(result, expected)
 
 
+@pytest.mark.notimpl(["flink", "polars"], raises=com.OperationNotDefinedError)
+def test_array_repeat_column(con):
+    t = ibis.memtable({"x": [[1.0, 2.0]]}, schema=ibis.schema({"x": "array<float64>"}))
+    expr = (t.x * 2).name("tmp")
+
+    result = con.execute(expr.name("tmp")).iat[0]
+    expected = np.array([1.0, 2.0, 1.0, 2.0])
+    assert Counter(result) == Counter(expected)
+
+
 def test_array_concat(con):
     left = ibis.literal([1, 2, 3])
     right = ibis.literal([2, 1])