Add support for grouping with counting in collect clauses

Introduced new opcode `OpCollectKc` for grouping with counting and updated the VM, compiler, and dataset logic to support this functionality. Enhanced the `CollectClause` to handle `WITH COUNT INTO` syntax, added `emitCollectCountProjection`, and implemented relevant test cases.

pkg/compiler/compiler_exec_test.go

@@ -2332,6 +2332,55 @@ LET users = [
 				},
 			},
 		}, "Should create default projection with custom KEEP with multiple custom names"),
+		CaseArray(
+			`LET users = [
+				{
+					active: true,
+					age: 31,
+					gender: "m",
+					married: true
+				},
+				{
+					active: true,
+					age: 25,
+					gender: "f",
+					married: false
+				},
+				{
+					active: true,
+					age: 36,
+					gender: "m",
+					married: false
+				},
+				{
+					active: false,
+					age: 69,
+					gender: "m",
+					married: true
+				},
+				{
+					active: true,
+					age: 45,
+					gender: "f",
+					married: true
+				}
+			]
+			FOR i IN users
+				COLLECT gender = i.gender WITH COUNT INTO numberOfUsers
+				RETURN {
+					gender,
+					values: numberOfUsers
+				}
+		`, []any{
+				map[string]any{
+					"gender": "f",
+					"values": 2,
+				},
+				map[string]any{
+					"gender": "m",
+					"values": 3,
+				},
+			}, "Should group and count result by a single key"),
 	})
 }
 

pkg/compiler/visitor.go

@@ -401,18 +401,6 @@ func (v *visitor) VisitLimitClause(ctx *fql.LimitClauseContext) interface{} {
 
 func (v *visitor) VisitCollectClause(ctx *fql.CollectClauseContext) interface{} {
 	if c := ctx.CollectGrouping(); c != nil {
-
-		if cvar := ctx.CollectGroupVariable(); cvar != nil {
-			return v.visitCollectGrouping(c.(*fql.CollectGroupingContext), cvar.(*fql.CollectGroupVariableContext))
-		}
-		// Collect by grouping
-		return v.visitCollectGrouping(c.(*fql.CollectGroupingContext), nil)
-	}
-
-	return nil
-}
-
-func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fql.CollectGroupVariableContext) interface{} {
 		// TODO: Undefine original loop variables
 		loop := v.loops.Loop()
 
@@ -421,7 +409,7 @@ func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fq
 		// Where a key is either a single value or a list of values
 		// These KeyValuePairs are then added to the dataset
 		var kvKeyReg vm.Operand
-	selectors := ctx.AllCollectSelector()
+		selectors := c.AllCollectSelector()
 		isMultiSelector := len(selectors) > 1
 
 		if isMultiSelector {
@@ -453,14 +441,17 @@ func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fq
 
 		var projectionVariableName string
 
-	if cvar != nil {
-		if identifier := cvar.Identifier(); identifier != nil {
-			projectionVariableName = v.emitDefaultCollectProjection(loop, kvValReg, identifier, cvar.CollectGroupVariableKeeper())
-		} else if selector := cvar.CollectSelector(); selector != nil {
-			projectionVariableName = v.emitCustomCollectProjection(loop, kvValReg, selector)
+		if groupVar := ctx.CollectGroupVariable(); groupVar != nil {
+			if identifier := groupVar.Identifier(); identifier != nil {
+				projectionVariableName = v.emitDefaultCollectGroupProjection(loop, kvValReg, identifier, groupVar.CollectGroupVariableKeeper())
+			} else if selector := groupVar.CollectSelector(); selector != nil {
+				projectionVariableName = v.emitCustomCollectGroupProjection(loop, kvValReg, selector)
 			}
 
 			v.emitter.EmitABC(vm.OpCollectKV, loop.Result, kvKeyReg, kvValReg)
+		} else if countVar := ctx.CollectCounter(); countVar != nil {
+			projectionVariableName = v.emitCollectCountProjection(loop, kvValReg, countVar)
+			v.emitter.EmitABC(vm.OpCollectKc, loop.Result, kvKeyReg, kvValReg)
 		} else {
 			v.emitter.EmitABC(vm.OpCollectK, loop.Result, kvKeyReg, kvValReg)
 		}
@@ -527,10 +518,13 @@ func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fq
 			v.emitter.EmitAB(vm.OpMove, varReg, kvValReg)
 		}
 
+		return nil
+	}
+
 	return nil
 }
 
-func (v *visitor) emitDefaultCollectProjection(loop *Loop, kvValReg vm.Operand, identifier antlr.TerminalNode, keeper fql.ICollectGroupVariableKeeperContext) string {
+func (v *visitor) emitDefaultCollectGroupProjection(loop *Loop, kvValReg vm.Operand, identifier antlr.TerminalNode, keeper fql.ICollectGroupVariableKeeperContext) string {
 	if keeper == nil {
 		seq := v.registers.AllocateSequence(2) // Key and Value for Map
 
@@ -558,8 +552,7 @@ func (v *visitor) emitDefaultCollectProjection(loop *Loop, kvValReg vm.Operand,
 	return identifier.GetText()
 }
 
-func (v *visitor) emitCustomCollectProjection(_ *Loop, kvValReg vm.Operand, selector fql.ICollectSelectorContext) string {
-	selector.Identifier().GetText()
+func (v *visitor) emitCustomCollectGroupProjection(_ *Loop, kvValReg vm.Operand, selector fql.ICollectSelectorContext) string {
 	selectorReg := selector.Expression().Accept(v).(vm.Operand)
 	v.emitter.EmitAB(vm.OpMove, kvValReg, selectorReg)
 	v.registers.Free(selectorReg)
@@ -567,6 +560,14 @@ func (v *visitor) emitCustomCollectProjection(_ *Loop, kvValReg vm.Operand, sele
 	return selector.Identifier().GetText()
 }
 
+func (v *visitor) emitCollectCountProjection(_ *Loop, kvValReg vm.Operand, selector fql.ICollectCounterContext) string {
+	//selectorReg := selector.Expression().Accept(v).(vm.Operand)
+	//v.emitter.EmitAB(vm.OpMove, kvValReg, selectorReg)
+	//v.registers.Free(selectorReg)
+
+	return selector.Identifier().GetText()
+}
+
 func (v *visitor) VisitCollectSelector(ctx *fql.CollectSelectorContext) interface{} {
 	if c := ctx.Expression(); c != nil {
 		return c.Accept(v)

pkg/vm/internal/dataset.go

@@ -103,7 +103,7 @@ func (ds *DataSet) AddKV(ctx context.Context, key, value runtime.Value) error {
 	return nil
 }
 
-func (ds *DataSet) CollectKey(ctx context.Context, key runtime.Value) error {
+func (ds *DataSet) CollectK(ctx context.Context, key runtime.Value) error {
 	k, err := Stringify(ctx, key)
 
 	if err != nil {
@@ -126,6 +126,38 @@ func (ds *DataSet) CollectKey(ctx context.Context, key runtime.Value) error {
 	return nil
 }
 
+func (ds *DataSet) CollectKc(ctx context.Context, key runtime.Value) error {
+	k, err := Stringify(ctx, key)
+
+	if err != nil {
+		return err
+	}
+
+	if ds.grouping == nil {
+		ds.grouping = make(map[string]runtime.Value)
+	}
+
+	group, exists := ds.grouping[k]
+
+	if !exists {
+		group = NewKV(key, runtime.ZeroInt)
+		ds.grouping[k] = group
+		_ = ds.values.Add(ctx, group)
+	}
+
+	kv := group.(*KV)
+	if count, ok := kv.Value.(runtime.Int); ok {
+		sum := count + 1
+		kv.Value = sum
+	} else {
+		kv.Value = runtime.NewInt(1)
+	}
+
+	ds.keyed = true
+
+	return nil
+}
+
 func (ds *DataSet) CollectKV(ctx context.Context, key, value runtime.Value) error {
 	k, err := Stringify(ctx, key)
 

pkg/vm/opcode.go

@@ -79,6 +79,7 @@ const (
 	OpPush      // Adds a value to a dataset
 	OpPushKV    // Adds a key-value pair to a dataset
 	OpCollectK  // Adds a key to a group
+	OpCollectKc // Adds a key to a group and counts it
 	OpCollectKV // Adds a value to a group using key
 	OpLimit
 	OpSkip

pkg/vm/vm.go

@@ -387,7 +387,18 @@ loop:
 			ds := reg[dst].(*internal.DataSet)
 			key := reg[src1]
 
-			if err := ds.CollectKey(ctx, key); err != nil {
+			if err := ds.CollectK(ctx, key); err != nil {
+				if _, catch := tryCatch(vm.pc); catch {
+					continue
+				}
+
+				return nil, err
+			}
+		case OpCollectKc:
+			ds := reg[dst].(*internal.DataSet)
+			key := reg[src1]
+
+			if err := ds.CollectKc(ctx, key); err != nil {
 				if _, catch := tryCatch(vm.pc); catch {
 					continue
 				}