IndexError for ak.argsort on arrays with None

[kmohrman]

Version of Awkward Array

2.4.6

Description and code to reproduce

It seems that ak.argsort does not work on arrays that are entirely None, though ak.sort works fine. I would have naively expected that if sort works on an array, argsort would also work on the array, so it is not clear to me if this behavior would be expected or not. Here is an example that reproduces the behavior:

import awkward as ak
  
a = ak.Array([None,None,[None,None]])
print(ak.sort(a,axis=1))
print(ak.argsort(a,axis=1))

The ak.sort line is ok (and prints [None, None, [None, None]]), while the ak.argsort line fails with the error below:

Traceback (most recent call last):
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 301, in _getitem_at
    out = self._data[where]
          ~~~~~~~~~~^^^^^^^
IndexError: index 2 is out of bounds for axis 0 with size 2

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/k.mohrman/coffea_dir/ewkcoffea/analysis/wwz/mre2.py", line 6, in <module>
    print(ak.argsort(a,axis=1))
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/highlevel.py", line 1233, in __str__
    return awkward._prettyprint.valuestr(self, 1, 80)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 240, in valuestr
    _, strs = valuestr_horiz(data, limit_cols)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 135, in valuestr_horiz
    cols_taken, strs = valuestr_horiz(current, limit_cols - 2)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 117, in valuestr_horiz
    current = get_at(data, index)
              ^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 43, in get_at
    out = data._layout._getitem_at(index)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/indexedoptionarray.py", line 308, in _getitem_at
    return self._content._getitem_at(self._index[where])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 303, in _getitem_at
    raise ak._errors.index_error(self, where, str(err)) from err
IndexError: cannot slice NumpyArray (of length 2) with 2: index 2 is out of bounds for axis 0 with size 2

[agoose77]

From the error, it looks like ak.argsort is actually succeeding; it's the pretty-printing that's failing. Let me take a look at that!

[kmohrman]

Thank you! Just to mention one more thing, it seems that an error also occurs when trying to use the result of argsort to actually sort an array, e.g. in this code:

a = ak.Array([None,None,[None,None]])
b = ak.argsort(a,axis=1)
c = a[b]

There is an error (pasted below) on the c = a[b] line (though it runs fine if a is not all None, e.g. if its [None,None,[None,1.1]]). I'm wondering if the #2769 will also fix this issue?

Traceback (most recent call last):
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 333, in _carry
    nextdata = self._data[carry.data]
               ~~~~~~~~~~^^^^^^^^^^^^
IndexError: index 2 is out of bounds for axis 0 with size 2

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/highlevel.py", line 1014, in __getitem__
    prepare_layout(self._layout[where]), self._behavior, allow_other=True
                   ~~~~~~~~~~~~^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/content.py", line 519, in __getitem__
    return self._getitem(where)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/content.py", line 572, in _getitem
    return self._getitem(where.layout)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/content.py", line 647, in _getitem
    return self._getitem((where,))
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/content.py", line 553, in _getitem
    items = ak._slicing.normalise_items(where, backend)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 320, in normalise_items
    return [normalise_item(x, backend=backend) for x in where]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 269, in normalise_item
    out = _normalise_item_bool_to_int(_normalise_item_nested(item), backend)
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 428, in _normalise_item_nested
    _normalise_item_nested(projected),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 380, in _normalise_item_nested
    return _normalise_item_nested(next)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 367, in _normalise_item_nested
    _normalise_item_nested(item.content),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_slicing.py", line 419, in _normalise_item_nested
    projected = item.content._carry(ak.index.Index64(nextindex[nonnull]), False)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 335, in _carry
    raise ak._errors.index_error(self, carry.data, str(err)) from err
IndexError: cannot slice NumpyArray (of length 2) with array([1, 2]): index 2 is out of bounds for axis 0 with size 2

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 301, in _getitem_at
    out = self._data[where]
          ~~~~~~~~~~^^^^^^^
IndexError: index 2 is out of bounds for axis 0 with size 2

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/k.mohrman/coffea_dir/ewkcoffea/analysis/wwz/mre2.py", line 8, in <module>
    c = a[b]
        ~^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/highlevel.py", line 1012, in __getitem__
    with ak._errors.SlicingErrorContext(self, where):
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 67, in __exit__
    self.handle_exception(exception_type, exception_value)
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 80, in handle_exception
    self.decorate_exception(cls, exception)
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 91, in decorate_exception
    exception.add_note(self.note)
                       ^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 342, in note
    {self.where}"""
     ^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 327, in where
    out = self._kwargs["where"] = out()
                                  ^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 33, in __call__
    return self.func(*self.args, **self.kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_errors.py", line 380, in format_slice
    return repr(x)
           ^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/highlevel.py", line 1236, in __repr__
    return self._repr(80)
           ^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/highlevel.py", line 1262, in _repr
    valuestr = valuestr + " " + awkward._prettyprint.valuestr(self, 1, strwidth)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 240, in valuestr
    _, strs = valuestr_horiz(data, limit_cols)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 135, in valuestr_horiz
    cols_taken, strs = valuestr_horiz(current, limit_cols - 2)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 117, in valuestr_horiz
    current = get_at(data, index)
              ^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/_prettyprint.py", line 43, in get_at
    out = data._layout._getitem_at(index)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/indexedoptionarray.py", line 308, in _getitem_at
    return self._content._getitem_at(self._index[where])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/k.mohrman/miniconda3/envs/ak-env/lib/python3.12/site-packages/awkward/contents/numpyarray.py", line 303, in _getitem_at
    raise ak._errors.index_error(self, where, str(err)) from err
IndexError: cannot slice NumpyArray (of length 2) with 2: index 2 is out of bounds for axis 0 with size 2