Open Policy Agent - Improve performance of a grouping comprehension

Question

I have a id and role mapping in below format

{
    "ra": [
      {
        "id": 168,
        "code": "AFAP"
      },
      {
        "id": 180,
        "code": "ABC"
      },
      {
        "id": 180,
        "code": "ABCMND"
      }
   ]
}

I need the output to be like below

{
    "roleactions": {
        "168": [
            "AFAP"
        ],
        "180": [
            "ABC",
            "ABCMND",
            "DCRMP"
        ]
    }
}

So i wrote the below code

roleactions = r_map {
    r := data.ra
    r_map := {id: list |
        some i
        id := r[i].id
        list := [obj |
            some j
            r[j].id == id
            obj := r[j].code
        ]
    }
}

But when I run this for it almost takes 5-6 seconds

Found 1 result in 5682526.465 µs.

Can someone guide on how to make write this policy map to improve performance?

Answer 1

OPA can evaluate comprehensions like this in linear-time: https://www.openpolicyagent.org/docs/latest/policy-performance/#comprehension-indexing. The problem in this case is that the local variable r is not safe when considering the comprehensions in isolation.

If you refactor the comprehension like below, the runtime should be linear:

roleactions := r_map {
    r_map := {id: list |
        some i
        id := data.ra[i].id
        list := [obj |
            some j
            data.ra[j].id == id
            obj := data.ra[j].code
        ]
    }
}