sortingoparego

Sorting objects in Rego, based on multiple attributes


I'm trying to write some code in OPA Rego to sort a collection of objects based on the value of multiple attributes. I start with the objects in a dictionary as below and I'd like to sort the objects inside and return their relative ids. So having the dictionary

dict = { 
    "1": {"name": "ccc", "foo": "4"},
    "2": {"name": "aaa", "foo": "1"},
    "3": {"name": "bbb", "foo": "6"},
    "4": {"name": "eee", "foo": "5"},
    "5": {"name": "eee", "foo": "2"},
    "6": {"name": "ddd", "foo": "3"} 
}

sorting first by name and then by foo I would expect to return [ 2 3 1 6 5 4] Notice that for ids 4 and 5 the objects have the same name, so the order should be decided by foo

My first attempt is

_sorted = res{ 
    orderBy = ["name", "foo"]
    sorted1 = sort([ x | x := dict[_][orderBy[0]] ])
    res = [id | sorted1[_] == dict[id][orderBy[0]] ]
}

but this approach has problems when there are objects with the same name, therefore the following result "_sorted": ["2","3","1","6",**"4","5","4","5"**]

The second attempt was using sets instead, which solves the duplication issue

_sorted = res{ 
    orderBy = ["name", "foo"]
    sorted1 = { x | x := dict[_][orderBy[0]] }
    res = [id | sorted1[_] == dict[id][orderBy[0]] ]
}

but I don't know how to make it work with sorting on 2 attributes - last attempt

_sorted = res{ 
    orderBy = ["name", "foo"]
    sorted1 = { x | x := dict[_][orderBy[0]] }
    sorted2 = { x | x := dict[_][orderBy[1]] }
    res = [id | sorted1[_] == dict[id][orderBy[0]]; sorted2[_] == dict[sid][orderBy[1]] ]
}

Any suggestions are much appreciated :-) Thanks!


Solution

  • I think the key to sorting declaratively is to calculate the rank of each element, and then build an array accordingly.

    I've given your example a go here: playground

    rank(x, k, coll) := count(smaller) {
      smaller := { y | some y in coll; y[k] <= x[k] }
    }
    
    sorted(k, coll) := { r: sorted_("foo", es) |
        some r in numbers.range(1, count(coll))
        es := { e | some e in coll; r == rank(e, k, coll) }
    }
    
    # same as sorted, without the inner sorted_ call
    sorted_(k, coll) := { r: es |
        some r in numbers.range(1, count(coll))
        es := { e | some e in coll; r == rank(e, k, coll) }
    }
    
    result := [ x | x := sorted("name", dict)[_][_][_] ]
    

    It is not generic in the orderBy array -- name and foo are hardcoded there.

    I think it might be cool to add a new built-in function: sort_by(any<array<object>, set<object>) -> array<object>