4040
4141import pyarrow as pa
4242import pyarrow .compute as pc
43- from pyarrow .lib import ArrowNotImplementedError
43+ from pyarrow .lib import ArrowNotImplementedError , ArrowIndexError
4444
4545try :
4646 import pyarrow .substrait as pas
@@ -1596,19 +1596,14 @@ def test_inverse_permutation():
15961596 arr = pa .chunked_array ([
15971597 arr0 , [9 , 7 , 5 , 3 , 1 ], [0 ], [2 , 4 , 6 ], [8 ], arr0 ,
15981598 ])
1599- result = pc .inverse_permutation (arr )
1600- print (result )
16011599 expected = pa .chunked_array ([[5 , 4 , 6 , 3 , 7 , 2 , 8 , 1 , 9 , 0 ]], type = pa .int32 ())
1602- assert result .equals (expected )
1600+ assert pc . inverse_permutation ( arr ) .equals (expected )
16031601
1604- # `inverse_permutation` kernel currently does not accept options
1605- options = pc .InversePermutationOptions (max_index = 4 , output_type = pa .int64 ())
1606- print (options )
1607- with pytest .raises (TypeError , match = "an unexpected keyword argument \' options\' " ):
1608- pc .inverse_permutation (arr , options = options )
1602+ options = pc .InversePermutationOptions (max_index = 9 , output_type = pa .int32 ())
1603+ assert pc .inverse_permutation (arr , options = options ).equals (expected )
1604+ assert pc .inverse_permutation (arr , max_index = - 1 ).equals (expected )
16091605
1610- # `inverse_permutation` kernel currently won't accept max_index
1611- with pytest .raises (TypeError , match = "an unexpected keyword argument \' max_index\' " ):
1606+ with pytest .raises (ArrowIndexError , match = "Index out of bounds: 9" ):
16121607 pc .inverse_permutation (arr , max_index = 4 )
16131608
16141609
@@ -1620,14 +1615,12 @@ def test_scatter():
16201615 result = pc .scatter (values , indices )
16211616 assert result .equals (expected )
16221617
1623- # `scatter` kernel currently does not accept options
1624- options = pc .ScatterOptions (max_index = 4 )
1625- with pytest .raises (TypeError , match = "unexpected keyword argument \' options\' " ):
1626- pc .scatter (values , indices , options = options )
1618+ options = pc .ScatterOptions (max_index = - 1 )
1619+ assert pc .scatter (values , indices , options = options ).equals (expected )
1620+ assert pc .scatter (values , indices , max_index = 9 ).equals (expected )
16271621
1628- # `scatter` kernel currently won't accept max_index
1629- with pytest .raises (TypeError , match = "unexpected keyword argument \' max_index\' " ):
1630- pc .scatter (values , indices , max_index = 4 )
1622+ with pytest .raises (ArrowIndexError , match = "Index out of bounds: 9" ):
1623+ pc .scatter (values , indices , max_index = 4 )
16311624
16321625
16331626@pytest .mark .parametrize ("typ" , ["array" , "chunked_array" ])
0 commit comments