feat(pyspark): add ArrayFilter operation

ibis/backends/pyspark/compiler.py

@@ -1655,6 +1655,20 @@ def compile_array_collect(t, op, **kwargs):
     return F.collect_list(src_column)
 
 
+@compiles(ops.Argument)
+def compile_argument(t, op, arg_columns, **kwargs):
+    return arg_columns[op.name]
+
+
+@compiles(ops.ArrayFilter)
+def compile_array_filter(t, op, **kwargs):
+    src_column = t.translate(op.arg, **kwargs)
+    return F.filter(
+        src_column,
+        lambda x: t.translate(op.result, arg_columns={op.parameter: x}, **kwargs),
+    )
+
+
 # --------------------------- Null Operations -----------------------------
 
 

ibis/backends/pyspark/tests/test_array.py

@@ -166,3 +166,17 @@ def test_array_collect(client):
         .rename(columns={'array_int': 'collected'})
     )
     tm.assert_frame_equal(result, expected)
+
+
+def test_array_filter(client):
+    table = client.table('array_table')
+    expr = table.select(
+        table.array_int.filter(lambda item: item != 3).name('array_int')
+    )
+    result = expr.execute()
+    df = table.compile().toPandas()
+    df['array_int'] = df['array_int'].apply(
+        lambda ar: [item for item in ar if item != 3]
+    )
+    expected = df[['array_int']]
+    tm.assert_frame_equal(result, expected)

ibis/backends/tests/test_array.py

@@ -503,7 +503,6 @@ def test_array_map(backend, con):
         "mssql",
         "polars",
         "postgres",
-        "pyspark",
         "snowflake",
     ],
     raises=com.OperationNotDefinedError,