apacheGH-42013 [Python] Allow Array.filter() to take general array in…

…put (apache#42051) ### What changes are included in this PR? Allow Array.filter() to take general array input. ### Are these changes tested? Unit test added, via CI. ### Are there any user-facing changes? No * GitHub Issue: apache#42013 Authored-by: Kelvin Wu <[email protected]> Signed-off-by: Joris Van den Bossche <[email protected]>
mathworks · Jun 11, 2024 · 03a960d · 03a960d
1 parent baf4089
commit 03a960d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
@@ -1444,7 +1444,7 @@ cdef class Array(_PandasConvertible):
         """
         return _pc().drop_null(self)
 
-    def filter(self, Array mask, *, null_selection_behavior='drop'):
+    def filter(self, object mask, *, null_selection_behavior='drop'):
         """
         Select values from an array.
 

diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py
@@ -1305,6 +1305,12 @@ def test_filter(ty, values):
     result.validate()
     assert result.equals(pa.array([values[0], values[3], None], type=ty))
 
+    # same test with different array type
+    mask = np.array([True, False, False, True, None])
+    result = arr.filter(mask, null_selection_behavior='drop')
+    result.validate()
+    assert result.equals(pa.array([values[0], values[3]], type=ty))
+
     # non-boolean dtype
     mask = pa.array([0, 1, 0, 1, 0])
     with pytest.raises(NotImplementedError):