Feature/#5 collect keyword alt (#141)

Implemented COLLECT key word

pkg/compiler/compiler_collect_aggregate_test.go

@@ -0,0 +1,309 @@
+package compiler_test
+
+import (
+	"context"
+	"github.com/MontFerret/ferret/pkg/compiler"
+	"github.com/MontFerret/ferret/pkg/runtime"
+	. "github.com/smartystreets/goconvey/convey"
+	"testing"
+)
+
+func TestAggregate(t *testing.T) {
+	Convey("Should aggregate values without grouping", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+  				RETURN {
+    				minAge, 
+    				maxAge 
+  				}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"maxAge":69,"minAge":25}]`)
+	})
+
+	Convey("Should aggregate values without grouping with multiple arguments", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 AGGREGATE ages = UNION(u.age, u.age)
+  				RETURN ages
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[[31,25,36,69,45,31,25,36,69,45]]`)
+	})
+
+	Convey("Should aggregate values with grouping", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 ageGroup = FLOOR(u.age / 5) * 5 
+  				AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+  				RETURN {
+					ageGroup,
+    				minAge, 
+    				maxAge 
+  				}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"ageGroup":25,"maxAge":25,"minAge":25},{"ageGroup":30,"maxAge":31,"minAge":31},{"ageGroup":35,"maxAge":36,"minAge":36},{"ageGroup":45,"maxAge":45,"minAge":45},{"ageGroup":65,"maxAge":69,"minAge":69}]`)
+	})
+}
+
+func BenchmarkAggregate(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+  				RETURN {
+    				minAge, 
+    				maxAge 
+  				}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}
+
+func BenchmarkAggregate2(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 AGGREGATE ages = UNION(u.age, u.age)
+  				RETURN ages
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}
+
+func BenchmarkAggregate3(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 ageGroup = FLOOR(u.age / 5) * 5 
+  				AGGREGATE minAge = MIN(u.age), maxAge = MAX(u.age)
+  				RETURN {
+					ageGroup,
+    				minAge, 
+    				maxAge 
+  				}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}

pkg/compiler/compiler_collect_count_test.go

@@ -0,0 +1,105 @@
+package compiler_test
+
+import (
+	"context"
+	"github.com/MontFerret/ferret/pkg/compiler"
+	"github.com/MontFerret/ferret/pkg/runtime"
+	. "github.com/smartystreets/goconvey/convey"
+	"testing"
+)
+
+func TestCollectCount(t *testing.T) {
+	Convey("Should count grouped values", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 WITH COUNT INTO c
+				RETURN c
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[5]`)
+	})
+}
+
+func BenchmarkCollectCount(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 WITH COUNT INTO c
+				RETURN c
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}

pkg/compiler/compiler_collect_into_test.go

@@ -0,0 +1,265 @@
+package compiler_test
+
+import (
+	"context"
+	"github.com/MontFerret/ferret/pkg/compiler"
+	"github.com/MontFerret/ferret/pkg/runtime"
+	. "github.com/smartystreets/goconvey/convey"
+	"testing"
+)
+
+func TestCollectInto(t *testing.T) {
+	Convey("Should create default projection", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 INTO genders
+				RETURN {
+					gender,
+					values: genders
+				}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"gender":"f","values":[{"i":{"active":true,"age":25,"gender":"f","married":false}},{"i":{"active":true,"age":45,"gender":"f","married":true}}]},{"gender":"m","values":[{"i":{"active":true,"age":31,"gender":"m","married":true}},{"i":{"active":true,"age":36,"gender":"m","married":false}},{"i":{"active":false,"age":69,"gender":"m","married":true}}]}]`)
+	})
+
+	Convey("Should create custom projection", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 INTO genders = { active: i.active }
+				RETURN {
+					gender,
+					values: genders
+				}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"gender":"f","values":[{"active":true},{"active":true}]},{"gender":"m","values":[{"active":true},{"active":true},{"active":false}]}]`)
+	})
+
+	Convey("Should create custom projection grouped by miltiple keys", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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, age = i.age INTO genders = { active: i.active }
+				RETURN {
+					age,
+					gender,
+					values: genders
+				}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"age":25,"gender":"f","values":[{"active":true}]},{"age":45,"gender":"f","values":[{"active":true}]},{"age":31,"gender":"m","values":[{"active":true}]},{"age":36,"gender":"m","values":[{"active":true}]},{"age":69,"gender":"m","values":[{"active":false}]}]`)
+	})
+}
+
+func BenchmarkCollectInto(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 INTO genders
+				RETURN {
+					gender,
+					values: genders
+				}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}
+
+func BenchmarkCollectInto2(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 INTO genders = { active: i.active }
+				RETURN {
+					gender,
+					values: genders
+				}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}

pkg/compiler/compiler_collect_test.go

@@ -0,0 +1,246 @@
+package compiler_test
+
+import (
+	"context"
+	"github.com/MontFerret/ferret/pkg/compiler"
+	"github.com/MontFerret/ferret/pkg/runtime"
+	. "github.com/smartystreets/goconvey/convey"
+	"testing"
+)
+
+func TestCollect(t *testing.T) {
+	Convey("Should not have access to initial variables", t, func() {
+		c := compiler.New()
+
+		_, err := c.Compile(`
+			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 users
+				COLLECT gender = i.gender
+				RETURN {
+					user: i,
+					gender: gender
+				}
+		`)
+
+		So(err, ShouldNotBeNil)
+	})
+
+	Convey("Should group result by a single key", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 users
+				COLLECT gender = i.gender
+				RETURN gender
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `["f","m"]`)
+	})
+
+	Convey("Should group result by multiple keys", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 users
+				COLLECT gender = i.gender, age = i.age
+				RETURN {age, gender}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"age":25,"gender":"f"},{"age":45,"gender":"f"},{"age":31,"gender":"m"},{"age":36,"gender":"m"},{"age":69,"gender":"m"}]`)
+	})
+}
+
+func BenchmarkCollect(b *testing.B) {
+	p := compiler.New().MustCompile(`
+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 users
+				COLLECT gender = i.gender
+				RETURN gender
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}
+
+func BenchmarkCollect2(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 users
+				COLLECT gender = i.gender, age = i.age
+				RETURN {age, gender}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}

pkg/compiler/compiler_collect_with_count_test.go

@@ -0,0 +1,199 @@
+package compiler_test
+
+import (
+	"context"
+	"github.com/MontFerret/ferret/pkg/compiler"
+	"github.com/MontFerret/ferret/pkg/runtime"
+	. "github.com/smartystreets/goconvey/convey"
+	"testing"
+)
+
+func TestCollectWithCount(t *testing.T) {
+	Convey("Should count grouped values", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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 genders
+				RETURN {gender, genders}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"gender":"f","genders":2},{"gender":"m","genders":3}]`)
+	})
+
+	Convey("Should count grouped values with multiple keys", t, func() {
+		c := compiler.New()
+
+		prog, err := c.Compile(`
+			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, married = i.married WITH COUNT INTO genders
+				RETURN {gender, married, genders}
+		`)
+
+		So(err, ShouldBeNil)
+		So(prog, ShouldHaveSameTypeAs, &runtime.Program{})
+
+		out, err := prog.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[{"gender":"f","genders":1,"married":false},{"gender":"f","genders":1,"married":true},{"gender":"m","genders":1,"married":false},{"gender":"m","genders":2,"married":true}]`)
+	})
+}
+
+func BenchmarkWithCount(b *testing.B) {
+	p := compiler.New().MustCompile(`
+			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 genders
+				RETURN {gender, genders}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}
+
+func BenchmarkWithCount2(b *testing.B) {
+	p := compiler.New().MustCompile(`
+						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, married = i.married WITH COUNT INTO genders
+				RETURN {gender, married, genders}
+		`)
+
+	for n := 0; n < b.N; n++ {
+		p.Run(context.Background())
+	}
+}

pkg/compiler/compiler_param_test.go

@@ -47,17 +47,18 @@ func TestParam(t *testing.T) {
 	})
 
 	Convey("Should be possible to use in range", t, func() {
-		out := compiler.New().
+		prog := compiler.New().
 			MustCompile(`
 			FOR i IN @start..@end
 			SORT i
 			RETURN i
-		`).
-			MustRun(
-				context.Background(),
-				runtime.WithParam("start", 1),
-				runtime.WithParam("end", 4),
-			)
+		`)
+
+		out := prog.MustRun(
+			context.Background(),
+			runtime.WithParam("start", 1),
+			runtime.WithParam("end", 4),
+		)
 
 		So(string(out), ShouldEqual, `[1,2,3,4]`)
 

pkg/compiler/scope.go

@@ -58,6 +58,18 @@ func (s *scope) SetVariable(name string) error {
 	return nil
 }
 
+func (s *scope) RemoveVariable(name string) error {
+	_, exists := s.vars[name]
+
+	if !exists {
+		return errors.Wrap(ErrVariableNotFound, name)
+	}
+
+	delete(s.vars, name)
+
+	return nil
+}
+
 func (s *scope) Fork() *scope {
 	return newScope(s)
 }

pkg/compiler/visitor.go

@@ -16,11 +16,15 @@ import (
 	"strings"
 )
 
-type visitor struct {
-	*fql.BaseFqlParserVisitor
-	src   string
-	funcs map[string]core.Function
-}
+type (
+	forOption func(f *expressions.ForExpression) error
+
+	visitor struct {
+		*fql.BaseFqlParserVisitor
+		src   string
+		funcs map[string]core.Function
+	}
+)
 
 func newVisitor(src string, funcs map[string]core.Function) *visitor {
 	return &visitor{
@@ -155,6 +159,7 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 	var valVarName string
 	var keyVarName string
 
+	parsedClauses := make([]forOption, 0, 10)
 	valVar := ctx.ForExpressionValueVariable()
 	valVarName = valVar.GetText()
 	forInScope := scope.Fork()
@@ -167,12 +172,95 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 		forInScope.SetVariable(keyVarName)
 	}
 
-	src, err := v.doVisitForExpressionSource(ctx.ForExpressionSource().(*fql.ForExpressionSourceContext), forInScope)
+	srcCtx := ctx.ForExpressionSource().(*fql.ForExpressionSourceContext)
+	srcExp, err := v.doVisitForExpressionSource(srcCtx, forInScope)
 
 	if err != nil {
 		return nil, err
 	}
 
+	src, err := expressions.NewDataSource(
+		v.getSourceMap(srcCtx),
+		valVarName,
+		keyVarName,
+		srcExp,
+	)
+
+	// Clauses.
+	// We put clauses parsing before parsing the query body because COLLECT clause overrides scope variables
+	for _, clause := range ctx.AllForExpressionClause() {
+		clause := clause.(*fql.ForExpressionClauseContext)
+
+		limitCtx := clause.LimitClause()
+
+		if limitCtx != nil {
+			limit, offset, err := v.createLimit(limitCtx.(*fql.LimitClauseContext))
+
+			if err != nil {
+				return nil, err
+			}
+
+			parsedClauses = append(parsedClauses, func(f *expressions.ForExpression) error {
+				return f.AddLimit(v.getSourceMap(limitCtx), limit, offset)
+			})
+
+			continue
+		}
+
+		filterCtx := clause.FilterClause()
+
+		if filterCtx != nil {
+			filterExp, err := v.createFilter(filterCtx.(*fql.FilterClauseContext), forInScope)
+
+			if err != nil {
+				return nil, err
+			}
+
+			parsedClauses = append(parsedClauses, func(f *expressions.ForExpression) error {
+				return f.AddFilter(v.getSourceMap(filterCtx), filterExp)
+			})
+
+			continue
+		}
+
+		sortCtx := clause.SortClause()
+
+		if sortCtx != nil {
+			sortCtx := sortCtx.(*fql.SortClauseContext)
+			sortExps, err := v.createSort(sortCtx, forInScope)
+
+			if err != nil {
+				return nil, err
+			}
+
+			parsedClauses = append(parsedClauses, func(f *expressions.ForExpression) error {
+				return f.AddSort(v.getSourceMap(sortCtx), sortExps...)
+			})
+		}
+
+		collectCtx := clause.CollectClause()
+
+		if collectCtx != nil {
+			collectCtx := collectCtx.(*fql.CollectClauseContext)
+
+			params, err := v.createCollect(collectCtx, forInScope, valVarName)
+
+			if err != nil {
+				return nil, err
+			}
+
+			forInScope.RemoveVariable(valVarName)
+
+			if keyVarName != "" {
+				forInScope.RemoveVariable(keyVarName)
+			}
+
+			parsedClauses = append(parsedClauses, func(f *expressions.ForExpression) error {
+				return f.AddCollect(v.getSourceMap(collectCtx), params)
+			})
+		}
+	}
+
 	body := ctx.AllForExpressionBody()
 	predicate := expressions.NewBlockExpression(len(body) + 1)
 
@@ -192,7 +280,6 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 
 	var spread bool
 	var distinct bool
-	var distinctSrc core.SourceMap
 	forRetCtx := ctx.ForExpressionReturn().(*fql.ForExpressionReturnContext)
 	returnCtx := forRetCtx.ReturnExpression()
 
@@ -208,8 +295,6 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 
 		if distinctCtx != nil {
 			distinct = true
-			token := distinctCtx.GetSymbol()
-			distinctSrc = core.NewSourceMap(token.GetText(), token.GetLine(), token.GetColumn())
 		}
 
 		predicate.Add(returnExp)
@@ -228,10 +313,9 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 
 	forExp, err := expressions.NewForExpression(
 		v.getSourceMap(ctx),
-		valVarName,
-		keyVarName,
 		src,
 		predicate,
+		distinct,
 		spread,
 	)
 
@@ -239,52 +323,10 @@ func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *sco
 		return nil, err
 	}
 
-	if distinct {
-		forExp.AddDistinct(distinctSrc)
-	}
-
-	for _, clause := range ctx.AllForExpressionClause() {
-		clause := clause.(*fql.ForExpressionClauseContext)
-
-		limitCtx := clause.LimitClause()
-
-		if limitCtx != nil {
-			limit, offset, err := v.createLimit(limitCtx.(*fql.LimitClauseContext))
-
-			if err != nil {
-				return nil, err
-			}
-
-			forExp.AddLimit(v.getSourceMap(limitCtx), limit, offset)
-
-			continue
-		}
-
-		filterCtx := clause.FilterClause()
-
-		if filterCtx != nil {
-			filterExp, err := v.createFilter(filterCtx.(*fql.FilterClauseContext), forInScope)
-
-			if err != nil {
-				return nil, err
-			}
-
-			forExp.AddFilter(v.getSourceMap(filterCtx), filterExp)
-
-			continue
-		}
-
-		sortCtx := clause.SortClause()
-
-		if sortCtx != nil {
-			sortCtx := sortCtx.(*fql.SortClauseContext)
-			sortExps, err := v.createSort(sortCtx, forInScope)
-
-			if err != nil {
-				return nil, err
-			}
-
-			forExp.AddSort(v.getSourceMap(sortCtx), sortExps...)
+	// add all available clauses
+	for _, clause := range parsedClauses {
+		if err := clause(forExp); err != nil {
+			return nil, err
 		}
 	}
 
@@ -392,7 +434,190 @@ func (v *visitor) createSort(ctx *fql.SortClauseContext, scope *scope) ([]*claus
 	return res, nil
 }
 
-func (v *visitor) doVisitForExpressionSource(ctx *fql.ForExpressionSourceContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) createCollect(ctx *fql.CollectClauseContext, scope *scope, valVarName string) (*clauses.Collect, error) {
+	var err error
+	var selectors []*clauses.CollectSelector
+	var projection *clauses.CollectProjection
+	var count *clauses.CollectCount
+	var aggregate *clauses.CollectAggregate
+
+	groupingCtx := ctx.CollectGrouping()
+
+	if groupingCtx != nil {
+		groupingCtx := groupingCtx.(*fql.CollectGroupingContext)
+		collectSelectors := groupingCtx.AllCollectSelector()
+
+		// group selectors
+		if collectSelectors != nil && len(collectSelectors) > 0 {
+			selectors = make([]*clauses.CollectSelector, 0, len(collectSelectors))
+
+			for _, cs := range collectSelectors {
+				selector, err := v.createCollectSelector(cs.(*fql.CollectSelectorContext), scope)
+
+				if err != nil {
+					return nil, err
+				}
+
+				selectors = append(selectors, selector)
+
+				if err := scope.SetVariable(selector.Variable()); err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		projectionCtx := ctx.CollectGroupVariable()
+
+		if projectionCtx != nil {
+			projectionCtx := projectionCtx.(*fql.CollectGroupVariableContext)
+			projectionSelectorCtx := projectionCtx.CollectSelector()
+			var projectionSelector *clauses.CollectSelector
+
+			// if projection expression is defined like WITH group = { foo: i.bar }
+			if projectionSelectorCtx != nil {
+				selector, err := v.createCollectSelector(projectionSelectorCtx.(*fql.CollectSelectorContext), scope)
+
+				if err != nil {
+					return nil, err
+				}
+
+				projectionSelector = selector
+			} else {
+				// otherwise, use default expression WITH group = { i }
+				projectionIdentifier := projectionCtx.Identifier(0)
+
+				if projectionIdentifier != nil {
+					varExp, err := expressions.NewVariableExpression(v.getSourceMap(projectionCtx), valVarName)
+
+					if err != nil {
+						return nil, err
+					}
+
+					strLitExp := literals.NewStringLiteral(valVarName)
+
+					propExp, err := literals.NewObjectPropertyAssignment(
+						strLitExp,
+						varExp,
+					)
+
+					if err != nil {
+						return nil, err
+					}
+
+					projectionSelectorExp := literals.NewObjectLiteralWith(propExp)
+
+					if err != nil {
+						return nil, err
+					}
+
+					selector, err := clauses.NewCollectSelector(projectionIdentifier.GetText(), projectionSelectorExp)
+
+					if err != nil {
+						return nil, err
+					}
+
+					projectionSelector = selector
+				}
+			}
+
+			if projectionSelector != nil {
+				if err := scope.SetVariable(projectionSelector.Variable()); err != nil {
+					return nil, err
+				}
+
+				projection, err = clauses.NewCollectProjection(projectionSelector)
+
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+	}
+
+	countCtx := ctx.CollectCounter()
+
+	if countCtx != nil {
+		countCtx := countCtx.(*fql.CollectCounterContext)
+		variable := countCtx.Identifier().GetText()
+
+		if err := scope.SetVariable(variable); err != nil {
+			return nil, err
+		}
+
+		count, err = clauses.NewCollectCount(variable)
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	aggrCtx := ctx.CollectAggregator()
+
+	if aggrCtx != nil {
+		aggrCtx := aggrCtx.(*fql.CollectAggregatorContext)
+
+		selectorCtxs := aggrCtx.AllCollectAggregateSelector()
+		selectors := make([]*clauses.CollectAggregateSelector, 0, len(selectorCtxs))
+
+		for _, sc := range selectorCtxs {
+			selector, err := v.createCollectAggregateSelector(sc.(*fql.CollectAggregateSelectorContext), scope)
+
+			if err != nil {
+				return nil, err
+			}
+
+			selectors = append(selectors, selector)
+		}
+
+		aggregate, err = clauses.NewCollectAggregate(selectors)
+
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return clauses.NewCollect(selectors, projection, count, aggregate)
+}
+
+func (v *visitor) createCollectSelector(ctx *fql.CollectSelectorContext, scope *scope) (*clauses.CollectSelector, error) {
+	variable := ctx.Identifier().GetText()
+	exp, err := v.doVisitExpression(ctx.Expression().(*fql.ExpressionContext), scope)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return clauses.NewCollectSelector(variable, exp)
+}
+
+func (v *visitor) createCollectAggregateSelector(ctx *fql.CollectAggregateSelectorContext, scope *scope) (*clauses.CollectAggregateSelector, error) {
+	variable := ctx.Identifier().GetText()
+	fnCtx := ctx.FunctionCallExpression()
+
+	if fnCtx != nil {
+		exp, err := v.doVisitFunctionCallExpression(fnCtx.(*fql.FunctionCallExpressionContext), scope)
+
+		if err != nil {
+			return nil, err
+		}
+
+		fnExp, ok := exp.(*expressions.FunctionCallExpression)
+
+		if !ok {
+			return nil, core.Error(core.ErrInvalidType, "expected function expression")
+		}
+
+		if err := scope.SetVariable(variable); err != nil {
+			return nil, err
+		}
+
+		return clauses.NewCollectAggregateSelector(variable, fnExp.Arguments(), fnExp.Function())
+	}
+
+	return nil, core.Error(core.ErrNotFound, "function expression")
+}
+
+func (v *visitor) doVisitForExpressionSource(ctx *fql.ForExpressionSourceContext, scope *scope) (core.Expression, error) {
 	arr := ctx.ArrayLiteral()
 
 	if arr != nil {
@@ -454,7 +679,7 @@ func (v *visitor) doVisitForExpressionBody(ctx *fql.ForExpressionBodyContext, sc
 	return nil, v.unexpectedToken(ctx)
 }
 
-func (v *visitor) doVisitMemberExpression(ctx *fql.MemberExpressionContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitMemberExpression(ctx *fql.MemberExpressionContext, scope *scope) (core.Expression, error) {
 	varName := ctx.Identifier().GetText()
 
 	_, err := scope.GetVariable(varName)
@@ -516,7 +741,7 @@ func (v *visitor) doVisitMemberExpression(ctx *fql.MemberExpressionContext, scop
 	return member, nil
 }
 
-func (v *visitor) doVisitObjectLiteral(ctx *fql.ObjectLiteralContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitObjectLiteral(ctx *fql.ObjectLiteralContext, scope *scope) (core.Expression, error) {
 	assignments := ctx.AllPropertyAssignment()
 	props := make([]*literals.ObjectPropertyAssignment, 0, len(assignments))
 
@@ -553,7 +778,13 @@ func (v *visitor) doVisitObjectLiteral(ctx *fql.ObjectLiteralContext, scope *sco
 			return nil, err
 		}
 
-		props = append(props, literals.NewObjectPropertyAssignment(name, value))
+		pa, err := literals.NewObjectPropertyAssignment(name, value)
+
+		if err != nil {
+			return nil, err
+		}
+
+		props = append(props, pa)
 	}
 
 	return literals.NewObjectLiteralWith(props...), nil
@@ -598,7 +829,7 @@ func (v *visitor) doVisitShorthandPropertyNameContext(ctx *fql.ShorthandProperty
 	return literals.NewStringLiteral(ctx.Variable().GetText()), nil
 }
 
-func (v *visitor) doVisitArrayLiteral(ctx *fql.ArrayLiteralContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitArrayLiteral(ctx *fql.ArrayLiteralContext, scope *scope) (core.Expression, error) {
 	listCtx := ctx.ArrayElementList()
 
 	if listCtx == nil {
@@ -657,7 +888,7 @@ func (v *visitor) doVisitNoneLiteral(_ *fql.NoneLiteralContext) (core.Expression
 	return literals.None, nil
 }
 
-func (v *visitor) doVisitVariable(ctx *fql.VariableContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitVariable(ctx *fql.VariableContext, scope *scope) (core.Expression, error) {
 	name := ctx.Identifier().GetText()
 
 	// check whether the variable is defined
@@ -711,7 +942,7 @@ func (v *visitor) doVisitVariableDeclaration(ctx *fql.VariableDeclarationContext
 	)
 }
 
-func (v *visitor) doVisitRangeOperator(ctx *fql.RangeOperatorContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitRangeOperator(ctx *fql.RangeOperatorContext, scope *scope) (core.Expression, error) {
 	exp, err := v.doVisitChildren(ctx, scope)
 
 	if err != nil {
@@ -732,7 +963,7 @@ func (v *visitor) doVisitRangeOperator(ctx *fql.RangeOperatorContext, scope *sco
 	)
 }
 
-func (v *visitor) doVisitFunctionCallExpression(context *fql.FunctionCallExpressionContext, scope *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitFunctionCallExpression(context *fql.FunctionCallExpressionContext, scope *scope) (core.Expression, error) {
 	args := make([]core.Expression, 0, 5)
 	argsCtx := context.Arguments()
 
@@ -765,7 +996,7 @@ func (v *visitor) doVisitFunctionCallExpression(context *fql.FunctionCallExpress
 	)
 }
 
-func (v *visitor) doVisitParamContext(context *fql.ParamContext, _ *scope) (collections.IterableExpression, error) {
+func (v *visitor) doVisitParamContext(context *fql.ParamContext, _ *scope) (core.Expression, error) {
 	name := context.Identifier().GetText()
 
 	return expressions.NewParameterExpression(