This is an automated email from the ASF dual-hosted git repository.
AlenkaF pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 4f3b0f01b5 GH-50072:[Python] Add tests for replace_with_mask kernel
(#50102)
4f3b0f01b5 is described below
commit 4f3b0f01b51c8d596334027f939fa97f45b5ac05
Author: AnkitAhlawat <[email protected]>
AuthorDate: Thu Jun 11 16:11:14 2026 +0530
GH-50072:[Python] Add tests for replace_with_mask kernel (#50102)
### Fix : https://github.com/apache/arrow/issues/50072
### Rationale for this change
The `replace_with_mask` compute function in PyArrow currently has no
Python-level tests for regular use cases.
This PR adds test cases for standard usage of `replace_with_mask` to ensure
the function works correctly.
### What changes are included in this PR?
Added parameterized test cases in `python/pyarrow/tests/test_compute.py`
### Are these changes tested?
Yes, Manually
### Are there any user-facing changes?
No user-facing changes.
* GitHub Issue: #50072
Authored-by: [email protected] <[email protected]>
Signed-off-by: AlenkaF <[email protected]>
---
python/pyarrow/tests/test_compute.py | 118 +++++++++++++++++++++++++++++++++++
1 file changed, 118 insertions(+)
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index 4e7f506ba5..478a0c3dc5 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -1297,6 +1297,124 @@ def test_replace_with_mask_null_type():
assert result.to_pylist() == [None]
+def test_replace_with_mask_basic():
+ """Test basic replacement with array mask."""
+ arr = pa.array([1, 2, 3, 4, 5])
+ mask = pa.array([True, False, True, False, True])
+ replacements = pa.array([10, 20, 30])
+ expected = pa.array([10, 2, 20, 4, 30])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_scalar_mask_true():
+ """Test replacement with scalar mask True."""
+ arr = pa.array([1, 2, 3])
+ mask = True
+ replacements = pa.array([10, 20, 30])
+ expected = pa.array([10, 20, 30])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_scalar_mask_false():
+ """Test replacement with scalar mask False."""
+ arr = pa.array([1, 2, 3])
+ mask = False
+ replacements = pa.array([], type=pa.int64())
+ expected = pa.array([1, 2, 3])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_scalar_replacement():
+ """Test replacement with scalar replacement value."""
+ arr = pa.array([1, 2, 3, 4])
+ mask = pa.array([True, False, True, False])
+ replacements = pa.scalar(99)
+ expected = pa.array([99, 2, 99, 4])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_null_in_array():
+ """Test null handling in input array."""
+ arr = pa.array([1, None, 3, None, 5])
+ mask = pa.array([False, True, False, True, True])
+ replacements = pa.array([10, 20, 30])
+ expected = pa.array([1, 10, 3, 20, 30])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_null_in_mask():
+ """Test null handling in mask."""
+ arr = pa.array([1, 2, 3, 4, 5, 6])
+ mask = pa.array([False, False, None, None, True, True])
+ replacements = pa.array([10, None])
+ expected = pa.array([1, 2, None, None, 10, None])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_string_type():
+ """Test replacement with string type."""
+ arr = pa.array(['a', 'b', 'c', 'd'])
+ mask = pa.array([True, False, True, False])
+ replacements = pa.array(['x', 'y'])
+ expected = pa.array(['x', 'b', 'y', 'd'])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_float_type():
+ """Test replacement with float type."""
+ arr = pa.array([1.1, 2.2, 3.3, 4.4])
+ mask = pa.array([True, False, True, False])
+ replacements = pa.array([10.5, 20.5])
+ expected = pa.array([10.5, 2.2, 20.5, 4.4])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_chunked_array_multiple_chunks():
+ """Test replace_with_mask with ChunkedArray with multiple chunks."""
+ arr = pa.chunked_array([[1, 2, 3], [4, 5, 6]])
+ mask = pa.array([True, False, False, False, True, False])
+ replacements = pa.array([10, 20])
+ expected = pa.chunked_array([[10, 2, 3], [4, 20, 6]])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_chunked_array_empty_chunks():
+ """Test replace_with_mask with ChunkedArray with empty chunks."""
+ arr = pa.chunked_array([[1, 2], [], [3, 4]])
+ mask = pa.array([True, False, True, False])
+ replacements = pa.array([10, 20])
+ expected = pa.chunked_array([[10, 2], [20, 4]])
+ result = pc.replace_with_mask(arr, mask, replacements)
+ assert result.equals(expected)
+
+
+def test_replace_with_mask_error_replacement_count_mismatch():
+ """Replacement count does not match true values in mask."""
+ arr = pa.array([1, 2, 3])
+ mask = pa.array([True, True, False])
+ replacements = pa.array([10])
+ with pytest.raises(pa.ArrowInvalid, match="expected 2.*but got 1"):
+ pc.replace_with_mask(arr, mask, replacements)
+
+
+def test_replace_with_mask_error_mask_length_mismatch():
+ """Mask length does not match input array length."""
+ arr = pa.array([1, 2, 3])
+ mask = pa.array([True, False])
+ replacements = pa.array([10])
+ with pytest.raises(pa.ArrowInvalid):
+ pc.replace_with_mask(arr, mask, replacements)
+
+
def test_binary_join():
ar_list = pa.array([['foo', 'bar'], None, []])
expected = pa.array(['foo-bar', None, ''])
