def test_EmptyArray(): array = ak.layout.EmptyArray() assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] assert str(ak.type(ak.sort(array))) == "0 * float64" assert str(ak.type(ak.argsort(array))) == "0 * int64" array2 = ak.Array([[], [], []]) assert ak.to_list(ak.argsort(array2)) == [[], [], []] assert str(ak.type(ak.argsort(array2))) == "3 * var * int64"
def test_UnionArray(): content0 = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) content1 = ak.from_iter( [["one"], ["two"], ["three"], ["four"], ["five"]], highlevel=False ) tags = ak.layout.Index8(np.array([1, 1, 0, 0, 1, 0, 1, 1], dtype=np.int8)) index = ak.layout.Index32(np.array([0, 1, 0, 1, 2, 2, 4, 3], dtype=np.int32)) array = ak.layout.UnionArray8_32(tags, index, [content0, content1]) with pytest.raises(ValueError) as err: ak.sort(array, axis=1, ascending=True, stable=False) assert str(err.value).startswith("cannot sort UnionArray8_32")
def test_3d(): array = ak.layout.NumpyArray( np.array( [ # axis 2: 0 1 2 3 4 # axis 1: [ [1.1, 2.2, 3.3, 4.4, 5.5], # 0 [6.6, 7.7, 8.8, 9.9, 10.10], # 1 [11.11, 12.12, 13.13, 14.14, 15.15], ], # 2 [ [-1.1, -2.2, -3.3, -4.4, -5.5], # 3 [-6.6, -7.7, -8.8, -9.9, -10.1], # 4 [-11.11, -12.12, -13.13, -14.14, -15.15], ], ] ) ) # 5 assert ak.to_list( ak.argsort(array, axis=2, ascending=True, stable=False) ) == ak.to_list(np.argsort(array, 2)) assert ak.to_list( ak.sort(array, axis=2, ascending=True, stable=False) ) == ak.to_list(np.sort(array, 2)) assert ak.to_list( ak.argsort(array, axis=1, ascending=True, stable=False) ) == ak.to_list(np.argsort(array, 1)) assert ak.to_list( ak.sort(array, axis=1, ascending=True, stable=False) ) == ak.to_list(np.sort(array, 1)) assert ak.to_list(ak.sort(array, axis=1, ascending=False, stable=False)) == [ [ [11.11, 12.12, 13.13, 14.14, 15.15], [6.6, 7.7, 8.8, 9.9, 10.1], [1.1, 2.2, 3.3, 4.4, 5.5], ], [ [-1.1, -2.2, -3.3, -4.4, -5.5], [-6.6, -7.7, -8.8, -9.9, -10.1], [-11.11, -12.12, -13.13, -14.14, -15.15], ], ] assert ak.to_list( ak.sort(array, axis=0, ascending=True, stable=False) ) == ak.to_list(np.sort(array, 0)) assert ak.to_list( ak.argsort(array, axis=0, ascending=True, stable=False) ) == ak.to_list(np.argsort(array, 0))
def test_ByteMaskedArray(): content = ak.from_iter( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False) mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8)) array = ak.layout.ByteMaskedArray(mask, content, valid_when=False) assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=False)) == [ [0, 0, 0], [], [2, 2, 2, 2], None, None, ] assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [ [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9], None, None, ] assert ak.to_list(ak.sort(array, axis=0, ascending=False, stable=False)) == [ [6.6, 7.7, 8.8], [], [0.0, 1.1, 2.2, 9.9], None, None, ] assert ak.to_list(ak.argsort(array, axis=1, ascending=True, stable=False)) == [ [0, 1, 2], [], None, None, [0, 1, 2, 3], ] assert ak.to_list(array.sort(1, False, False)) == [ [2.2, 1.1, 0.0], [], None, None, [9.9, 8.8, 7.7, 6.6], ]
def test_sort_bytestrings(): array = ak.from_iter( [b"one", b"two", b"three", b"two", b"two", b"one", b"three"], highlevel=False ) assert ak.to_list(array) == [ b"one", b"two", b"three", b"two", b"two", b"one", b"three", ] assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [ b"one", b"one", b"three", b"three", b"two", b"two", b"two", ] assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=True)) == [ 0, 5, 2, 6, 1, 3, 4, ]
def test_bool_sort(): array = ak.layout.NumpyArray(np.array([True, False, True, False, False])) assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [ False, False, False, True, True, ]
def test_sort_strings(): content1 = ak.from_iter(["one", "two", "three", "four", "five"], highlevel=False) assert ak.to_list(content1) == ["one", "two", "three", "four", "five"] assert ak.to_list(ak.sort(content1, axis=0, ascending=True, stable=False)) == [ "five", "four", "one", "three", "two", ] assert ak.to_list(ak.sort(content1, axis=0, ascending=False, stable=False)) == [ "two", "three", "one", "four", "five", ]
def test_keep_None_in_place_test(): array = ak.Array([[3, 2, 1], [], None, [4, 5]]) assert ak.to_list(ak.argsort(array, axis=1)) == [ [2, 1, 0], [], None, [0, 1], ] assert ak.to_list(ak.sort(array, axis=1)) == [ [1, 2, 3], [], None, [4, 5], ] assert ak.to_list(array[ak.argsort(array, axis=1)]) == ak.to_list( ak.sort(array, axis=1))
def test_sort(): assert ak.sort(ak.from_iter([[2 + 4j, 1 + 5j], [], [3 + 3j]])).tolist() == [ [1 + 5j, 2 + 4j], [], [3 + 3j], ] assert ak.argsort(ak.from_iter([[2 + 4j, 1 + 5j], [], [3 + 3j]])).tolist() == [ [1, 0], [], [0], ]
def test_nan(): array = ak.Array([1, 2, np.nan, 3, 0, np.nan]) assert ak.argsort(array).tolist() == [ 2, 5, 4, 0, 1, 3, ] # Note, `nan` comparison with `nan` returns False assert str(ak.sort(array).tolist()) == "[nan, nan, 0.0, 1.0, 2.0, 3.0]"
def test(): array = ak.Array([[0, 1, 2, 3], [3, 3, 3, 2, 1]]) is_valid = array != 3 assert array.mask[is_valid].tolist() == [[0, 1, 2, None], [None, None, None, 2, 1]] assert ak.sort(array.mask[is_valid]).tolist() == [ [0, 1, 2, None], [1, 2, None, None, None], ] assert ak.argsort(array.mask[is_valid]).tolist() == [ [0, 1, 2, 3], [4, 3, 0, 1, 2], ]
def test_sort(): array = ak.Array(["one", "two", "three", "four", "five", "six", "seven", "eight"]) assert ak.sort(array, axis=-1).tolist() == [ "eight", "five", "four", "one", "seven", "six", "three", "two", ] array = ak.Array( [["twotwo", "two", "three"], ["four", "five"], [], ["six", "seven", "eight"]] ) assert ak.sort(array, axis=-1).tolist() == [ ["three", "two", "twotwo"], ["five", "four"], [], ["eight", "seven", "six"], ] array = ak.Array( [ [["twotwo", "two"], ["three"]], [["four", "five"]], [], [["six"], ["seven", "eight"]], ] ) assert ak.sort(array, axis=-1).tolist() == [ [["two", "twotwo"], ["three"]], [["five", "four"]], [], [["six"], ["eight", "seven"]], ]
def test_NumpyArray(): array = ak.layout.NumpyArray(np.array([3.3, 2.2, 1.1, 5.5, 4.4])) assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=False)) == [ 2, 1, 0, 4, 3, ] assert ak.to_list(ak.argsort(array, axis=0, ascending=False, stable=False)) == [ 3, 4, 0, 1, 2, ] assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [ 1.1, 2.2, 3.3, 4.4, 5.5, ] assert ak.to_list(ak.sort(array, axis=0, ascending=False, stable=False)) == [ 5.5, 4.4, 3.3, 2.2, 1.1, ] array2 = ak.layout.NumpyArray(np.array([[3.3, 2.2, 4.4], [1.1, 5.5, 3.3]])) assert ak.to_list(ak.sort(array2, axis=1, ascending=True, stable=False)) == ak.to_list( np.sort(array2, axis=1)) assert ak.to_list(ak.sort(array2, axis=0, ascending=True, stable=False)) == ak.to_list( np.sort(array2, axis=0)) assert ak.to_list(ak.argsort(array2, axis=1, ascending=True, stable=False)) == ak.to_list( np.argsort(array2, 1)) assert ak.to_list(ak.argsort(array2, axis=0, ascending=True, stable=False)) == ak.to_list( np.argsort(array2, 0)) with pytest.raises(ValueError) as err: ak.sort(array2, axis=2, ascending=True, stable=False) assert str(err.value).startswith( "axis=2 exceeds the depth of the nested list structure (which is 2)")
def test_sort_zero_length_arrays(): array = ak.layout.IndexedArray64(ak.layout.Index64([]), ak.layout.NumpyArray([1, 2, 3])) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] content0 = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) content1 = ak.from_iter(["one", "two", "three", "four", "five"], highlevel=False) tags = ak.layout.Index8([]) index = ak.layout.Index32([]) array = ak.layout.UnionArray8_32(tags, index, [content0, content1]) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] content = ak.from_iter( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False) mask = ak.layout.Index8([]) array = ak.layout.ByteMaskedArray(mask, content, valid_when=False) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] array = ak.layout.NumpyArray([]) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] array = ak.layout.RecordArray([]) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] content = ak.layout.NumpyArray( np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) starts1 = ak.layout.Index64([]) stops1 = ak.layout.Index64([]) offsets1 = ak.layout.Index64(np.array([0])) array = ak.layout.ListArray64(starts1, stops1, content) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == [] array = ak.layout.ListOffsetArray64(offsets1, content) assert ak.to_list(array) == [] assert ak.to_list(ak.sort(array)) == [] assert ak.to_list(ak.argsort(array)) == []
def test_bool(): array = ak.Array([True, False, False, True, True, True]) assert ak.argsort(array).tolist() == [1, 2, 0, 3, 4, 5] assert ak.sort(array).tolist() == [False, False, True, True, True, True]
def test_but_first_fix_sort(): assert ak.is_valid(ak.sort(ak.Array(["one", "two", "three"]), axis=-1))
def test(): data = ak.Array([[7, 5, 7], [], [2], [8, 2]]) assert ak.to_list(ak.sort(data)) == [[5, 7, 7], [], [2], [2, 8]] index = ak.argsort(data) assert ak.to_list(data[index]) == [[5, 7, 7], [], [2], [2, 8]]
def test_IndexedOptionArray(): array = ak.Array( [ [None, None, 2.2, 1.1, 3.3], [None, None, None], [4.4, None, 5.5], [5.5, None, None], [-4.4, -5.5, -6.6], ] ) assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [ [-4.4, -5.5, -6.6, 1.1, 3.3], [4.4, None, 2.2], [5.5, None, 5.5], [None, None, None], [None, None, None], ] assert ak.to_list(ak.sort(array, axis=1, ascending=True, stable=False)) == [ [1.1, 2.2, 3.3, None, None], [None, None, None], [4.4, 5.5, None], [5.5, None, None], [-6.6, -5.5, -4.4], ] assert ak.to_list(ak.sort(array, axis=1, ascending=False, stable=True)) == [ [3.3, 2.2, 1.1, None, None], [None, None, None], [5.5, 4.4, None], [5.5, None, None], [-4.4, -5.5, -6.6], ] assert ak.to_list(ak.sort(array, axis=1, ascending=False, stable=False)) == [ [3.3, 2.2, 1.1, None, None], [None, None, None], [5.5, 4.4, None], [5.5, None, None], [-4.4, -5.5, -6.6], ] assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=True)) == [ [4, 4, 4, 0, 0], [2, 0, 0], [3, 1, 2], [0, 2, 1], [1, 3, 3], ] assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=False)) == [ [4, 4, 4, 0, 0], [2, 0, 0], [3, 1, 2], [0, 2, 1], [1, 3, 3], ] assert ak.to_list(ak.argsort(array, axis=0, ascending=False, stable=True)) == [ [3, 4, 2, 0, 0], [2, 0, 0], [4, 1, 4], [0, 2, 1], [1, 3, 3], ] assert ak.to_list(ak.argsort(array, axis=0, ascending=False, stable=False)) == [ [3, 4, 2, 0, 0], [2, 0, 0], [4, 1, 4], [0, 2, 1], [1, 3, 3], ] assert ak.to_list(ak.argsort(array, axis=1, ascending=True, stable=True)) == [ [3, 2, 4, 0, 1], [0, 1, 2], [0, 2, 1], [0, 1, 2], [2, 1, 0], ] assert ak.to_list(ak.argsort(array, axis=1, ascending=True, stable=False)) == [ [3, 2, 4, 0, 1], [0, 1, 2], [0, 2, 1], [0, 1, 2], [2, 1, 0], ] assert ak.to_list(ak.argsort(array, axis=1, ascending=False, stable=True)) == [ [4, 2, 3, 0, 1], [0, 1, 2], [2, 0, 1], [0, 1, 2], [0, 1, 2], ] array2 = ak.Array([None, None, 1, -1, 30]) assert ak.to_list(ak.argsort(array2, axis=0, ascending=True, stable=True)) == [ 3, 2, 4, 0, 1, ] array3 = ak.Array( [[2.2, 1.1, 3.3], [], [4.4, 5.5], [5.5], [-4.4, -5.5, -6.6]] ).layout assert ak.to_list(ak.sort(array3, axis=1, ascending=False, stable=False)) == [ [3.3, 2.2, 1.1], [], [5.5, 4.4], [5.5], [-4.4, -5.5, -6.6], ] assert ak.to_list(ak.sort(array3, axis=0, ascending=True, stable=False)) == [ [-4.4, -5.5, -6.6], [], [2.2, 1.1], [4.4], [5.5, 5.5, 3.3], ] # FIXME: Based on NumPy list sorting: # # array([list([2.2, 1.1, 3.3]), list([]), list([4.4, 5.5]), list([5.5]), # list([-4.4, -5.5, -6.6])], dtype=object) # np.sort(array, axis=0) # array([list([]), list([-4.4, -5.5, -6.6]), list([2.2, 1.1, 3.3]), # list([4.4, 5.5]), list([5.5])], dtype=object) # # the result should be: # # [[ -4.4, -5.5, -6.6 ], # [ 2.2, 1.1, 3.3 ], # [ 4.4, 5.5 ], # [ 5.5 ], # []] # This can be done following the steps: pad, sort, # and dropna to strip off the None's # array4 = array3.rpad(3, 1) assert ak.to_list(array4) == [ [2.2, 1.1, 3.3], [None, None, None], [4.4, 5.5, None], [5.5, None, None], [-4.4, -5.5, -6.6], ] array5 = ak.sort(array4, axis=0, ascending=True, stable=False) assert ak.to_list(array5) == [ [-4.4, -5.5, -6.6], [2.2, 1.1, 3.3], [4.4, 5.5, None], [5.5, None, None], [None, None, None], ] array4 = array3.rpad(5, 1) assert ak.to_list(array4) == [ [2.2, 1.1, 3.3, None, None], [None, None, None, None, None], [4.4, 5.5, None, None, None], [5.5, None, None, None, None], [-4.4, -5.5, -6.6, None, None], ] array5 = ak.sort(array4, axis=0, ascending=True, stable=False) assert ak.to_list(array5) == [ [-4.4, -5.5, -6.6, None, None], [2.2, 1.1, 3.3, None, None], [4.4, 5.5, None, None, None], [5.5, None, None, None, None], [None, None, None, None, None], ] array5 = ak.argsort(array4, axis=0, ascending=True, stable=False) assert ak.to_list(array5) == [ [4, 4, 4, 0, 0], [0, 0, 0, 1, 1], [2, 2, 1, 2, 2], [3, 1, 2, 3, 3], [1, 3, 3, 4, 4], ] # FIXME: implement dropna to strip off the None's # # array6 = array5.dropna(0) # assert ak.to_list(array6) == [ # [ -4.4, -5.5, -6.6 ], # [ 2.2, 1.1, 3.3 ], # [ 4.4, 5.5 ], # [ 5.5 ], # []] content = ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])) index1 = ak.layout.Index32(np.array([1, 2, 3, 4], dtype=np.int32)) indexedarray1 = ak.layout.IndexedArray32(index1, content) assert ak.to_list( ak.argsort(indexedarray1, axis=0, ascending=True, stable=False) ) == [0, 1, 2, 3] index2 = ak.layout.Index64(np.array([1, 2, 3], dtype=np.int64)) indexedarray2 = ak.layout.IndexedArray64(index2, indexedarray1) assert ak.to_list( ak.sort(indexedarray2, axis=0, ascending=False, stable=False) ) == [5.5, 4.4, 3.3] index3 = ak.layout.Index32(np.array([1, 2], dtype=np.int32)) indexedarray3 = ak.layout.IndexedArray32(index3, indexedarray2) assert ak.to_list(ak.sort(indexedarray3, axis=0, ascending=True, stable=False)) == [ 4.4, 5.5, ]