sorting an indexed option array in axis0 bug fix; argsort to accoun…

…t `None`s (#946)

* do not drop nones in when sorting an indexed option array in axis0

* remove unused parameter from sort_next method

* do the same for argsort

* add shifts for indexed option array argsorted indices

* fix spec python kernel function

* compact offsets, thanks @agoose77

* argsort and sort arrays containing `nan`s

* format

* missing header for std::nan

* try to fix bool comparison with nan

* replace Nones with their indices for argsort

* fix kernel test specs

* move for loops to kernel function

* disable the kernel test

include/awkward/Content.h

@@ -813,7 +813,7 @@ namespace awkward {
                   bool mask,
                   bool keepdims) const = 0;
 
-    /// @brief This array with one axis removed by applying a Reducer
+    /// @brief This array sorted
     ///
     /// The user's entry point for this operation is #sort.
     ///
@@ -824,7 +824,7 @@ namespace awkward {
     /// structure into this structure with the same meaning as in
     /// {@link ListArrayOf ListArray}.
     /// @param parents Groups to combine as an {@link IndexOf Index} of
-    /// upward pointers from this structure to the outer structure to reduce.
+    /// upward pointers from this structure to the outer structure to sort.
     /// @param outlength The length of the array, after the operation
     /// completes.
     /// @param ascending If `true`, the values will be sorted in an ascending
@@ -838,17 +838,38 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const = 0;
+                bool stable) const = 0;
 
+    /// @brief This array sorted indices
+    ///
+    /// The user's entry point for this operation is #argsort.
+    ///
+    /// @param negaxis The negative axis: `-axis`. That is, `negaxis = 1`
+    /// means the deepest axis level.
+    /// @param starts Staring positions of each group to combine as an
+    /// {@link IndexOf Index}. These are downward pointers from an outer
+    /// structure into this structure with the same meaning as in
+    /// {@link ListArrayOf ListArray}.
+    /// @param shifts Per-element adjustments that allows
+    /// for variable-length lists with axis != -1 and for missing values
+    /// (None).
+    /// @param parents Groups to combine as an {@link IndexOf Index} of
+    /// upward pointers from this structure to the outer structure to sort.
+    /// @param outlength The length of the array, after the operation
+    /// completes.
+    /// @param ascending If `true`, the values will be sorted in an ascending
+    /// order.
+    /// @param stable If `true`, the values will be sorted by a
+    /// stable sort algorithm to maintain the relative order of records with
+    /// equal keys (i.e. values).
     virtual const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const = 0;
+                   bool stable) const = 0;
 
     /// @brief A (possibly nested) array of integers indicating the
     /// positions of elements within each nested list.

include/awkward/array/BitMaskedArray.h

@@ -391,17 +391,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/ByteMaskedArray.h

@@ -357,17 +357,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/EmptyArray.h

@@ -266,17 +266,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/IndexedArray.h

@@ -445,17 +445,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/ListArray.h

@@ -379,17 +379,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/ListOffsetArray.h

@@ -367,17 +367,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/None.h

@@ -222,17 +222,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     /// @exception std::runtime_error is always thrown
     const ContentPtr

include/awkward/array/NumpyArray.h

@@ -471,20 +471,19 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       as_unique_strings(const Index64& offsets) const;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;
@@ -806,6 +805,7 @@ namespace awkward {
     const std::shared_ptr<void> index_sort(const T* data,
                                            int64_t length,
                                            const Index64& starts,
+                                           const Index64& shifts,
                                            const Index64& parents,
                                            int64_t outlength,
                                            bool ascending,

include/awkward/array/RawArray.h

@@ -1049,8 +1049,7 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override {
+                bool stable) const override {
       std::shared_ptr<T> ptr = kernel::malloc<T>(kernel::lib::cpu,   // DERIVE
                                                  length_*sizeof(T));
       Index64 offsets(2);
@@ -1088,11 +1087,11 @@ namespace awkward {
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override {
+                   bool stable) const override {
       std::shared_ptr<int64_t> ptr =
           kernel::malloc<int64_t>(kernel::lib::cpu,   // DERIVE
                                   length_*sizeof(int64_t));

include/awkward/array/Record.h

@@ -232,17 +232,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/RecordArray.h

@@ -361,17 +361,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/RegularArray.h

@@ -335,17 +335,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/UnionArray.h

@@ -381,17 +381,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/UnmaskedArray.h

@@ -315,17 +315,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
     const ContentPtr
       localindex(int64_t axis, int64_t depth) const override;

include/awkward/array/VirtualArray.h

@@ -332,17 +332,16 @@ namespace awkward {
                 const Index64& parents,
                 int64_t outlength,
                 bool ascending,
-                bool stable,
-                bool keepdims) const override;
+                bool stable) const override;
 
     const ContentPtr
       argsort_next(int64_t negaxis,
                    const Index64& starts,
+                   const Index64& shifts,
                    const Index64& parents,
                    int64_t outlength,
                    bool ascending,
-                   bool stable,
-                   bool keepdims) const override;
+                   bool stable) const override;
 
 
     const ContentPtr