perf: update the accumulated value of a reduce expression in place

_assign and _modify create an allocator on every call, so the value they
are given is never one that allocator created, and updateObject and
updateArrayIndex always take the copying branch. Accumulating into an
object or an array with a reduce expression therefore copies the whole
accumulated value on every iteration, which is quadratic in the number
of elements. The query of #241,

    [range(20000)] | reduce .[] as $item ({}; .[($item|tostring)] = {})

took between 10 and 21 seconds here over three runs, against 0.03
seconds for jq.

Create the allocator once per reduce expression and pass it to the
assignment as a third argument, so that the accumulated value is
allocated on the first iteration and modified in place afterwards. The
query above now takes 0.03 seconds, and the memory allocated by an
accumulating reduction grows linearly instead of quadratically: 2.1 MB
instead of 545 MB for 4000 elements.

This is only correct while the accumulated value, and the values the
allocator created inside it, are reachable through the reduction alone.
Query#inPlaceUpdate therefore requires the update query to be a single
assignment, and its right-hand side to neither observe the input value,
which would let the assigned value alias the value it is assigned into,
nor yield more than one value, which would run the assignment again
against a value it already modified. Everything else keeps copying. A
foreach expression is left alone, because it yields every intermediate
accumulated value.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
dualfrozandClaude Opus 5 committed 2026-09-04 08:23:49 +02:00
1 parent c5274c6af4
commit 70b3f7d698
4 files changed
+477 -27

No files matched your search

+62
View File
@@ -3609,6 +3609,68 @@
expected: |
[[1,1],[1]]
- name: reduce with assignment in update
args:
- -c
- 'reduce range(3) as $i ({}; .[$i|tostring] = {}),
reduce range(3) as $i ({}; .a[$i|tostring] = $i),
reduce range(3) as $i ([]; .[$i] = $i),
reduce range(3) as $i ({}; .[$i|tostring] += 1)'
input: 'null'
expected: |
{"0":{},"1":{},"2":{}}
{"a":{"0":0,"1":1,"2":2}}
[0,1,2]
{"0":1,"1":1,"2":1}
- name: reduce with assignment of the accumulated value
args:
- -c
- 'reduce range(3) as $i ({}; .a = .),
reduce range(3) as $i ({}; .a = {v: .}),
reduce range(3) as $i ({}; . as $p | .a = $i | .b = $p)'
input: 'null'
expected: |
{"a":{"a":{"a":{}}}}
{"a":{"v":{"a":{"v":{"a":{"v":{}}}}}}}
{"a":2,"b":{"a":1,"b":{"a":0,"b":{}}}}
- name: reduce with assignment of multiple values
args:
- -c
- 'reduce range(2) as $i ({"k":"k"}; .[.k] = ("x","y")),
reduce range(2) as $i ({}; (.a = 1), (.b = 2))'
input: 'null'
expected: |
{"k":"y","y":"y"}
{"b":2}
- name: reduce with assignment of a shared value
args:
- -c
- '{"a":1} as $x | reduce range(2) as $i ({}; .[$i|tostring] = $x) |
[(.["0"].a = 9), $x]'
input: 'null'
expected: |
[{"0":{"a":9},"1":{"a":1}},{"a":1}]
- name: reduce with assignment result reused
args:
- -c
- 'reduce range(3) as $i ({}; .[$i|tostring] = $i) as $r |
[($r | .["0"] = 9), $r]'
input: 'null'
expected: |
[{"0":9,"1":1,"2":2},{"0":0,"1":1,"2":2}]
- name: foreach with assignment in update
args:
- -c
- '[foreach range(3) as $i ({}; .[$i|tostring] = $i)]'
input: 'null'
expected: |
[{"0":0},{"0":0,"1":1},{"0":0,"1":1,"2":2}]
- name: foreach iterator
args:
- 'foreach range(5) as $item (0; $item)'
+136 -27
View File
@@ -397,7 +397,7 @@ func (c *compiler) compileQuery(e *Query) error {
return c.compileAlt(e.Left, e.Right)
case OpAssign, OpModify, OpUpdateAdd, OpUpdateSub,
OpUpdateMul, OpUpdateDiv, OpUpdateMod, OpUpdateAlt:
return c.compileQueryUpdate(e.Left, e.Right, e.Op)
return c.compileQueryUpdate(e.Left, e.Right, e.Op, nil)
case OpOr:
return c.compileIf(
&If{
@@ -471,7 +471,11 @@ func (c *compiler) compileAlt(l, r *Query) error {
return c.compileQuery(r)
}
func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
// compileQueryUpdate appends the compiled code for an assignment operator. The
// alloc query is nil unless the assignment may update the input value in
// place, in which case it yields the allocator of the enclosing reduce
// expression; see compiler#compileReduce.
func (c *compiler) compileQueryUpdate(l, r *Query, op Operator, alloc *Query) error {
switch op {
case OpAssign:
// optimize assignment operator with constant indexing and slicing
@@ -491,12 +495,7 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
}
fallthrough
case OpModify:
return c.compileFunc(
&Func{
Name: op.getFunc(),
Args: []*Query{l, r},
},
)
return c.compileUpdateFunc(op.getFunc(), []*Query{l, r}, alloc)
default:
name := "$%0"
c.append(&code{op: opdup})
@@ -504,10 +503,9 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
return err
}
c.append(&code{op: opstore, v: c.pushVariable(name)})
return c.compileFunc(
&Func{
Name: "_modify",
Args: []*Query{
return c.compileUpdateFunc(
"_modify",
[]*Query{
l,
{Term: &Term{
Type: TermTypeFunc,
@@ -520,11 +518,44 @@ func (c *compiler) compileQueryUpdate(l, r *Query, op Operator) error {
},
}},
},
},
alloc,
)
}
}
const (
// Appended to the names of the _assign and _modify variants taking an
// allocator as an additional argument. The resulting name is not a valid
// function name, so that a query cannot call such a variant directly with
// an arbitrary value in place of the allocator.
allocatorArgSuffix = "/3"
// The name of the variable holding the allocator of a reduce expression.
// It is not a valid variable name, just like the variable of the operand
// of an update-assignment operator.
allocatorName = "$%allocator"
)
// compileUpdateFunc appends a call to _assign or _modify, or to the variant
// taking the allocator to update the input value with as its third argument.
func (c *compiler) compileUpdateFunc(name string, args []*Query, alloc *Query) error {
if alloc == nil {
return c.compileFunc(&Func{Name: name, Args: args})
}
args = append(args, alloc)
fn := c.lookupBuiltin(name+allocatorArgSuffix, len(args))
if fn == nil {
switch name {
case "_assign":
c.compileAssign(true)
case "_modify":
c.compileModify(true)
}
fn = c.lookupBuiltin(name+allocatorArgSuffix, len(args))
}
return c.compileCallPc(fn, args)
}
func (c *compiler) compileBind(l, r *Query, patterns []*Pattern) error {
defer c.newScopeDepth()()
c.append(&code{op: opdup})
@@ -720,6 +751,25 @@ func (c *compiler) compileReduce(e *Reduce) error {
}
f()
c.append(&code{op: opstore, v: v})
// The accumulated value is passed to the update query and replaced by its
// result, so an assignment forming the whole update query is the only
// thing that observes it. Such an assignment may modify the accumulated
// value in place, instead of copying it on every iteration, as long as it
// keeps track of the values it allocated; see Query#inPlaceUpdate. The
// allocator is created once, outside the loop, and lives until the
// reduction finishes.
l, r, op, inPlace := e.Update.inPlaceUpdate()
var alloc *Query
if inPlace {
c.appends(
&code{op: oppush, v: nil},
&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
&code{op: opstore, v: c.pushVariable(allocatorName)},
)
alloc = &Query{Term: &Term{
Type: TermTypeFunc, Func: &Func{Name: allocatorName},
}}
}
setfork := c.lazy(func() *code {
return &code{op: opfork, v: len(c.codes)}
})
@@ -731,7 +781,13 @@ func (c *compiler) compileReduce(e *Reduce) error {
}
c.append(&code{op: opload, v: v})
f = c.newScopeDepth()
if err := c.compileQuery(e.Update); err != nil {
var err error
if inPlace {
err = c.compileQueryUpdate(l, r, op, alloc)
} else {
err = c.compileQuery(e.Update)
}
if err != nil {
return err
}
f()
@@ -936,9 +992,9 @@ func (c *compiler) compileFunc(e *Func) error {
if !compiled {
switch e.Name {
case "_assign":
c.compileAssign()
c.compileAssign(false)
case "_modify":
c.compileModify()
c.compileModify(false)
case "_last":
c.compileLast()
}
@@ -1034,29 +1090,56 @@ func (c *compiler) compileFunc(e *Func) error {
//
// To overcome the difficulty of reducing allocations on `setpath`, we use the
// `allocator` type and track the allocated addresses during the reduction.
func (c *compiler) compileAssign() {
defer c.appendBuiltin("_assign", 2)()
//
// The shared variant takes the allocator of the enclosing reduce expression as
// its third argument instead of creating a new one, so that the value being
// assigned to is modified in place on every iteration but the first; see
// compiler#compileReduce.
func (c *compiler) compileAssign(shared bool) {
name, argcnt := "_assign", 2
if shared {
name, argcnt = name+allocatorArgSuffix, 3
}
defer c.appendBuiltin(name, argcnt)()
scope := c.newScope()
v, p := [2]int{scope.id, 0}, [2]int{scope.id, 1}
x, a := [2]int{scope.id, 2}, [2]int{scope.id, 3}
// Cannot reuse v, p due to backtracking in x.
w, q := [2]int{scope.id, 4}, [2]int{scope.id, 5}
c.appends(
&code{op: opscope, v: [3]int{scope.id, 6, 2}},
&code{op: opscope, v: [3]int{scope.id, 6, argcnt}},
&code{op: opstore, v: v}, // def _assign(p; $x):
&code{op: opstore, v: p},
&code{op: opstore, v: x},
)
if shared {
c.append(&code{op: opstore, v: a}) // the allocator argument
}
c.appends(
&code{op: opload, v: v},
&code{op: opexpbegin},
&code{op: opload, v: x},
&code{op: opcallpc},
&code{op: opstore, v: x},
&code{op: opexpend},
)
if shared {
c.appends(
&code{op: opload, v: v},
&code{op: opload, v: a},
&code{op: opcallpc},
&code{op: opstore, v: a},
)
} else {
c.appends(
&code{op: oppush, v: nil},
&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
&code{op: opstore, v: a},
)
}
c.appends(
&code{op: opload, v: v},
&code{op: opfork, v: len(c.codes) + 30}, // reduce [L1]
&code{op: opfork, v: len(c.codes) + 17}, // reduce [L1]
&code{op: opdup},
&code{op: opstore, v: w},
&code{op: oppathbegin}, // path(p)
@@ -1080,24 +1163,50 @@ func (c *compiler) compileAssign() {
// Appends the compiled code for the update-assignment operator (`|=`) to
// maximize performance. We use the `allocator` type, just like `_assign/2`.
func (c *compiler) compileModify() {
defer c.appendBuiltin("_modify", 2)()
//
// The shared variant takes the allocator of the enclosing reduce expression as
// its third argument instead of creating a new one; see
// compiler#compileReduce and compiler#compileAssign.
func (c *compiler) compileModify(shared bool) {
name, argcnt := "_modify", 2
if shared {
name, argcnt = name+allocatorArgSuffix, 3
}
defer c.appendBuiltin(name, argcnt)()
scope := c.newScope()
v, p := [2]int{scope.id, 0}, [2]int{scope.id, 1}
f, d := [2]int{scope.id, 2}, [2]int{scope.id, 3}
a, l := [2]int{scope.id, 4}, [2]int{scope.id, 5}
c.appends(
&code{op: opscope, v: [3]int{scope.id, 6, 2}},
&code{op: opscope, v: [3]int{scope.id, 6, argcnt}},
&code{op: opstore, v: v}, // def _modify(p; f):
&code{op: opstore, v: p},
&code{op: opstore, v: f},
)
if shared {
c.append(&code{op: opstore, v: a}) // the allocator argument
}
c.appends(
&code{op: oppush, v: []any{}},
&code{op: opstore, v: d},
)
if shared {
c.appends(
&code{op: opload, v: v},
&code{op: opload, v: a},
&code{op: opcallpc},
&code{op: opstore, v: a},
)
} else {
c.appends(
&code{op: oppush, v: nil},
&code{op: opcall, v: [3]any{funcAllocator, 0, "_allocator"}},
&code{op: opstore, v: a},
)
}
c.appends(
&code{op: opload, v: v},
&code{op: opfork, v: len(c.codes) + 40}, // reduce [L1]
&code{op: opfork, v: len(c.codes) + 31}, // reduce [L1]
&code{op: oppathbegin}, // path(p)
&code{op: opload, v: p},
&code{op: opcallpc},
@@ -1106,7 +1215,7 @@ func (c *compiler) compileModify() {
&code{op: opstore, v: p}, // as $p (.;
&code{op: opforklabel, v: l}, // label $l |
&code{op: opload, v: v}, //
&code{op: opfork, v: len(c.codes) + 37}, // [L2]
&code{op: opfork, v: len(c.codes) + 28}, // [L2]
&code{op: oppop}, // (getpath($p) |
&code{op: opload, v: a},
&code{op: opload, v: a},
@@ -1120,8 +1229,8 @@ func (c *compiler) compileModify() {
&code{op: opcall, v: [3]any{funcSetpathWithAllocator, 3, "_setpath"}},
&code{op: opstore, v: v},
&code{op: opload, v: v}, // ., break $l
&code{op: opfork, v: len(c.codes) + 35}, // [L4]
&code{op: opjump, v: len(c.codes) + 39}, // [L3]
&code{op: opfork, v: len(c.codes) + 26}, // [L4]
&code{op: opjump, v: len(c.codes) + 30}, // [L3]
&code{op: opload, v: l}, // [L4]
&code{op: opcall, v: [3]any{funcBreak(""), 0, "_break"}},
&code{op: opload, v: p}, // append $p to $d [L2]
+54
View File
@@ -7,6 +7,7 @@ import (
"log"
"os"
"reflect"
"runtime"
"sync"
"testing"
"time"
@@ -374,6 +375,59 @@ func TestCodeRun_RaceRegexp(t *testing.T) {
wg.Wait()
}
// An assignment forming the whole update query of a reduce expression modifies
// the accumulated value in place, so building an object or an array of n
// elements allocates memory proportional to n, not to n squared.
func TestCodeRun_ReduceAssignmentAllocation(t *testing.T) {
for _, src := range []string{
`reduce .[] as $x ({}; .[$x|tostring] = $x)`,
`reduce .[] as $x ({}; .[$x|tostring] |= $x)`,
`reduce .[] as $x ({}; .[$x|tostring] += 1)`,
`reduce .[] as $x ([]; .[$x] = $x)`,
} {
t.Run(src, func(t *testing.T) {
query, err := gojq.Parse(src)
if err != nil {
t.Fatal(err)
}
code, err := gojq.Compile(query)
if err != nil {
t.Fatal(err)
}
small, large := allocatedBytes(t, code, 1000), allocatedBytes(t, code, 4000)
// Copying the accumulated value on every iteration allocates about
// sixteen times as much memory for four times as many elements.
if expected := small * 8; large > expected {
t.Errorf("expected at most %v bytes for 4000 elements, got %v (%v bytes for 1000 elements)",
expected, large, small)
}
})
}
}
func allocatedBytes(t *testing.T, code *gojq.Code, n int) uint64 {
t.Helper()
input := make([]any, n)
for i := range input {
input[i] = i
}
var before, after runtime.MemStats
runtime.GC()
runtime.ReadMemStats(&before)
iter := code.Run(input)
for {
v, ok := iter.Next()
if !ok {
break
}
if err, ok := v.(error); ok {
t.Fatal(err)
}
}
runtime.ReadMemStats(&after)
return after.TotalAlloc - before.TotalAlloc
}
func BenchmarkCompile(b *testing.B) {
cnt, err := os.ReadFile("builtin.jq")
if err != nil {
+225
View File
@@ -0,0 +1,225 @@
package gojq
// This file implements the static analysis deciding whether the accumulated
// value of a reduce expression may be updated in place; see
// compiler#compileReduce for how the result is used, and the comment on
// compiler#compileAssign for why an allocator makes assignments cheap.
//
// Reusing an allocator across the iterations of a reduce expression is only
// correct while the accumulated value, and the values the allocator created
// inside it, are reachable through the accumulated value alone. The checks
// below are deliberately conservative; failing them only means the assignment
// keeps copying the accumulated value, which is the previous behavior.
// inPlaceUpdate reports whether the query is a single assignment applied to
// the input value, and returns its operands. Such an assignment receives the
// accumulated value of the enclosing reduce expression, and passes it on as
// the next accumulated value, without any other part of the query observing
// it in between.
func (e *Query) inPlaceUpdate() (l, r *Query, op Operator, ok bool) {
if e.Term != nil || len(e.FuncDefs) > 0 {
return nil, nil, 0, false
}
switch e.Op {
case OpModify:
// The update function is applied to the values found at the target
// paths, so it never observes the input value itself. When a target
// path points to a value the allocator created, _modify releases the
// allocator before calling the update function.
return e.Left, e.Right, e.Op, true
case OpAssign:
// An assignment to a constant path is compiled to a setpath call,
// which does not use an allocator; see compiler#compileQueryUpdate.
if e.Left.toIndices(nil) != nil {
return nil, nil, 0, false
}
fallthrough
case OpUpdateAdd, OpUpdateSub, OpUpdateMul,
OpUpdateDiv, OpUpdateMod, OpUpdateAlt:
// The right-hand side is evaluated against the input value, so it must
// not observe it; the assigned value would otherwise be able to alias
// the value it is assigned into. It must not yield multiple values
// either, because the assignment then runs once per value, each time
// against an input value the previous run already modified.
if !e.Right.readsInput() && e.Right.singleValued() {
return e.Left, e.Right, e.Op, true
}
}
return nil, nil, 0, false
}
// readsInput reports whether evaluating the query may observe its input value.
// It answers true unless it can prove otherwise.
func (e *Query) readsInput() bool {
if len(e.FuncDefs) > 0 {
return true
}
if e.Term != nil {
return e.Term.readsInput()
}
switch e.Op {
case OpPipe:
// The right-hand side is evaluated against the output of the
// left-hand side, unless the query binds variables, in which case it
// is evaluated against the input value as well.
if len(e.Patterns) > 0 {
return e.Left.readsInput() || e.Right.readsInput()
}
return e.Left.readsInput()
case OpComma, OpAdd, OpSub, OpMul, OpDiv, OpMod, OpEq, OpNe,
OpGt, OpLt, OpGe, OpLe, OpAnd, OpOr, OpAlt:
return e.Left.readsInput() || e.Right.readsInput()
default:
return true
}
}
// singleValued reports whether the query yields at most one value.
// It answers false unless it can prove otherwise.
func (e *Query) singleValued() bool {
if len(e.FuncDefs) > 0 {
return false
}
if e.Term != nil {
return e.Term.singleValued()
}
switch e.Op {
case OpPipe:
return len(e.Patterns) == 0 &&
e.Left.singleValued() && e.Right.singleValued()
case OpAdd, OpSub, OpMul, OpDiv, OpMod, OpEq, OpNe,
OpGt, OpLt, OpGe, OpLe, OpAnd, OpOr:
return e.Left.singleValued() && e.Right.singleValued()
default:
return false
}
}
func (e *Term) readsInput() bool {
for _, s := range e.SuffixList {
if s.readsInput() {
return true
}
}
switch e.Type {
case TermTypeNull, TermTypeTrue, TermTypeFalse, TermTypeNumber:
return false
case TermTypeString:
return e.Str.readsInput()
case TermTypeObject:
for _, kv := range e.Object.KeyVals {
if kv.readsInput() {
return true
}
}
return false
case TermTypeArray:
return e.Array.Query != nil && e.Array.Query.readsInput()
case TermTypeUnary:
return e.Unary.Term.readsInput()
case TermTypeFunc:
// Only a variable reference is known not to read the input value; any
// function may be redefined, so its name tells us nothing.
return e.Func.Name[0] != '$' || len(e.Func.Args) > 0
case TermTypeQuery:
return e.Query.readsInput()
default:
return true
}
}
func (e *Term) singleValued() bool {
for _, s := range e.SuffixList {
if !s.singleValued() {
return false
}
}
switch e.Type {
case TermTypeNull, TermTypeTrue, TermTypeFalse, TermTypeNumber:
return true
case TermTypeString:
return e.Str.singleValued()
case TermTypeObject:
for _, kv := range e.Object.KeyVals {
if !kv.singleValued() {
return false
}
}
return true
case TermTypeArray:
return true
case TermTypeUnary:
return e.Unary.Term.singleValued()
case TermTypeFunc:
return e.Func.Name[0] == '$' && len(e.Func.Args) == 0
case TermTypeQuery:
return e.Query.singleValued()
default:
return false
}
}
func (e *Suffix) readsInput() bool {
// Both iterating and the optional operator work on the value of the term,
// while an index is evaluated against the input value.
return e.Index != nil && e.Index.readsInput()
}
func (e *Suffix) singleValued() bool {
return !e.Iter && (e.Index == nil || e.Index.singleValued())
}
func (e *Index) readsInput() bool {
if e.Str != nil {
return e.Str.readsInput()
}
return e.Start != nil && e.Start.readsInput() ||
e.End != nil && e.End.readsInput()
}
func (e *Index) singleValued() bool {
if e.Str != nil {
return e.Str.singleValued()
}
return (e.Start == nil || e.Start.singleValued()) &&
(e.End == nil || e.End.singleValued())
}
func (e *String) readsInput() bool {
for _, q := range e.Queries {
if q.readsInput() {
return true
}
}
return false
}
func (e *String) singleValued() bool {
for _, q := range e.Queries {
if !q.singleValued() {
return false
}
}
return true
}
func (e *ObjectKeyVal) readsInput() bool {
if e.Val == nil {
// The shorthand {foo} and {"foo"} read the input value, while {$foo}
// does not.
return e.Key == "" || e.Key[0] != '$'
}
if e.KeyString != nil && e.KeyString.readsInput() ||
e.KeyQuery != nil && e.KeyQuery.readsInput() {
return true
}
return e.Val.readsInput()
}
func (e *ObjectKeyVal) singleValued() bool {
if e.KeyString != nil && !e.KeyString.singleValued() ||
e.KeyQuery != nil && !e.KeyQuery.singleValued() {
return false
}
return e.Val == nil || e.Val.singleValued()
}