diff --git a/pkg/compiler/internal/loop_collect.go b/pkg/compiler/internal/loop_collect.go index 19e0fabe..31e00b3d 100644 --- a/pkg/compiler/internal/loop_collect.go +++ b/pkg/compiler/internal/loop_collect.go @@ -40,10 +40,8 @@ func (c *LoopCollectCompiler) compileCollect(ctx fql.ICollectClauseContext, aggr return core.NewKV(vm.NoopOperand, vm.NoopOperand), nil } - loop := c.ctx.Loops.Current() - kv, groupSelectors := c.initializeCollector(grouping) - projectionVarName, collectorType := c.initializeProjection(ctx, loop, kv, counter, grouping != nil) + projectionVarName, collectorType := c.initializeProjection(ctx, kv, counter, grouping != nil) // If we use aggregators, we need to collect group items by key if aggregation && collectorType != core.CollectorTypeKeyGroup { @@ -51,6 +49,7 @@ func (c *LoopCollectCompiler) compileCollect(ctx fql.ICollectClauseContext, aggr collectorType = core.CollectorTypeKeyGroup } + loop := c.ctx.Loops.Current() c.finalizeCollector(loop, collectorType, kv) // If we are using a projection, we need to ensure the loop is set to ForInLoop @@ -115,13 +114,13 @@ func (c *LoopCollectCompiler) finalizeCollector(loop *core.Loop, collectorType c // initializeProjection handles the projection setup for group variables and counters. // Returns the projection variable name and the appropriate collector type. -func (c *LoopCollectCompiler) initializeProjection(ctx fql.ICollectClauseContext, loop *core.Loop, kv *core.KV, counter fql.ICollectCounterContext, hasGrouping bool) (string, core.CollectorType) { +func (c *LoopCollectCompiler) initializeProjection(ctx fql.ICollectClauseContext, kv *core.KV, counter fql.ICollectCounterContext, hasGrouping bool) (string, core.CollectorType) { projectionVariableName := "" collectorType := core.CollectorTypeKey // Handle group variable projection if groupVar := ctx.CollectGroupVariable(); groupVar != nil { - projectionVariableName = c.compileGroupVariableProjection(loop, kv, groupVar) + projectionVariableName = c.compileGroupVariableProjection(kv, groupVar) collectorType = core.CollectorTypeKeyGroup return projectionVariableName, collectorType } @@ -176,15 +175,15 @@ func (c *LoopCollectCompiler) compileGrouping(ctx fql.ICollectGroupingContext) ( } // compileGroupVariableProjection processes group variable projections (both default and custom). -func (c *LoopCollectCompiler) compileGroupVariableProjection(loop *core.Loop, kv *core.KV, groupVar fql.ICollectGroupVariableContext) string { +func (c *LoopCollectCompiler) compileGroupVariableProjection(kv *core.KV, groupVar fql.ICollectGroupVariableContext) string { // Handle default projection (identifier) if identifier := groupVar.Identifier(); identifier != nil { - return c.compileDefaultGroupProjection(loop, kv, identifier, groupVar.CollectGroupVariableKeeper()) + return c.compileDefaultGroupProjection(kv, identifier, groupVar.CollectGroupVariableKeeper()) } // Handle custom projection (selector expression) if selector := groupVar.CollectSelector(); selector != nil { - return c.compileCustomGroupProjection(loop, kv, selector) + return c.compileCustomGroupProjection(kv, selector) } return "" @@ -226,7 +225,7 @@ func (c *LoopCollectCompiler) compileGroupSelectorVariables(selectors []fql.ICol } } -func (c *LoopCollectCompiler) compileDefaultGroupProjection(loop *core.Loop, kv *core.KV, identifier antlr.TerminalNode, keeper fql.ICollectGroupVariableKeeperContext) string { +func (c *LoopCollectCompiler) compileDefaultGroupProjection(kv *core.KV, identifier antlr.TerminalNode, keeper fql.ICollectGroupVariableKeeperContext) string { if keeper == nil { variables := c.ctx.Symbols.LocalVariables() scope := core.NewScopeProjection(c.ctx.Registers, c.ctx.Emitter, c.ctx.Symbols, variables) @@ -255,7 +254,7 @@ func (c *LoopCollectCompiler) compileDefaultGroupProjection(loop *core.Loop, kv return identifier.GetText() } -func (c *LoopCollectCompiler) compileCustomGroupProjection(_ *core.Loop, kv *core.KV, selector fql.ICollectSelectorContext) string { +func (c *LoopCollectCompiler) compileCustomGroupProjection(kv *core.KV, selector fql.ICollectSelectorContext) string { selectorReg := c.ctx.ExprCompiler.Compile(selector.Expression()) c.ctx.Emitter.EmitMove(kv.Value, selectorReg) c.ctx.Registers.Free(selectorReg) diff --git a/test/integration/vm/vm_for_in_collect_agg_test.go b/test/integration/vm/vm_for_in_collect_agg_test.go index 4c16b20c..e981988e 100644 --- a/test/integration/vm/vm_for_in_collect_agg_test.go +++ b/test/integration/vm/vm_for_in_collect_agg_test.go @@ -513,5 +513,60 @@ FOR u IN users "uniqueSkillCount": 4, }, }, "Should aggregate with array operations"), + 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 i IN 0..4 + LET u = users[i] + COLLECT gender = u.gender + AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age) + + RETURN { + gender: gender, + minAge, + maxAge + } +`, []any{ + map[string]any{ + "gender": "f", + "maxAge": 45, + "minAge": 25, + }, + map[string]any{ + "gender": "m", + "maxAge": 69, + "minAge": 31, + }, + }), }) } diff --git a/test/integration/vm/vm_for_while_nested_test.go b/test/integration/vm/vm_for_while_nested_test.go index 53e72965..8c1e9b9e 100644 --- a/test/integration/vm/vm_for_while_nested_test.go +++ b/test/integration/vm/vm_for_while_nested_test.go @@ -610,7 +610,7 @@ func TestForWhileNested(t *testing.T) { }, }, }), - SkipCaseArray(` + CaseArray(` LET users = [ { active: true, @@ -652,40 +652,45 @@ func TestForWhileNested(t *testing.T) { gender: CONCAT(gender, n), values: genders } -`, []any{map[string]any{ - "gender": "f0", - "values": []map[string]any{ - { - "i": map[string]any{ - "active": true, - "age": 25, - "gender": "f", - "married": false, - }, - }, - { - "i": map[string]any{ - "active": true, - "age": 45, - "gender": "f", - "married": true, - }, - }, - }, - }, +`, []any{ map[string]any{ - "gender": "f1", - "values": []map[string]any{ - { - "i": map[string]any{ + "gender": "f0", + "values": []any{ + map[string]any{ + "i": 1, + "u": map[string]any{ "active": true, "age": 25, "gender": "f", "married": false, }, }, - { - "i": map[string]any{ + map[string]any{ + "i": 4, + "u": map[string]any{ + "active": true, + "age": 45, + "gender": "f", + "married": true, + }, + }, + }, + }, + map[string]any{ + "gender": "f1", + "values": []any{ + map[string]any{ + "i": 1, + "u": map[string]any{ + "active": true, + "age": 25, + "gender": "f", + "married": false, + }, + }, + map[string]any{ + "i": 4, + "u": map[string]any{ "active": true, "age": 45, "gender": "f", @@ -696,25 +701,28 @@ func TestForWhileNested(t *testing.T) { }, map[string]any{ "gender": "m0", - "values": []map[string]any{ - { - "i": map[string]any{ + "values": []any{ + map[string]any{ + "i": 0, + "u": map[string]any{ "active": true, "age": 31, "gender": "m", "married": true, }, }, - { - "i": map[string]any{ + map[string]any{ + "i": 2, + "u": map[string]any{ "active": true, "age": 36, "gender": "m", "married": false, }, }, - { - "i": map[string]any{ + map[string]any{ + "i": 3, + "u": map[string]any{ "active": false, "age": 69, "gender": "m", @@ -725,25 +733,28 @@ func TestForWhileNested(t *testing.T) { }, map[string]any{ "gender": "m1", - "values": []map[string]any{ - { - "i": map[string]any{ + "values": []any{ + map[string]any{ + "i": 0, + "u": map[string]any{ "active": true, "age": 31, "gender": "m", "married": true, }, }, - { - "i": map[string]any{ + map[string]any{ + "i": 2, + "u": map[string]any{ "active": true, "age": 36, "gender": "m", "married": false, }, }, - { - "i": map[string]any{ + map[string]any{ + "i": 3, + "u": map[string]any{ "active": false, "age": 69, "gender": "m", @@ -753,7 +764,7 @@ func TestForWhileNested(t *testing.T) { }, }, }), - SkipCaseArray(` + CaseArray(` LET users = [ { active: true,