From 70c8746cedde3956f025dbbfac64c0f3d16ee4cc Mon Sep 17 00:00:00 2001 From: Ville Vesilehto Date: Thu, 8 Jan 2026 20:52:57 +0200 Subject: [PATCH] perf(topdown): optimize json.remove and .filter (#8193) Optimize `json.remove` and `json.filter` to reduce memory allocations and copy overhead. Changes: - Pre-allocate result collections in `jsonRemove` using the source size as an upper bound. - Refactor `jsonRemove` for Arrays to accumulate elements in a pre-allocated Go slice before creating the AST Array, avoiding intermediate allocations. - Pre-allocate the root object in `pathsToObject` used by both built-ins. - Optimize `ast.Object.Filter` (used by `json.filter`) for Arrays to avoid incremental `Append` allocations. Signed-off-by: Ville Vesilehto --- v1/ast/term.go | 6 +- v1/topdown/json.go | 12 ++-- v1/topdown/json_bench_test.go | 124 ++++++++++++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 9 deletions(-) diff --git a/v1/ast/term.go b/v1/ast/term.go index cbe2b191e2..01184f2fbf 100644 --- a/v1/ast/term.go +++ b/v1/ast/term.go @@ -2570,7 +2570,7 @@ func filterObject(o Value, filter Value) (Value, error) { case String, Number, Boolean, Null: return o, nil case *Array: - values := NewArray() + values := make([]*Term, 0, v.Len()) for i := range v.Len() { subFilter := filteredObj.Get(InternedIntegerString(i)) if subFilter != nil { @@ -2578,10 +2578,10 @@ func filterObject(o Value, filter Value) (Value, error) { if err != nil { return nil, err } - values = values.Append(NewTerm(filteredValue)) + values = append(values, NewTerm(filteredValue)) } } - return values, nil + return NewArray(values...), nil case Set: terms := make([]*Term, 0, v.Len()) for _, t := range v.Slice() { diff --git a/v1/topdown/json.go b/v1/topdown/json.go index 2c7d642883..d37b70674d 100644 --- a/v1/topdown/json.go +++ b/v1/topdown/json.go @@ -65,7 +65,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) { case ast.String, ast.Number, ast.Boolean, ast.Null: return a, nil case ast.Object: - newObj := ast.NewObject() + newObj := ast.NewObjectWithCapacity(aValue.Len()) err := aValue.Iter(func(k *ast.Term, v *ast.Term) error { // recurse and add the diff of sub objects as needed diffValue, err := jsonRemove(v, bObj.Get(k)) @@ -80,7 +80,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) { } return ast.NewTerm(newObj), nil case ast.Set: - newSet := ast.NewSet() + newSet := ast.NewSetWithCapacity(aValue.Len()) err := aValue.Iter(func(v *ast.Term) error { // recurse and add the diff of sub objects as needed diffValue, err := jsonRemove(v, bObj.Get(v)) @@ -97,7 +97,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) { case *ast.Array: // When indexes are removed we shift left to close empty spots in the array // as per the JSON patch spec. - newArray := ast.NewArray() + newArraySlice := make([]*ast.Term, 0, aValue.Len()) for i := range aValue.Len() { v := aValue.Elem(i) // recurse and add the diff of sub objects as needed @@ -107,10 +107,10 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) { return nil, err } if diffValue != nil { - newArray = newArray.Append(diffValue) + newArraySlice = append(newArraySlice, diffValue) } } - return ast.NewTerm(newArray), nil + return ast.NewTerm(ast.NewArray(newArraySlice...)), nil default: return nil, fmt.Errorf("invalid value type %T", a) } @@ -196,7 +196,7 @@ func parsePath(path *ast.Term) (ast.Ref, error) { } func pathsToObject(paths []ast.Ref) ast.Object { - root := ast.NewObject() + root := ast.NewObjectWithCapacity(len(paths)) for _, path := range paths { node := root diff --git a/v1/topdown/json_bench_test.go b/v1/topdown/json_bench_test.go index 6e2011b0ad..1cf7fa38ee 100644 --- a/v1/topdown/json_bench_test.go +++ b/v1/topdown/json_bench_test.go @@ -17,6 +17,130 @@ import ( inmem "github.com/open-policy-agent/opa/v1/storage/inmem/test" ) +func BenchmarkJSONRemoveArray(b *testing.B) { + sizes := []int{10, 100, 1000, 5000} + + for _, n := range sizes { + b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) { + // Create an object wrapping the array: {"a": [0, 1, ...]} + terms := make([]*ast.Term, n) + for i := range n { + terms[i] = ast.IntNumberTerm(i) + } + arr := ast.NewArray(terms...) + obj := ast.NewObject([2]*ast.Term{ast.StringTerm("a"), ast.NewTerm(arr)}) + + // Remove something inside the array to force traversal + paths := ast.NewSet(ast.StringTerm("a/nonexistent")) + + operands := []*ast.Term{ + ast.NewTerm(obj), + ast.NewTerm(paths), + } + + for b.Loop() { + if err := builtinJSONRemove( + BuiltinContext{ + Context: context.Background(), + }, + operands, + func(*ast.Term) error { + return nil + }, + ); err != nil { + b.Fatal(err) + } + } + }) + } +} + +func BenchmarkJSONFilterArray(b *testing.B) { + sizes := []int{10, 100, 1000, 5000} + + for _, n := range sizes { + b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) { + // Create an object with n keys + obj := ast.NewObject() + pathSlice := make([]*ast.Term, n) + for i := range n { + k := ast.StringTerm(fmt.Sprintf("k%d", i)) + obj.Insert(k, ast.IntNumberTerm(i)) + pathSlice[i] = k + } + // Filter all keys: json.filter(obj, ["k0", "k1", ...]) + // This stresses pathsToObject (creating the filter mask) + paths := ast.NewSet(pathSlice...) + + operands := []*ast.Term{ + ast.NewTerm(obj), + ast.NewTerm(paths), + } + + for b.Loop() { + if err := builtinJSONFilter( + BuiltinContext{ + Context: context.Background(), + }, + operands, + func(*ast.Term) error { + return nil + }, + ); err != nil { + b.Fatal(err) + } + } + }) + } +} + +func BenchmarkJSONFilterArrayIndices(b *testing.B) { + sizes := []int{10, 100, 1000, 5000} + + for _, n := range sizes { + b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) { + // Create an object wrapping an array: {"a": [0, 1, ...]} + terms := make([]*ast.Term, n) + for i := range n { + terms[i] = ast.IntNumberTerm(i) + } + arr := ast.NewArray(terms...) + obj := ast.NewObject([2]*ast.Term{ast.StringTerm("a"), ast.NewTerm(arr)}) + + // Filter to keep the first half of the array elements + // json.filter(obj, ["a/0", "a/1", ... "a/n/2"]) + filterSize := n / 2 + if filterSize == 0 { + filterSize = 1 + } + pathSlice := make([]*ast.Term, filterSize) + for i := range filterSize { + pathSlice[i] = ast.StringTerm(fmt.Sprintf("a/%d", i)) + } + paths := ast.NewSet(pathSlice...) + + operands := []*ast.Term{ + ast.NewTerm(obj), + ast.NewTerm(paths), + } + + for b.Loop() { + if err := builtinJSONFilter( + BuiltinContext{ + Context: context.Background(), + }, + operands, + func(*ast.Term) error { + return nil + }, + ); err != nil { + b.Fatal(err) + } + } + }) + } +} + func BenchmarkJSONPatchAddShallowScalar(b *testing.B) { ctx := b.Context()