Add unit tests that exercise rules; fix bugs found along the way (#140)

Author: Christopher M. Wolff

functions/transformations/range.go

@@ -128,6 +128,9 @@ func (s *RangeProcedureSpec) Kind() plan.ProcedureKind {
 func (s *RangeProcedureSpec) Copy() plan.ProcedureSpec {
 	ns := new(RangeProcedureSpec)
 	ns.Bounds = s.Bounds
+	ns.TimeCol = s.TimeCol
+	ns.StartCol = s.StartCol
+	ns.StopCol = s.StopCol
 	return ns
 }
 

planner/heuristic_planner.go

@@ -31,9 +31,11 @@ func (p *heuristicPlanner) matchRules(node PlanNode) (PlanNode, bool) {
 	}
 
 	for _, rule := range p.rules[node.Kind()] {
-		newNode, changed := rule.Rewrite(node)
-		anyChanged = anyChanged || changed
-		node = newNode
+		if rule.Pattern().Match(node) {
+			newNode, changed := rule.Rewrite(node)
+			anyChanged = anyChanged || changed
+			node = newNode
+		}
 	}
 
 	return node, anyChanged
@@ -103,12 +105,15 @@ func updateSuccessors(plan *PlanSpec, oldNode, newNode PlanNode) {
 	}
 
 	for _, succ := range oldNode.Successors() {
-		i := 0
-		for ; i < len(succ.Predecessors()); i++ {
-			succ.Predecessors()[i] = newNode
+		found := false
+		for i, succPred := range succ.Predecessors() {
+			if succPred == oldNode {
+				found = true
+				succ.Predecessors()[i] = newNode
+			}
 		}
 
-		if i == len(succ.Predecessors()) {
+		if !found {
 			panic("Inconsistent plan graph: successor does not have edge back to predecessor")
 		}
 	}

planner/logical.go

@@ -107,6 +107,14 @@ func (lpn *LogicalPlanNode) ProcedureSpec() ProcedureSpec {
 	return lpn.Spec
 }
 
+func (lpn *LogicalPlanNode) ShallowCopy() PlanNode {
+	newNode := new(LogicalPlanNode)
+	newNode.edges = lpn.edges.shallowCopy()
+	newNode.id = lpn.id + "_copy"
+	newNode.Spec = lpn.Spec.Copy()
+	return newNode
+}
+
 // createLogicalPlan creates a logical query plan from a flux spec
 func createLogicalPlan(spec *flux.Spec, a Administration) (*PlanSpec, error) {
 	nodes := make(map[flux.OperationID]PlanNode, len(spec.Operations))

planner/logical_test.go

@@ -6,13 +6,19 @@ import (
 	"time"
 
 	"github.com/influxdata/flux"
+	"github.com/influxdata/flux/ast"
 	"github.com/influxdata/flux/functions/inputs"
 	"github.com/influxdata/flux/functions/transformations"
 	"github.com/influxdata/flux/planner"
 	"github.com/influxdata/flux/planner/plantest"
 	"github.com/influxdata/flux/semantic"
 )
 
+func compile(fluxText string) (*flux.Spec, error) {
+	now := time.Now().UTC()
+	return flux.Compile(context.Background(), fluxText, now)
+}
+
 // Test the translation of Flux query to logical plan
 func TestFluxSpecToLogicalPlan(t *testing.T) {
 	testcases := []struct {
@@ -99,14 +105,12 @@ func TestFluxSpecToLogicalPlan(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			now := time.Now().UTC()
-			spec, err := flux.Compile(context.Background(), tc.query, now)
+			spec, err := compile(tc.query)
 
 			if err != nil {
 				t.Fatal(err)
 			}
 
-			tc.spec.Now = now
 			want := plantest.CreateLogicalPlanSpec(tc.spec)
 
 			thePlanner := planner.NewLogicalPlanner()
@@ -125,3 +129,197 @@ func TestFluxSpecToLogicalPlan(t *testing.T) {
 		})
 	}
 }
+
+type MergeFiltersRule struct {
+}
+
+func (MergeFiltersRule) Name() string {
+	return "mergeFilters"
+}
+
+func (MergeFiltersRule) Pattern() planner.Pattern {
+	return planner.Pat(transformations.FilterKind,
+		planner.Pat(transformations.FilterKind,
+			planner.Any()))
+}
+
+func (MergeFiltersRule) Rewrite(pn planner.PlanNode) (planner.PlanNode, bool) {
+	specTop := pn.ProcedureSpec()
+
+	filterSpecTop := specTop.(*transformations.FilterProcedureSpec)
+	filterSpecBottom := pn.Predecessors()[0].ProcedureSpec().(*transformations.FilterProcedureSpec)
+	mergedFilterSpec := mergeFilterSpecs(filterSpecTop, filterSpecBottom)
+
+	return planner.MergeLogicalPlanNodes(pn, pn.Predecessors()[0], mergedFilterSpec), true
+}
+
+func mergeFilterSpecs(a, b *transformations.FilterProcedureSpec) planner.ProcedureSpec {
+	fn := a.Fn.Copy().(*semantic.FunctionExpression)
+
+	aExp, aOK := a.Fn.Body.(semantic.Expression)
+	bExp, bOK := b.Fn.Body.(semantic.Expression)
+
+	if !aOK || !bOK {
+		// Note that this is just a unit test, so "return" statements are not handled.
+		panic("function body not expression")
+	}
+
+	fn.Body = &semantic.LogicalExpression{
+		Operator: ast.AndOperator,
+		Left:     aExp,
+		Right:    bExp,
+	}
+
+	return &transformations.FilterProcedureSpec{
+		Fn: fn,
+	}
+}
+
+type PushFilterThroughMapRule struct {
+}
+
+func (PushFilterThroughMapRule) Name() string {
+	return "pushFilterThroughMap"
+}
+
+func (PushFilterThroughMapRule) Pattern() planner.Pattern {
+	return planner.Pat(transformations.FilterKind,
+		planner.Pat(transformations.MapKind,
+			planner.Any()))
+}
+
+func (PushFilterThroughMapRule) Rewrite(pn planner.PlanNode) (planner.PlanNode, bool) {
+	// It will not always be possible to push a filter through a map... but this is just a unit test.
+	return planner.SwapPlanNodes(pn, pn.Predecessors()[0]), true
+}
+
+func init() {
+	planner.RegisterLogicalRule(MergeFiltersRule{})
+	planner.RegisterLogicalRule(PushFilterThroughMapRule{})
+}
+
+func TestLogicalPlanner(t *testing.T) {
+	testcases := []struct {
+		name     string
+		flux     string
+		wantPlan plantest.LogicalPlanSpec
+	}{{
+		name: "with merge-able filters",
+		flux: `
+			from(bucket: "telegraf") |>
+				filter(fn: (r) => r._measurement == "cpu") |>
+				filter(fn: (r) => r._value > 0.5) |>
+				filter(fn: (r) => r._value < 0.9) |>
+				yield(name: "result")`,
+		wantPlan: plantest.LogicalPlanSpec{
+			Nodes: []planner.PlanNode{
+				planner.CreateLogicalNode("from0", &inputs.FromProcedureSpec{Bucket: "telegraf"}),
+				planner.CreateLogicalNode("merged_filter1_merged_filter2_filter3", &transformations.FilterProcedureSpec{Fn: &semantic.FunctionExpression{
+					Params: []*semantic.FunctionParam{{Key: &semantic.Identifier{Name: "r"}}},
+					Body: &semantic.LogicalExpression{Operator: ast.AndOperator,
+						Left: &semantic.LogicalExpression{Operator: ast.AndOperator,
+							Left: &semantic.BinaryExpression{Operator: ast.LessThanOperator,
+								Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+								Right: &semantic.FloatLiteral{Value: 0.9}},
+							Right: &semantic.BinaryExpression{Operator: ast.GreaterThanOperator,
+								Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+								Right: &semantic.FloatLiteral{Value: 0.5}}},
+						Right: &semantic.BinaryExpression{Operator: ast.EqualOperator,
+							Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_measurement"},
+							Right: &semantic.StringLiteral{Value: "cpu"}}}},
+				}),
+				planner.CreateLogicalNode("yield4", &transformations.YieldProcedureSpec{Name: "result"}),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{1, 2},
+			},
+		},
+	},
+		{
+			name: "with swappable map and filter",
+			flux: `from(bucket: "telegraf") |> map(fn: (r) => r._value * 2.0) |> filter(fn: (r) => r._value < 10.0) |> yield(name: "result")`,
+			wantPlan: plantest.LogicalPlanSpec{
+				Nodes: []planner.PlanNode{
+					planner.CreateLogicalNode("from0", &inputs.FromProcedureSpec{Bucket: "telegraf"}),
+					planner.CreateLogicalNode("filter2_copy", &transformations.FilterProcedureSpec{Fn: &semantic.FunctionExpression{
+						Params: []*semantic.FunctionParam{{Key: &semantic.Identifier{Name: "r"}}},
+						Body: &semantic.BinaryExpression{Operator: ast.LessThanOperator,
+							Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+							Right: &semantic.FloatLiteral{Value: 10}},
+					}}),
+					planner.CreateLogicalNode("map1", &transformations.MapProcedureSpec{
+						Fn: &semantic.FunctionExpression{
+							Params: []*semantic.FunctionParam{{Key: &semantic.Identifier{Name: "r"}}},
+							Body: &semantic.BinaryExpression{Operator: ast.MultiplicationOperator,
+								Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+								Right: &semantic.FloatLiteral{Value: 2}}},
+						MergeKey: true,
+					}),
+					planner.CreateLogicalNode("yield3", &transformations.YieldProcedureSpec{Name: "result"}),
+				},
+				Edges: [][2]int{
+					{0, 1},
+					{1, 2},
+					{2, 3},
+				},
+			}},
+		{
+			name: "rules working together",
+			flux: `
+				from(bucket: "telegraf") |>
+					filter(fn: (r) => r._value != 0) |>
+					map(fn: (r) => r._value * 10) |>
+					filter(fn: (r) => f._value < 100) |>
+					yield(name: "result")`,
+			wantPlan: plantest.LogicalPlanSpec{
+				Nodes: []planner.PlanNode{
+					planner.CreateLogicalNode("from0", &inputs.FromProcedureSpec{Bucket: "telegraf"}),
+					planner.CreateLogicalNode("merged_filter1_filter3_copy", &transformations.FilterProcedureSpec{Fn: &semantic.FunctionExpression{
+						Params: []*semantic.FunctionParam{{Key: &semantic.Identifier{Name: "r"}}},
+						Body: &semantic.LogicalExpression{Operator: ast.AndOperator,
+							Left: &semantic.BinaryExpression{Operator: ast.LessThanOperator,
+								Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "f"}, Property: "_value"},
+								Right: &semantic.IntegerLiteral{Value: 100}},
+							Right: &semantic.BinaryExpression{Operator: ast.NotEqualOperator,
+								Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+								Right: &semantic.IntegerLiteral{}}},
+					}}),
+					planner.CreateLogicalNode("map2", &transformations.MapProcedureSpec{Fn: &semantic.FunctionExpression{
+						Params: []*semantic.FunctionParam{{Key: &semantic.Identifier{Name: "r"}}},
+						Body: &semantic.BinaryExpression{Operator: ast.MultiplicationOperator,
+							Left:  &semantic.MemberExpression{Object: &semantic.IdentifierExpression{Name: "r"}, Property: "_value"},
+							Right: &semantic.IntegerLiteral{Value: 10}}},
+						MergeKey: true,
+					}),
+					planner.CreateLogicalNode("yield4", &transformations.YieldProcedureSpec{Name: "result"}),
+				},
+				Edges: [][2]int{
+					{0, 1},
+					{1, 2},
+					{2, 3},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			fluxSpec, err := compile(tc.flux)
+			if err != nil {
+				t.Fatalf("could not compile flux query: %v", err)
+			}
+
+			logicalPlanner := planner.NewLogicalPlanner()
+			logicalPlan, err := logicalPlanner.Plan(fluxSpec)
+
+			wantPlan := plantest.CreateLogicalPlanSpec(&tc.wantPlan)
+			if err := plantest.ComparePlans(wantPlan, logicalPlan, plantest.CompareLogicalPlanNodes); err != nil {
+				t.Error(err)
+			}
+		})
+	}
+}

planner/physical.go

@@ -157,6 +157,15 @@ func (ppn *PhysicalPlanNode) Kind() ProcedureKind {
 	return ppn.Spec.Kind()
 }
 
+func (ppn *PhysicalPlanNode) ShallowCopy() PlanNode {
+	newNode := new(PhysicalPlanNode)
+	newNode.edges = ppn.edges.shallowCopy()
+	newNode.id = ppn.id + "_copy"
+	// TODO: the type assertion below... is it needed?
+	newNode.Spec = ppn.Spec.Copy().(PhysicalProcedureSpec)
+	return newNode
+}
+
 // Cost provides the self-cost (i.e., does not include the cost of its predecessors) for
 // this plan node.  Caller must provide statistics of predecessors to this node.
 func (ppn *PhysicalPlanNode) Cost(inStats []Statistics) (cost Cost, outStats Statistics) {

planner/plantest/cmp.go

@@ -2,7 +2,6 @@ package plantest
 
 import (
 	"fmt"
-
 	"github.com/google/go-cmp/cmp"
 	"github.com/influxdata/flux/planner"
 	"github.com/influxdata/flux/semantic/semantictest"
@@ -95,7 +94,7 @@ func cmpPlanNode(p, q planner.PlanNode) error {
 	// The specifications of both procedures must be the same
 	if !cmp.Equal(p.ProcedureSpec(), q.ProcedureSpec(), semantictest.CmpOptions...) {
 		return fmt.Errorf("procedure specs not equal -want(%s)/+got(%s) %s",
-			p.ID(), q.ID(), cmp.Diff(p.ProcedureSpec(), q.ProcedureSpec()))
+			p.ID(), q.ID(), cmp.Diff(p.ProcedureSpec(), q.ProcedureSpec(), semantictest.CmpOptions...))
 	}
 
 	return nil