Skip to content

Commit 4f3b0f0

Browse files
GH-50072:[Python] Add tests for replace_with_mask kernel (#50102)
### Fix : #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: Ankit.Ahlawat@ibm.com <Ankit.Ahlawat@ibm.com> Signed-off-by: AlenkaF <frim.alenka@gmail.com>
1 parent 108240f commit 4f3b0f0

1 file changed

Lines changed: 118 additions & 0 deletions

File tree

python/pyarrow/tests/test_compute.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1297,6 +1297,124 @@ def test_replace_with_mask_null_type():
12971297
assert result.to_pylist() == [None]
12981298

12991299

1300+
def test_replace_with_mask_basic():
1301+
"""Test basic replacement with array mask."""
1302+
arr = pa.array([1, 2, 3, 4, 5])
1303+
mask = pa.array([True, False, True, False, True])
1304+
replacements = pa.array([10, 20, 30])
1305+
expected = pa.array([10, 2, 20, 4, 30])
1306+
result = pc.replace_with_mask(arr, mask, replacements)
1307+
assert result.equals(expected)
1308+
1309+
1310+
def test_replace_with_mask_scalar_mask_true():
1311+
"""Test replacement with scalar mask True."""
1312+
arr = pa.array([1, 2, 3])
1313+
mask = True
1314+
replacements = pa.array([10, 20, 30])
1315+
expected = pa.array([10, 20, 30])
1316+
result = pc.replace_with_mask(arr, mask, replacements)
1317+
assert result.equals(expected)
1318+
1319+
1320+
def test_replace_with_mask_scalar_mask_false():
1321+
"""Test replacement with scalar mask False."""
1322+
arr = pa.array([1, 2, 3])
1323+
mask = False
1324+
replacements = pa.array([], type=pa.int64())
1325+
expected = pa.array([1, 2, 3])
1326+
result = pc.replace_with_mask(arr, mask, replacements)
1327+
assert result.equals(expected)
1328+
1329+
1330+
def test_replace_with_mask_scalar_replacement():
1331+
"""Test replacement with scalar replacement value."""
1332+
arr = pa.array([1, 2, 3, 4])
1333+
mask = pa.array([True, False, True, False])
1334+
replacements = pa.scalar(99)
1335+
expected = pa.array([99, 2, 99, 4])
1336+
result = pc.replace_with_mask(arr, mask, replacements)
1337+
assert result.equals(expected)
1338+
1339+
1340+
def test_replace_with_mask_null_in_array():
1341+
"""Test null handling in input array."""
1342+
arr = pa.array([1, None, 3, None, 5])
1343+
mask = pa.array([False, True, False, True, True])
1344+
replacements = pa.array([10, 20, 30])
1345+
expected = pa.array([1, 10, 3, 20, 30])
1346+
result = pc.replace_with_mask(arr, mask, replacements)
1347+
assert result.equals(expected)
1348+
1349+
1350+
def test_replace_with_mask_null_in_mask():
1351+
"""Test null handling in mask."""
1352+
arr = pa.array([1, 2, 3, 4, 5, 6])
1353+
mask = pa.array([False, False, None, None, True, True])
1354+
replacements = pa.array([10, None])
1355+
expected = pa.array([1, 2, None, None, 10, None])
1356+
result = pc.replace_with_mask(arr, mask, replacements)
1357+
assert result.equals(expected)
1358+
1359+
1360+
def test_replace_with_mask_string_type():
1361+
"""Test replacement with string type."""
1362+
arr = pa.array(['a', 'b', 'c', 'd'])
1363+
mask = pa.array([True, False, True, False])
1364+
replacements = pa.array(['x', 'y'])
1365+
expected = pa.array(['x', 'b', 'y', 'd'])
1366+
result = pc.replace_with_mask(arr, mask, replacements)
1367+
assert result.equals(expected)
1368+
1369+
1370+
def test_replace_with_mask_float_type():
1371+
"""Test replacement with float type."""
1372+
arr = pa.array([1.1, 2.2, 3.3, 4.4])
1373+
mask = pa.array([True, False, True, False])
1374+
replacements = pa.array([10.5, 20.5])
1375+
expected = pa.array([10.5, 2.2, 20.5, 4.4])
1376+
result = pc.replace_with_mask(arr, mask, replacements)
1377+
assert result.equals(expected)
1378+
1379+
1380+
def test_replace_with_mask_chunked_array_multiple_chunks():
1381+
"""Test replace_with_mask with ChunkedArray with multiple chunks."""
1382+
arr = pa.chunked_array([[1, 2, 3], [4, 5, 6]])
1383+
mask = pa.array([True, False, False, False, True, False])
1384+
replacements = pa.array([10, 20])
1385+
expected = pa.chunked_array([[10, 2, 3], [4, 20, 6]])
1386+
result = pc.replace_with_mask(arr, mask, replacements)
1387+
assert result.equals(expected)
1388+
1389+
1390+
def test_replace_with_mask_chunked_array_empty_chunks():
1391+
"""Test replace_with_mask with ChunkedArray with empty chunks."""
1392+
arr = pa.chunked_array([[1, 2], [], [3, 4]])
1393+
mask = pa.array([True, False, True, False])
1394+
replacements = pa.array([10, 20])
1395+
expected = pa.chunked_array([[10, 2], [20, 4]])
1396+
result = pc.replace_with_mask(arr, mask, replacements)
1397+
assert result.equals(expected)
1398+
1399+
1400+
def test_replace_with_mask_error_replacement_count_mismatch():
1401+
"""Replacement count does not match true values in mask."""
1402+
arr = pa.array([1, 2, 3])
1403+
mask = pa.array([True, True, False])
1404+
replacements = pa.array([10])
1405+
with pytest.raises(pa.ArrowInvalid, match="expected 2.*but got 1"):
1406+
pc.replace_with_mask(arr, mask, replacements)
1407+
1408+
1409+
def test_replace_with_mask_error_mask_length_mismatch():
1410+
"""Mask length does not match input array length."""
1411+
arr = pa.array([1, 2, 3])
1412+
mask = pa.array([True, False])
1413+
replacements = pa.array([10])
1414+
with pytest.raises(pa.ArrowInvalid):
1415+
pc.replace_with_mask(arr, mask, replacements)
1416+
1417+
13001418
def test_binary_join():
13011419
ar_list = pa.array([['foo', 'bar'], None, []])
13021420
expected = pa.array(['foo-bar', None, ''])

0 commit comments

Comments
 (0)