@@ -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+
13001418def test_binary_join ():
13011419 ar_list = pa .array ([['foo' , 'bar' ], None , []])
13021420 expected = pa .array (['foo-bar' , None , '' ])
0 commit comments