feat(pyspark): add ArrayMap operation

ibis-project · Mar 26, 2023 · e2c159c · e2c159c
1 parent 2b1301e
commit e2c159c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -1669,6 +1669,15 @@ def compile_array_filter(t, op, **kwargs):
     )
 
 
+@compiles(ops.ArrayMap)
+def compile_array_map(t, op, **kwargs):
+    src_column = t.translate(op.arg, **kwargs)
+    return F.transform(
+        src_column,
+        lambda x: t.translate(op.result, arg_columns={op.parameter: x}, **kwargs),
+    )
+
+
 # --------------------------- Null Operations -----------------------------
 
 

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -466,7 +466,6 @@ def test_array_slice(con, start, stop):
         "mssql",
         "polars",
         "postgres",
-        "pyspark",
         "snowflake",
         "sqlite",
     ],