feat(snowflake): implement ops.ArrayRepeat

ibis-project · Aug 5, 2023 · a93cbd6 · a93cbd6
1 parent 6f7e13d
commit a93cbd6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
@@ -17,7 +17,7 @@
 import sqlalchemy as sa
 import sqlalchemy.types as sat
 from snowflake.connector.constants import FIELD_ID_TO_NAME
-from snowflake.sqlalchemy import ARRAY, OBJECT, URL
+from snowflake.sqlalchemy import ARRAY, DOUBLE, OBJECT, URL
 from sqlalchemy.ext.compiler import compiles
 
 import ibis
@@ -96,6 +96,13 @@ class SnowflakeCompiler(AlchemyCompiler):
         "returns": ARRAY,
         "source": """return array.sort();""",
     },
+    "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();""",
+    },
 }
 
 

diff --git a/ibis/backends/snowflake/registry.py b/ibis/backends/snowflake/registry.py
@@ -462,6 +462,7 @@ def _map_get(t, op):
         ops.TableColumn: _table_column,
         ops.Levenshtein: fixed_arity(sa.func.editdistance, 2),
         ops.ArraySort: unary(sa.func.ibis_udfs.public.array_sort),
+        ops.ArrayRepeat: fixed_arity(sa.func.ibis_udfs.public.array_repeat, 2),
     }
 )
 
@@ -471,8 +472,6 @@ def _map_get(t, op):
     ops.CumulativeAny,
     ops.CumulativeOp,
     ops.NTile,
-    # ibis.expr.operations.array
-    ops.ArrayRepeat,
     # ibis.expr.operations.reductions
     ops.MultiQuantile,
     # ibis.expr.operations.strings

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -84,9 +84,7 @@ def test_array_scalar(con, backend):
         assert con.execute(expr.typeof()) == ARRAY_BACKEND_TYPES[backend_name]
 
 
-@pytest.mark.notimpl(
-    ["polars", "datafusion", "snowflake"], raises=com.OperationNotDefinedError
-)
+@pytest.mark.notimpl(["polars", "datafusion"], raises=com.OperationNotDefinedError)
 def test_array_repeat(con):
     expr = ibis.array([1.0, 2.0]) * 2