wip sorting

pkg/compiler/compiler_exec_test.go

@@ -1957,6 +1957,381 @@ LET users = [
 					},
 				},
 			}, "Should create custom projection grouped by multiple keys"),
+		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
+				LET married = i.married
+				COLLECT gender = i.gender INTO genders KEEP married
+				RETURN {
+					gender,
+					values: genders
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{"married": false},
+					map[string]any{"married": true},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{"married": true},
+					map[string]any{"married": false},
+					map[string]any{"married": true},
+				},
+			},
+		}, "Should create default projection with default KEEP"),
+		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
+				LET married = i.married
+				LET age = i.age
+				COLLECT gender = i.gender INTO values KEEP married, age
+				RETURN {
+					gender,
+					values
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{
+						"married": false,
+						"age":     25,
+					},
+					map[string]any{
+						"married": true,
+						"age":     45,
+					},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{
+						"married": true,
+						"age":     31,
+					},
+					map[string]any{
+						"married": false,
+						"age":     36,
+					},
+					map[string]any{
+						"married": true,
+						"age":     69,
+					},
+				},
+			},
+		}, "Should create default projection with default KEEP using multiple keys"),
+		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
+				LET married = "foo"
+				COLLECT gender = i.gender INTO values KEEP married
+				RETURN {
+					gender,
+					values
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{"married": "foo"},
+					map[string]any{"married": "foo"},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{"married": "foo"},
+					map[string]any{"married": "foo"},
+					map[string]any{"married": "foo"},
+				},
+			},
+		}, "Should create default projection with custom KEEP"),
+		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
+				LET married = "foo"
+				LET age = "bar"
+				COLLECT gender = i.gender INTO values KEEP married, age
+				RETURN {
+					gender,
+					values
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{
+						"married": "foo",
+						"age":     "bar",
+					},
+					map[string]any{
+						"married": "foo",
+						"age":     "bar",
+					},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{
+						"married": "foo",
+						"age":     "bar",
+					},
+					map[string]any{
+						"married": "foo",
+						"age":     "bar",
+					},
+					map[string]any{
+						"married": "foo",
+						"age":     "bar",
+					},
+				},
+			},
+		}, "Should create default projection with custom KEEP using multiple keys"),
+		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
+				LET bar = "foo"
+				COLLECT gender = i.gender INTO values KEEP bar
+				RETURN {
+					gender,
+					values
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{"bar": "foo"},
+					map[string]any{"bar": "foo"},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{"bar": "foo"},
+					map[string]any{"bar": "foo"},
+					map[string]any{"bar": "foo"},
+				},
+			},
+		}, "Should create default projection with custom KEEP with custom name"),
+		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
+				LET bar = "foo"
+				LET baz = "bar"
+				COLLECT gender = i.gender INTO values KEEP bar, baz
+				RETURN {
+					gender,
+					values
+				}
+`, []any{
+			map[string]any{
+				"gender": "f",
+				"values": []any{
+					map[string]any{"bar": "foo", "baz": "bar"},
+					map[string]any{"bar": "foo", "baz": "bar"},
+				},
+			},
+			map[string]any{
+				"gender": "m",
+				"values": []any{
+					map[string]any{"bar": "foo", "baz": "bar"},
+					map[string]any{"bar": "foo", "baz": "bar"},
+					map[string]any{"bar": "foo", "baz": "bar"},
+				},
+			},
+		}, "Should create default projection with custom KEEP with multiple custom names"),
 	})
 }
 

pkg/compiler/visitor.go

@@ -454,8 +454,8 @@ func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fq
 	var projectionVariableName string
 
 	if cvar != nil {
-		if identifiers := cvar.AllIdentifier(); len(identifiers) > 0 {
-			projectionVariableName = v.emitDefaultCollectProjection(loop, kvValReg, identifiers)
+		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)
 		}
@@ -530,18 +530,32 @@ func (v *visitor) visitCollectGrouping(ctx *fql.CollectGroupingContext, cvar *fq
 	return nil
 }
 
-func (v *visitor) emitDefaultCollectProjection(loop *Loop, kvValReg vm.Operand, identifiers []antlr.TerminalNode) string {
-	seq := v.registers.AllocateSequence(2) // Key and Value for Map
+func (v *visitor) emitDefaultCollectProjection(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
 
-	// TODO: Review this. It's quite a questionable ArrangoDB feature of wrapping group items by a nested object
-	// We will keep it for now for backward compatibility.
-	v.loadConstantTo(runtime.String(loop.ValueName), seq.Registers[0]) // Map key
-	v.emitter.EmitAB(vm.OpMove, seq.Registers[1], kvValReg)            // Map value
-	v.emitter.EmitAs(vm.OpLoadMap, kvValReg, seq)
+		// TODO: Review this. It's quite a questionable ArrangoDB feature of wrapping group items by a nested object
+		// We will keep it for now for backward compatibility.
+		v.loadConstantTo(runtime.String(loop.ValueName), seq.Registers[0]) // Map key
+		v.emitter.EmitAB(vm.OpMove, seq.Registers[1], kvValReg)            // Map value
+		v.emitter.EmitAs(vm.OpLoadMap, kvValReg, seq)
 
-	v.registers.FreeSequence(seq)
+		v.registers.FreeSequence(seq)
+	} else {
+		variables := keeper.AllIdentifier()
+		seq := v.registers.AllocateSequence(len(variables) * 2)
 
-	return identifiers[0].GetText()
+		for i, j := 0, 0; i < len(variables); i, j = i+1, j+2 {
+			varName := variables[i].GetText()
+			v.loadConstantTo(runtime.String(varName), seq.Registers[j])
+			v.emitter.EmitAB(vm.OpMove, seq.Registers[j+1], v.symbols.Variable(varName))
+		}
+
+		v.emitter.EmitAs(vm.OpLoadMap, kvValReg, seq)
+		v.registers.FreeSequence(seq)
+	}
+
+	return identifier.GetText()
 }
 
 func (v *visitor) emitCustomCollectProjection(_ *Loop, kvValReg vm.Operand, selector fql.ICollectSelectorContext) string {

pkg/parser/antlr/FqlParser.g4

@@ -136,7 +136,11 @@ collectAggregateSelector
 
 collectGroupVariable
     : Into collectSelector
-    | Into Identifier (Keep Identifier)?
+    | Into Identifier (collectGroupVariableKeeper)?
+    ;
+
+collectGroupVariableKeeper
+    : Keep Identifier (Comma Identifier)*
     ;
 
 collectCounter

pkg/parser/fql/fql_lexer.go

@@ -4,10 +4,9 @@ package fql
 
 import (
 	"fmt"
+	"github.com/antlr4-go/antlr/v4"
 	"sync"
 	"unicode"
-
-	"github.com/antlr4-go/antlr/v4"
 )
 
 // Suppress unused import error

pkg/parser/fql/fqlparser_base_listener.go

@@ -176,6 +176,14 @@ func (s *BaseFqlParserListener) EnterCollectGroupVariable(ctx *CollectGroupVaria
 // ExitCollectGroupVariable is called when production collectGroupVariable is exited.
 func (s *BaseFqlParserListener) ExitCollectGroupVariable(ctx *CollectGroupVariableContext) {}
 
+// EnterCollectGroupVariableKeeper is called when production collectGroupVariableKeeper is entered.
+func (s *BaseFqlParserListener) EnterCollectGroupVariableKeeper(ctx *CollectGroupVariableKeeperContext) {
+}
+
+// ExitCollectGroupVariableKeeper is called when production collectGroupVariableKeeper is exited.
+func (s *BaseFqlParserListener) ExitCollectGroupVariableKeeper(ctx *CollectGroupVariableKeeperContext) {
+}
+
 // EnterCollectCounter is called when production collectCounter is entered.
 func (s *BaseFqlParserListener) EnterCollectCounter(ctx *CollectCounterContext) {}
 

pkg/parser/fql/fqlparser_base_visitor.go

@@ -111,6 +111,10 @@ func (v *BaseFqlParserVisitor) VisitCollectGroupVariable(ctx *CollectGroupVariab
 	return v.VisitChildren(ctx)
 }
 
+func (v *BaseFqlParserVisitor) VisitCollectGroupVariableKeeper(ctx *CollectGroupVariableKeeperContext) interface{} {
+	return v.VisitChildren(ctx)
+}
+
 func (v *BaseFqlParserVisitor) VisitCollectCounter(ctx *CollectCounterContext) interface{} {
 	return v.VisitChildren(ctx)
 }

pkg/parser/fql/fqlparser_listener.go

@@ -85,6 +85,9 @@ type FqlParserListener interface {
 	// EnterCollectGroupVariable is called when entering the collectGroupVariable production.
 	EnterCollectGroupVariable(c *CollectGroupVariableContext)
 
+	// EnterCollectGroupVariableKeeper is called when entering the collectGroupVariableKeeper production.
+	EnterCollectGroupVariableKeeper(c *CollectGroupVariableKeeperContext)
+
 	// EnterCollectCounter is called when entering the collectCounter production.
 	EnterCollectCounter(c *CollectCounterContext)
 
@@ -301,6 +304,9 @@ type FqlParserListener interface {
 	// ExitCollectGroupVariable is called when exiting the collectGroupVariable production.
 	ExitCollectGroupVariable(c *CollectGroupVariableContext)
 
+	// ExitCollectGroupVariableKeeper is called when exiting the collectGroupVariableKeeper production.
+	ExitCollectGroupVariableKeeper(c *CollectGroupVariableKeeperContext)
+
 	// ExitCollectCounter is called when exiting the collectCounter production.
 	ExitCollectCounter(c *CollectCounterContext)
 

pkg/parser/fql/fqlparser_visitor.go

@@ -85,6 +85,9 @@ type FqlParserVisitor interface {
 	// Visit a parse tree produced by FqlParser#collectGroupVariable.
 	VisitCollectGroupVariable(ctx *CollectGroupVariableContext) interface{}
 
+	// Visit a parse tree produced by FqlParser#collectGroupVariableKeeper.
+	VisitCollectGroupVariableKeeper(ctx *CollectGroupVariableKeeperContext) interface{}
+
 	// Visit a parse tree produced by FqlParser#collectCounter.
 	VisitCollectCounter(ctx *CollectCounterContext) interface{}
 

pkg/vm/opcode.go

@@ -12,6 +12,9 @@ const (
 	OpLoadGlobal                  // Load a global variable to a register A
 	OpStoreGlobal                 // Store a value from register A to a global variable
 	OpLoadParam                   // Load a parameter to a register A
+	OpLoadList                    // Load an array from a list of registers (ARR R2, R3 R5 - creates an array in R2 with elements from R3 to R5)
+	OpLoadMap                     // Load an object from a list of registers (OBJ R2, R3 R5 - creates an object in R2 with elements from R3 to R5)
+	OpLoadRange                   // Load a range from a list of registers (RNG R2, R3, R4 - creates a range in R2 with start from R3 and end at R4)
 	OpLoadIndex                   // Load a value from a list to a register (INDEX R1, R2, R3 - loads a value from a list in R2 to R1)
 	OpLoadIndexOptional           // Load a value from a list to a register, if it exists
 	OpLoadKey                     // Load a value from a map to a register (KEY R1, R2, R3 - loads a value from a map in R2 to R1)
@@ -19,9 +22,6 @@ const (
 	OpLoadProperty                // Load a property (key or index) from an object (map or list) to a register
 	OpLoadPropertyOptional        // Load a property (key or index) from an object (map or list) to a register, if it exists
 	OpLoadDataSet                 // Load a dataset to a register A
-	OpLoadList                    // Load an array from a list of registers (ARR R2, R3 R5 - creates an array in R2 with elements from R3 to R5)
-	OpLoadMap                     // Load an object from a list of registers (OBJ R2, R3 R5 - creates an object in R2 with elements from R3 to R5)
-	OpLoadRange
 
 	OpJump
 	OpJumpIfFalse