opa icon indicating copy to clipboard operation
opa copied to clipboard
verified publisher icon open-policy-agent

Partial sets can be O(n²) when accessed iteratively

Open charlesdaniels opened this issue 2 years ago • 0 comments

TL;DR: if you are checking for membership in a partial set iteratively (in another partial rule, or in a comprehension), and it seems very slow, try copying the partial set into an intermediate variable to force it to be cached (e.g. my_partial_set_cached := my_partial_set, then use my_partial_set_cached during iteration).

In some circumstances, using partial sets can be substantially slower compared to equivalent set comprehensions when the partial sets are accessed iteratively. This situation is more likely to occur when the size of the partial set is small, but checking for set membership is expensive.

To provide a motivational example: I found this because I was filtering some JSON-lines style logs; I first scanned every object in the log to select a small collection of identifiers, then later wanted to filter the log down to only objects containing those identifiers.

With some assistance from @tsandall , I have constructed a minimal example that demonstrates this problem. Consider the following Rego code snippet:

# repro.rego
package repl

p[x] {
	x := 7
}

p[x] {
	x := 8
}

p[x] {
	x := 9
}

a := [1,2,3,7]

# this is slow
b := {x | x := a[_]; p[x]}

# this is fast
c := p
d := {x | x := a[_]; c[x]}

We can run this snippet with opa eval --explain full -f pretty --data ./repro.rego 'data.repl' , to see the following output:

query:1             Enter data.repl = _
query:1             | Eval data.repl = _
query:1             | Index data.repl.a (matched 1 rule, early exit)
./repro.rego:15     | Enter data.repl.a
./repro.rego:15     | | Eval true
./repro.rego:15     | | Exit data.repl.a early
./repro.rego:15     | Redo data.repl.a
./repro.rego:15     | | Redo true
query:1             | Index data.repl.b (matched 1 rule)
./repro.rego:18     | Enter data.repl.b
./repro.rego:18     | | Eval true
./repro.rego:18     | | Eval __local5__ = {x | x = data.repl.a[_]; data.repl.p[x]}
./repro.rego:18     | | Enter x = data.repl.a[_]; data.repl.p[x]
./repro.rego:18     | | | Eval x = data.repl.a[_]
./repro.rego:18     | | | Index data.repl.a (matched 1 rule, early exit)
./repro.rego:18     | | | Eval data.repl.p[x]
./repro.rego:18     | | | Index data.repl.p (matched 3 rules)
./repro.rego:3      | | | Enter data.repl.p
./repro.rego:4      | | | | Eval x = 7
./repro.rego:4      | | | | Fail x = 7
./repro.rego:7      | | | Enter data.repl.p
./repro.rego:8      | | | | Eval x = 8
./repro.rego:8      | | | | Fail x = 8
./repro.rego:11     | | | Enter data.repl.p
./repro.rego:12     | | | | Eval x = 9
./repro.rego:12     | | | | Fail x = 9
./repro.rego:18     | | | Fail data.repl.p[x]
./repro.rego:18     | | | Redo x = data.repl.a[_]
./repro.rego:18     | | | Eval data.repl.p[x]
./repro.rego:18     | | | Index data.repl.p (matched 3 rules)
./repro.rego:3      | | | Enter data.repl.p
./repro.rego:4      | | | | Eval x = 7
./repro.rego:4      | | | | Fail x = 7
./repro.rego:7      | | | Enter data.repl.p
./repro.rego:8      | | | | Eval x = 8
./repro.rego:8      | | | | Fail x = 8
./repro.rego:11     | | | Enter data.repl.p
./repro.rego:12     | | | | Eval x = 9
./repro.rego:12     | | | | Fail x = 9
./repro.rego:18     | | | Fail data.repl.p[x]
./repro.rego:18     | | | Redo x = data.repl.a[_]
./repro.rego:18     | | | Eval data.repl.p[x]
./repro.rego:18     | | | Index data.repl.p (matched 3 rules)
./repro.rego:3      | | | Enter data.repl.p
./repro.rego:4      | | | | Eval x = 7
./repro.rego:4      | | | | Fail x = 7
./repro.rego:7      | | | Enter data.repl.p
./repro.rego:8      | | | | Eval x = 8
./repro.rego:8      | | | | Fail x = 8
./repro.rego:11     | | | Enter data.repl.p
./repro.rego:12     | | | | Eval x = 9
./repro.rego:12     | | | | Fail x = 9
./repro.rego:18     | | | Fail data.repl.p[x]
./repro.rego:18     | | | Redo x = data.repl.a[_]
./repro.rego:18     | | | Eval data.repl.p[x]
./repro.rego:18     | | | Index data.repl.p (matched 3 rules)
./repro.rego:3      | | | Enter data.repl.p
./repro.rego:4      | | | | Eval x = 7
./repro.rego:3      | | | | Exit data.repl.p
./repro.rego:3      | | | Redo data.repl.p
./repro.rego:4      | | | | Redo x = 7
./repro.rego:7      | | | Enter data.repl.p
./repro.rego:8      | | | | Eval x = 8
./repro.rego:8      | | | | Fail x = 8
./repro.rego:11     | | | Enter data.repl.p
./repro.rego:12     | | | | Eval x = 9
./repro.rego:12     | | | | Fail x = 9
./repro.rego:18     | | | Exit x = data.repl.a[_]; data.repl.p[x]
./repro.rego:18     | | Redo x = data.repl.a[_]; data.repl.p[x]
./repro.rego:18     | | | Redo data.repl.p[x]
./repro.rego:18     | | | Redo x = data.repl.a[_]
./repro.rego:18     | | Exit data.repl.b
./repro.rego:18     | Redo data.repl.b
./repro.rego:18     | | Redo __local5__ = {x | x = data.repl.a[_]; data.repl.p[x]}
./repro.rego:18     | | Redo true
query:1             | Index data.repl.c (matched 1 rule)
./repro.rego:21     | Enter data.repl.c
./repro.rego:21     | | Eval true
./repro.rego:21     | | Eval __local6__ = data.repl.p
./repro.rego:21     | | Index data.repl.p (matched 3 rules)
./repro.rego:3      | | Enter data.repl.p
./repro.rego:4      | | | Eval x = 7
./repro.rego:3      | | | Exit data.repl.p
./repro.rego:3      | | Redo data.repl.p
./repro.rego:4      | | | Redo x = 7
./repro.rego:7      | | Enter data.repl.p
./repro.rego:8      | | | Eval x = 8
./repro.rego:7      | | | Exit data.repl.p
./repro.rego:7      | | Redo data.repl.p
./repro.rego:8      | | | Redo x = 8
./repro.rego:11     | | Enter data.repl.p
./repro.rego:12     | | | Eval x = 9
./repro.rego:11     | | | Exit data.repl.p
./repro.rego:11     | | Redo data.repl.p
./repro.rego:12     | | | Redo x = 9
./repro.rego:21     | | Exit data.repl.c
./repro.rego:21     | Redo data.repl.c
./repro.rego:21     | | Redo __local6__ = data.repl.p
./repro.rego:21     | | Redo true
query:1             | Index data.repl.d (matched 1 rule)
./repro.rego:22     | Enter data.repl.d
./repro.rego:22     | | Eval true
./repro.rego:22     | | Eval __local7__ = {x | x = data.repl.a[_]; data.repl.c[x]}
./repro.rego:22     | | Enter x = data.repl.a[_]; data.repl.c[x]
./repro.rego:22     | | | Eval x = data.repl.a[_]
./repro.rego:22     | | | Index data.repl.a (matched 1 rule, early exit)
./repro.rego:22     | | | Eval data.repl.c[x]
./repro.rego:22     | | | Index data.repl.c (matched 1 rule)
./repro.rego:22     | | | Fail data.repl.c[x]
./repro.rego:22     | | | Redo x = data.repl.a[_]
./repro.rego:22     | | | Eval data.repl.c[x]
./repro.rego:22     | | | Index data.repl.c (matched 1 rule)
./repro.rego:22     | | | Fail data.repl.c[x]
./repro.rego:22     | | | Redo x = data.repl.a[_]
./repro.rego:22     | | | Eval data.repl.c[x]
./repro.rego:22     | | | Index data.repl.c (matched 1 rule)
./repro.rego:22     | | | Fail data.repl.c[x]
./repro.rego:22     | | | Redo x = data.repl.a[_]
./repro.rego:22     | | | Eval data.repl.c[x]
./repro.rego:22     | | | Index data.repl.c (matched 1 rule)
./repro.rego:22     | | | Exit x = data.repl.a[_]; data.repl.c[x]
./repro.rego:22     | | Redo x = data.repl.a[_]; data.repl.c[x]
./repro.rego:22     | | | Redo data.repl.c[x]
./repro.rego:22     | | | Redo x = data.repl.a[_]
./repro.rego:22     | | Exit data.repl.d
./repro.rego:22     | Redo data.repl.d
./repro.rego:22     | | Redo __local7__ = {x | x = data.repl.a[_]; data.repl.c[x]}
./repro.rego:22     | | Redo true
query:1             | Index data.repl.p (matched 3 rules)
query:1             | Exit data.repl = _
query:1             Redo data.repl = _
query:1             | Redo data.repl = _
{
  "a": [
    1,
    2,
    3,
    7
  ],
  "b": [
    7
  ],
  "c": [
    7,
    8,
    9
  ],
  "d": [
    7
  ],
  "p": [
    7,
    8,
    9
  ]
}

In this example computation of the b object is O(n²), while computation of the d object is linear. As a rough intuition for why, consider x := a[_]; p[x], here we are iterating over the a array, and during each iteration, we check if x is a member of the partial set p. At present, the Rego compiler will compute a full copy of p to determine if x is a set member, then discard the computed copy of p. On the next iteration, it will repeat this process. In a situation where a is large and computing p is expensive, this can result in very slow policy execution. Computation of c is linear time on the number of set members, and the result is cached in c. Once we reach d, checking c[x] is O(1), because it is now a simple hash set lookup as c has been fully precomputed.

We can verify that this intuition is correct by considering the --explain output shown above. Notice that during computation of b, we repeatedly see the block:

./repro.rego:3      | | | Enter data.repl.p
./repro.rego:4      | | | | Eval x = 7
./repro.rego:4      | | | | Fail x = 7
./repro.rego:7      | | | Enter data.repl.p
./repro.rego:8      | | | | Eval x = 8
./repro.rego:8      | | | | Fail x = 8
./repro.rego:11     | | | Enter data.repl.p
./repro.rego:12     | | | | Eval x = 9
./repro.rego:12     | | | | Fail x = 9

This is the Rego runtime recomputing p for each loop iteration.

During computation of b, we instead see lines like Index data.repl.c (matched 1 rule), because c has already been fully expanded in memory.

Ideally, OPA would have additional optimizations to detect this situation and speed up policy execution. In the meantime, the workaround is to simply ensure that the set you need to check iteratively is precomputed, either by assigning the partial set to a temporary variable as shown in this example, or by writing an equivalent set comprehension.

I used this OPA version for testing:

$ opa version
Version: 0.59.0
Build Commit: c8e7863c83e1767cf019b0e8429decc9c4bdb9cd
Build Timestamp: 2023-11-30T15:17:55Z
Build Hostname: 
Go Version: go1.21.3
Platform: linux/amd64
WebAssembly: unavailable

charlesdaniels avatar Dec 14 '23 21:12 charlesdaniels