Improve performance in queries that perform "group by" operations

[tsandall]

Rego doesn't support mutation and there's no explicit "group by" operation. As a result, when policies need to aggregate or group data by some key, they need to use a comprehension and close over values from inside the query. This leads to cross-products in Rego when implementations in other languages would have linear runtime. This can be seen going back all the way to the old scheduler policy. For example:

# this is from test/scheduler/scheduler_test.go
pods_on_node[node_id] = pds {
    node_name := nodes[node_id].metadata.name
    pds := [p | pods[i].spec.nodeName == node_name; p := pods[i]]
}

This rule has O(nodes*pods) runtime. Another common case is inverting an object (which ends up being O(object^2). For example:

v := obj[_]
ks := [k | some k; v==obj[k]]

The latter case is trivial and could be addressed by a built-in function. The scheduler case above is not as easy. E.g., the built-in would have to accept a parameter that indicates how to key any index it builds.

A better solution would be for OPA to identify comprehensions that could and should be indexed (as well as the keys to use for those indices.) During evaluation, the evaluator can lazily build the comprehension index that would allow the query to efficiently lookup the comprehension value based on local variable bindings.

For example, in the scheduler case above, the index would be keyed by "node_name". In the second example, the index would be keyed by "v". The index could be hierarchical to avoid needing to prune lookup results.

Only comprehensions that are safe when excluding the parent query can be indexed. For example:

x := p[_]
_ = {1 | y = x[_]}  # cannot be indexed because 'x' is not safe

Note, in some cases it'll be less efficient to index comprehensions. For example:

p[x]
ys := {y | q[x][y] = 1}

In this case if O(p)>>O(q) then comprehension indexing is effective. OTOH, if O(q)>>O(p) then indexing may be more expensive because all values of 'x' are not known at indexing time. A simple heuristic could be to avoid indexing when keys are used as ref operands because they are probably used to perform constant-time lookups that shrink the search space.

Misc. notes

• Comprehension index must be invalidated by with statements