add multi-dim sort

arkouda/sorting.py

@@ -251,6 +251,7 @@ def sort(pda: pdarray, algorithm: SortingAlgorithm = SortingAlgorithm.RadixSortL
     if pda.size == 0:
         return zeros(0, dtype=pda.dtype)
     repMsg = generic_msg(
-        cmd=f"sort<{pda.dtype.name},{pda.ndim}>", args={"alg": algorithm.name, "array": pda, "axis": axis}
+        cmd=f"sort<{pda.dtype.name},{pda.ndim}>",
+        args={"alg": algorithm.name, "array": pda, "axis": axis},
     )
     return create_pdarray(cast(str, repMsg))

src/SortMsg.chpl

@@ -16,7 +16,7 @@ module SortMsg
     use Logging;
     use Message;
     private use ArgSortMsg;
-    use NumPyDType;
+    use NumPyDType only whichDtype;
 
     private config const logLevel = ServerConfig.logLevel;
     private config const logChannel = ServerConfig.logChannel;
@@ -48,8 +48,53 @@ module SortMsg
       }
     }
 
+
+    proc sort(array: [?d] ?t, alg: string, axis: int): [d] t throws 
+    where ( (t == real) || (t==int) || (t==uint(64)) ) && ( d.rank > 1 ) {
+
+      var algorithm: SortingAlgorithm = ArgSortMsg.getSortingAlgorithm(alg);
+      const itemsize = dtypeSize(whichDtype(t));
+      overMemLimit(radixSortLSD_keys_memEst(d.size, itemsize));
+
+      const DD = domOffAxis(d, axis);
+      var sorted = makeDistArray((...d.shape), t);
+
+      if algorithm == SortingAlgorithm.TwoArrayRadixSort {
+        for idx in DD {
+          // make a copy of the array along the slice corresponding to idx
+          // TODO: create a twoArrayRadixSort that operates on a slice of the array
+          // in place instead of requiring the copy in/out
+          var slice = makeDistArray(d.dim(axis).size, t);
+          forall i in d.dim(axis) with (var perpIdx = idx) {
+            perpIdx[axis] = i;
+            slice[i] = array[perpIdx];
+          }
+
+          ArgSortMsg.dynamicTwoArrayRadixSort(slice, comparator=myDefaultComparator);
+
+          forall i in d.dim(axis) with (var perpIdx = idx) {
+            perpIdx[axis] = i;
+            sorted[perpIdx] = slice[i];
+          }
+        }
+      } else {
+        // TODO: make a version of radixSortLSD_keys that does the sort on
+        // slices of `e.a` directly instead of requiring a copy for each slice
+        for idx in DD {
+          const sliceDom = domOnAxis(d, idx, axis),
+                sliced1D = removeDegenRanks(array[sliceDom], 1),
+                sliceSorted = radixSortLSD_keys(sliced1D);
+
+          forall i in sliceDom do sorted[i] = sliceSorted[i[axis]];
+        }
+      }
+
+      return sorted;
+    }
+
+
     proc sort(array: [?d] ?t, alg: string, axis: int): [d] t throws 
-    where ( (t != real) && (t!=int) && (t!=uint(64)) ) || ( d.rank != 1 ) {
+    where ( (t != real) && (t!=int) && (t!=uint(64)) ) {
       throw new Error("Insightful error message.");
     }