Refactor loop aggregation logic; optimize destination register handling, adjust collector allocation, and improve resource management in nested aggregation scenarios. Update integration tests with consistent formatting for nested FOR cases.

2025-08-15 20:02:56 +02:00 · 2025-06-25 16:23:39 -04:00
parent dd817e7115
commit e60a02ec12
2 changed files with 431 additions and 245 deletions
--- a/pkg/compiler/internal/loop_collect_agg.go
+++ b/pkg/compiler/internal/loop_collect_agg.go
@@ -49,14 +49,12 @@ func (c *LoopCollectCompiler) compileGroupedAggregation(ctx fql.ICollectAggregat
 func (c *LoopCollectCompiler) compileGlobalAggregation(ctx fql.ICollectAggregatorContext) {
 	parentLoop := c.ctx.Loops.Current()
 	// we create a custom collector for aggregators
-	c.ctx.Emitter.PatchSwapAx(parentLoop.Pos, vm.OpDataSetCollector, parentLoop.Dst, int(core.CollectorTypeKeyGroup))
-
+	dst := parentLoop.PatchDestinationAx(c.ctx.Registers, c.ctx.Emitter, vm.OpDataSetCollector, int(core.CollectorTypeKeyGroup))
 	// Nested scope for aggregators
 	c.ctx.Symbols.EnterScope()
 	// Now we add value selectors to the collector
 	selectors := ctx.AllCollectAggregateSelector()
-	argsPkg := c.compileAggregationFuncArgs(selectors, parentLoop.Dst)
-
+	argsPkg := c.compileAggregationFuncArgs(selectors, dst)
 	parentLoop.EmitFinalization(c.ctx.Emitter)
 	c.ctx.Loops.Pop()
 	c.ctx.Symbols.ExitScope()
@@ -66,14 +64,18 @@ func (c *LoopCollectCompiler) compileGlobalAggregation(ctx fql.ICollectAggregato
 	c.ctx.Emitter.EmitA(vm.OpLoadZero, zero)
 	// We move the aggregator to a temporary register to access it later from the new loop
 	aggregator := c.ctx.Registers.Allocate(core.Temp)
-	c.ctx.Emitter.EmitAB(vm.OpMove, aggregator, parentLoop.Dst)
+	c.ctx.Emitter.EmitAB(vm.OpMove, aggregator, dst)
+
+	if parentLoop.Dst != dst && !parentLoop.Allocate {
+		c.ctx.Registers.Free(dst)
+	}

 	// CreateFor new loop with 1 iteration only
 	c.ctx.Symbols.EnterScope()
 	c.ctx.Emitter.EmitABC(vm.OpRange, parentLoop.Src, zero, zero)
 	loop := c.ctx.Loops.CreateFor(core.TemporalLoop, parentLoop.Src, parentLoop.Distinct)
 	loop.Dst = parentLoop.Dst
-	loop.Allocate = true
+	loop.Allocate = parentLoop.Allocate
 	c.ctx.Loops.Push(loop)
 	loop.EmitInitialization(c.ctx.Registers, c.ctx.Emitter)

--- a/test/integration/vm/vm_for_nested_test.go
+++ b/test/integration/vm/vm_for_nested_test.go
@@ -8,48 +8,48 @@ import (

 func TestForNested(t *testing.T) {
 	RunUseCases(t, []UseCase{
-		CaseArray(
-			`FOR prop IN ["a"]
-							FOR val IN [1, 2, 3]
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
-		),
-		CaseArray(
-			`FOR val IN 1..3
+		CaseArray(`
 			FOR prop IN ["a"]
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
+				FOR val IN [1, 2, 3]
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
 		),
-		CaseArray(
-			`FOR prop IN ["a"]
+		CaseArray(`
 			FOR val IN 1..3
-								RETURN {[prop]: val}`,
-			[]any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
+				FOR prop IN ["a"]
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
 		),
-		CaseArray(
-			`FOR prop IN ["a"]
+		CaseArray(`
+			FOR prop IN ["a"]
+				FOR val IN 1..3
+					RETURN {[prop]: val}
+`, []any{map[string]any{"a": 1}, map[string]any{"a": 2}, map[string]any{"a": 3}},
+		),
+		CaseArray(`
+			FOR prop IN ["a"]
 				FOR val IN [1, 2, 3]
 					FOR val2 IN [1, 2, 3]
-									RETURN { [prop]: [val, val2] }`,
-			[]any{map[string]any{"a": []int{1, 1}}, map[string]any{"a": []int{1, 2}}, map[string]any{"a": []int{1, 3}}, map[string]any{"a": []int{2, 1}}, map[string]any{"a": []int{2, 2}}, map[string]any{"a": []int{2, 3}}, map[string]any{"a": []int{3, 1}}, map[string]any{"a": []int{3, 2}}, map[string]any{"a": []int{3, 3}}},
+						RETURN { [prop]: [val, val2] }
+`, []any{map[string]any{"a": []int{1, 1}}, map[string]any{"a": []int{1, 2}}, map[string]any{"a": []int{1, 3}}, map[string]any{"a": []int{2, 1}}, map[string]any{"a": []int{2, 2}}, map[string]any{"a": []int{2, 3}}, map[string]any{"a": []int{3, 1}}, map[string]any{"a": []int{3, 2}}, map[string]any{"a": []int{3, 3}}},
 		),
-		CaseArray(
-			`FOR val IN [1, 2, 3]
+		CaseArray(`
+			FOR val IN [1, 2, 3]
 				RETURN (
 					FOR prop IN ["a", "b", "c"]
 						RETURN { [prop]: val }
-							)`,
-			[]any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
+				)
+`, []any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
 		),
-		CaseArray(
-			`FOR val IN [1, 2, 3]
+		CaseArray(`
+			FOR val IN [1, 2, 3]
 			LET sub = (
 				FOR prop IN ["a", "b", "c"]
 					RETURN { [prop]: val }
 			)

-							RETURN sub`,
-			[]any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
+			RETURN sub
+`, []any{[]any{map[string]any{"a": 1}, map[string]any{"b": 1}, map[string]any{"c": 1}}, []any{map[string]any{"a": 2}, map[string]any{"b": 2}, map[string]any{"c": 2}}, []any{map[string]any{"a": 3}, map[string]any{"b": 3}, map[string]any{"c": 3}}},
 		),
 		CaseArray(`
 			LET strs = ["foo", "bar", "qaz", "abc"]
@@ -467,5 +467,189 @@ FOR i IN users
 				},
 			},
 		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR u IN users
+				  COLLECT gender = u.gender
+					AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				
+				  RETURN {
+					gender: CONCAT(gender, n),
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"gender": "f0",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m0",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+			map[string]any{
+				"gender": "f1",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m1",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			
+			FOR u IN users
+			  COLLECT gender = u.gender
+				AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				FOR n IN 0..1
+				  RETURN {
+					gender: CONCAT(gender, n),
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"gender": "f0",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "f1",
+				"maxAge": 45,
+				"minAge": 25,
+			},
+			map[string]any{
+				"gender": "m0",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+			map[string]any{
+				"gender": "m1",
+				"maxAge": 69,
+				"minAge": 31,
+			},
+		}),
+		CaseArray(`
+			LET users = [
+				{
+					active: true,
+					married: true,
+					age: 31,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: false,
+					age: 25,
+					gender: "f"
+				},
+				{
+					active: true,
+					married: false,
+					age: 36,
+					gender: "m"
+				},
+				{
+					active: false,
+					married: true,
+					age: 69,
+					gender: "m"
+				},
+				{
+					active: true,
+					married: true,
+					age: 45,
+					gender: "f"
+				}
+			]
+			FOR n IN 0..1
+				FOR u IN users
+  				COLLECT AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+				
+				  RETURN {
+					iteration: n,
+					minAge,
+					maxAge
+				  }
+`, []any{
+			map[string]any{
+				"iteration": 0,
+				"maxAge":    69,
+				"minAge":    25,
+			},
+			map[string]any{
+				"iteration": 1,
+				"maxAge":    69,
+				"minAge":    25,
+			},
+		}),
 	})
 }