5
$\begingroup$

I am aware that the question of "how do I sort on a column" has been asked in various forms, most notably this one. While some powerful solutions have been presented, as far as I can find no one solution brings together all the elements of what most users would expect from a multi-column sort, similar to what one would find in Excel or SQL. The sought solution would:

  • Work equally well with matrices, lists of associations, or Dataset objects
  • Allow one to choose multiple existing or calculated columns, in a set precedence order
  • Allow one to independently set the comparison function for each column, at minimum ascending vs. descending

The fact that this question has been asked (and answered) in so many forms is a collective expression of surprise that this seemingly simple action should not have a more elegant solution built in. This should not be so hard in a language that can give us so much power, usually in a single line of code.

I will put my own solution below.

$\endgroup$
3
$\begingroup$

Mathematica provides the sorting function Sort and SortBy. Sort has a second argument that allows us to pass a function that compares two elements. SortBy has a second argument that lets us calculate a value (or calculate a list of values) to sort by, but does not let us provide the comparison function. There does not appear to be a straightforward way to combine these two.

Borrowing the example data from this post,

data = {{"a", 1, 1}, {"a", 1, 2}, {"a", 1, 3}, {"c", 2, 1}, {"b", 2, 2}, 
{"b", 2, 3}, {"c", 3, 1}, {"a", 3, 2}, {"a", 3, 3}};
data[[All, 2]] = data[[All, 2]] /. {1 -> "q", 2 -> "r", 3 -> "s"};
assoc = AssociationThread[{"first", "second", "third"} -> #] & /@ data

(* {<|"first" -> "a", "second" -> "q", "third" -> 1|>, <|"first" -> "a", 
"second" -> "q", "third" -> 2|>, <|"first" -> "a", "second" -> "q", 
"third" -> 3|>, <|"first" -> "c", "second" -> "r", 
"third" -> 1|>, <|"first" -> "b", "second" -> "r", 
"third" -> 2|>, <|"first" -> "b", "second" -> "r", 
"third" -> 3|>, <|"first" -> "c", "second" -> "s", 
"third" -> 1|>, <|"first" -> "a", "second" -> "s", 
"third" -> 2|>, <|"first" -> "a", "second" -> "s", "third" -> 3|>} *)

The second argument to SortBy can be used to select columns

SortBy[{#first &, #third &}][assoc]

(* {<|"first" -> "a", "second" -> "q", "third" -> 1|>, 
<|"first" -> "a", "second" -> "q", "third" -> 2|>, 
<|"first" -> "a", "second" -> "s", "third" -> 2|>,
<|"first" -> "a", "second" -> "q", "third" -> 3|>,
<|"first" -> "a", "second" -> "s", "third" -> 3|>, 
<|"first" -> "b", "second" -> "r", "third" -> 2|>, 
<|"first" -> "b", "second" -> "r", "third" -> 3|>, 
<|"first" -> "c", "second" -> "r", "third" -> 1|>, 
<|"first" -> "c", "second" -> "s", "third" -> 1|>} *)

but no way to choose the function by which these are ordered column by column.

The approach I use here is to use GroupBy to create a tree of associations with common values in each column. Then I do a depth-first traversal, sorting by the keys and unpacking the results until the table is sorted. First, I have to make a new version of KeySort that lets me set the comparison function as in Sort:

keySort[assoc_Association, p_: (OrderedQ[{#1, #2}] &)] := 
 Query[Ordering[Keys[assoc], All, p]]@assoc

The function orderby groups the data by column values and performs the depth-first traversal:

orderby[dat_, f_, p_] := Module[{len = Length[p]},
  Fold[
   Function[{aa, pp}, 
    Map[Flatten[Values[keySort[#, pp]], 1] &, aa, {--len}]],
   GroupBy[dat, f],
   Reverse[p]
   ]
  ]

The rest of the code provides a more Mathematica-idosyncratic interface. The complete code is here:

Ascending::usage = 
  "Ascending is the default ordering function for OrderBy, sorting \
the corresponding column into canonical order.";
Descending::usage = 
  "Descending is an ordering function for OrderBy, sorting the \
corresponding column into reverse canonical order.";

keySort::usage = 
  "keySort[assoc, p] orders the elements of an association by sorting \
its keys using the ordering function p.";

OrderBy::usage = 
  "OrderBy[f][list] sorts the elements of list in the order defined \
by applying f to each of them.

  OrderBy[f, p][list] sorts using the ordering function p.

  OrderBy[{f1, f2, ...}, {p1, p2, ...}][list] sorts using the \
ordering function pi for the corresponding fi.";

Ascending = OrderedQ[{#1, #2}] &;
Descending = OrderedQ[{#2, #1}] &;

keySort[assoc_Association, p_: (OrderedQ[{#1, #2}] &)] := 
 Query[Ordering[Keys[assoc], All, p]]@assoc

orderby[dat_, f_, p_] := Module[{len = Length[p]},
  Fold[
   Function[{aa, pp}, 
    Map[Flatten[Values[keySort[#, pp]], 1] &, aa, {--len}]],
   GroupBy[dat, f],
   Reverse[p]
   ]
  ]

OrderBy[fi_, pi_: Automatic][dat_] := Module[{f, p},
  f = If[MatchQ[fi, _List], fi, {fi}];
  p = If[MatchQ[pi, _List],
     PadRight[pi, Length[f], Automatic],
     ConstantArray[pi, Length[f]]
     ] /. Automatic -> Ascending;
  orderby[dat, f, p]
  ]

OrderBy is an operator. For example, to sort the example association in descending order by the first column then ascending order in the third:

OrderBy[{#first &, #third &}, {Descending, Ascending}]@assoc

(* {<|"first" -> "c", "second" -> "r", "third" -> 1|>, 
<|"first" -> "c", "second" -> "s", "third" -> 1|>, 
<|"first" -> "b", "second" -> "r", "third" -> 2|>, 
<|"first" -> "b", "second" -> "r", "third" -> 3|>, 
<|"first" -> "a", "second" -> "q", "third" -> 1|>, 
<|"first" -> "a", "second" -> "q", "third" -> 2|>, 
<|"first" -> "a", "second" -> "s", "third" -> 2|>, 
<|"first" -> "a", "second" -> "q", "third" -> 3|>, 
<|"first" -> "a", "second" -> "s", "third" -> 3|>} *)

The solution also works with matrix data. The example below sorts the matrix data ascending by column 2 then descending by column 1, showing two different ways to choose the sorting columns (using Extract or Query)

OrderBy[{Extract[2], Query[1]}, {Ascending, Descending}]@data

(* {{"a", "q", 1}, {"a", "q", 2}, {"a", "q", 3}, {"c", "r", 1}, {"b", 
  "r", 2}, {"b", "r", 3}, {"c", "s", 1}, {"a", "s", 2}, {"a", "s", 3}} *)

Sadly, OrderBy does not always work with Dataset. I am not sure why at this point. If I figure this out I will update the post.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.