feat: added PushDownGroupAggregate planner rewrite rule (#18085)

Added a (disabled) planner rule that matches:
   ReadGroupPhys -> { count }

It uses the same physical spec node for group to implement the aggregate. The
rule requires:
 * the pushDownGroupAggregateCount feature flag enabled
 * no existing aggregate present in the ReadGroup
 * use of the "_value" column only

flags.yml

@@ -52,3 +52,9 @@
   default: false
   contact: Lyon Hill
   expose: true
+
+- name: Push Down Group Aggregate Count
+  description: Enable the count variant of PushDownGroupAggregate planner rule
+  key: pushDownGroupAggregateCount
+  default: false
+  contact: Query Team

query/stdlib/influxdata/influxdb/rules.go

@@ -27,6 +27,10 @@ func init() {
 		// For this rule to take effect the appropriate capabilities must be
 		// added AND feature flags must be enabled.
 		// PushDownWindowAggregateRule{},
+
+		// For this rule to take effect the corresponding feature flags must be
+		// enabled.
+		// PushDownGroupAggregateRule{},
 	)
 }
 
@@ -690,7 +694,7 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		}
 
 		minSpec := fnNode.ProcedureSpec().(*universe.MinProcedureSpec)
-		if minSpec.Column != "_value" {
+		if minSpec.Column != execute.DefaultValueColLabel {
 			return pn, false, nil
 		}
 	case universe.MaxKind:
@@ -699,7 +703,7 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		}
 
 		maxSpec := fnNode.ProcedureSpec().(*universe.MaxProcedureSpec)
-		if maxSpec.Column != "_value" {
+		if maxSpec.Column != execute.DefaultValueColLabel {
 			return pn, false, nil
 		}
 	case universe.MeanKind:
@@ -708,7 +712,7 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		}
 
 		meanSpec := fnNode.ProcedureSpec().(*universe.MeanProcedureSpec)
-		if len(meanSpec.Columns) != 1 || meanSpec.Columns[0] != "_value" {
+		if len(meanSpec.Columns) != 1 || meanSpec.Columns[0] != execute.DefaultValueColLabel {
 			return pn, false, nil
 		}
 	case universe.CountKind:
@@ -717,7 +721,7 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		}
 
 		countSpec := fnNode.ProcedureSpec().(*universe.CountProcedureSpec)
-		if len(countSpec.Columns) != 1 || countSpec.Columns[0] != "_value" {
+		if len(countSpec.Columns) != 1 || countSpec.Columns[0] != execute.DefaultValueColLabel {
 			return pn, false, nil
 		}
 	case universe.SumKind:
@@ -726,7 +730,7 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		}
 
 		sumSpec := fnNode.ProcedureSpec().(*universe.SumProcedureSpec)
-		if len(sumSpec.Columns) != 1 || sumSpec.Columns[0] != "_value" {
+		if len(sumSpec.Columns) != 1 || sumSpec.Columns[0] != execute.DefaultValueColLabel {
 			return pn, false, nil
 		}
 	}
@@ -764,3 +768,52 @@ func (PushDownWindowAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (p
 		WindowEvery:       window.Every.Nanoseconds(),
 	}), true, nil
 }
+
+//
+// Push Down of group aggregates.
+// ReadGroupPhys |> { count }
+//
+type PushDownGroupAggregateRule struct{}
+
+func (PushDownGroupAggregateRule) Name() string {
+	return "PushDownGroupAggregateRule"
+}
+
+func (rule PushDownGroupAggregateRule) Pattern() plan.Pattern {
+	return plan.OneOf(
+		[]plan.ProcedureKind{
+			universe.CountKind,
+		},
+		plan.Pat(ReadGroupPhysKind))
+}
+
+func (PushDownGroupAggregateRule) Rewrite(ctx context.Context, pn plan.Node) (plan.Node, bool, error) {
+	// Check the aggregate function spec. Require operation on _value.
+	var aggregateMethod string
+	fnNode := pn
+	switch fnNode.Kind() {
+	case universe.CountKind:
+		if !feature.PushDownGroupAggregateCount().Enabled(ctx) {
+			return pn, false, nil
+		}
+
+		countSpec := fnNode.ProcedureSpec().(*universe.CountProcedureSpec)
+		if len(countSpec.Columns) != 1 || countSpec.Columns[0] != execute.DefaultValueColLabel {
+			return pn, false, nil
+		}
+		aggregateMethod = universe.CountKind
+	}
+
+	groupNode := fnNode.Predecessors()[0]
+	groupSpec := groupNode.ProcedureSpec().(*ReadGroupPhysSpec)
+
+	// Group spec must not already have an aggregate method.
+	if len(groupSpec.AggregateMethod) > 0 {
+		return pn, false, nil
+	}
+
+	// Rule passes.
+	rewrite := *groupSpec.Copy().(*ReadGroupPhysSpec)
+	rewrite.AggregateMethod = aggregateMethod
+	return plan.CreatePhysicalNode("ReadGroup", &rewrite), true, nil
+}

query/stdlib/influxdata/influxdb/rules_test.go

@@ -1207,7 +1207,7 @@ func TestPushDownWindowAggregateRule(t *testing.T) {
 
 	tests := make([]plantest.RuleTestCase, 0)
 
-	// construct a simple plan with a specific window
+	// construct a simple plan with a specific window and aggregate function
 	simplePlanWithWindowAgg := func( window universe.WindowProcedureSpec, agg plan.NodeID, spec plan.ProcedureSpec ) *plantest.PlanSpec {
 		return &plantest.PlanSpec{
 			Nodes: []plan.Node{
@@ -1568,3 +1568,217 @@ func TestPushDownWindowAggregateRule(t *testing.T) {
 		})
 	}
 }
+
+//
+// Group Aggregate Testing
+//
+func TestPushDownGroupAggregateRule(t *testing.T) {
+	// Turn on all flags
+	flagger := mock.NewFlagger(map[feature.Flag] interface{}{
+		feature.PushDownGroupAggregateCount(): true,
+	})
+
+	ctx, _ := feature.Annotate(context.Background(), flagger)
+
+	readGroupAgg := func(aggregateMethod string) *influxdb.ReadGroupPhysSpec{
+		return &influxdb.ReadGroupPhysSpec{
+			ReadRangePhysSpec: influxdb.ReadRangePhysSpec {
+				Bucket: "my-bucket",
+				Bounds: flux.Bounds{
+					Start: fluxTime(5),
+					Stop:  fluxTime(10),
+				},
+			},
+			GroupMode: flux.GroupModeBy,
+			GroupKeys: []string{"_measurement", "tag0", "tag1"},
+			AggregateMethod: aggregateMethod,
+		}
+	}
+	readGroup := func() *influxdb.ReadGroupPhysSpec {
+		return readGroupAgg("")
+	}
+
+	tests := make([]plantest.RuleTestCase, 0)
+
+	// construct a simple plan with a specific aggregate
+	simplePlanWithAgg := func( agg plan.NodeID, spec plan.ProcedureSpec ) *plantest.PlanSpec {
+		return &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreateLogicalNode("ReadGroup", readGroup()),
+				plan.CreateLogicalNode(agg, spec),
+			},
+			Edges: [][2]int{
+				{0, 1},
+			},
+		}
+	}
+
+	// construct a simple result
+	simpleResult := func( proc string ) *plantest.PlanSpec {
+		return &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreatePhysicalNode("ReadGroup", readGroupAgg(proc) ),
+			},
+		}
+	}
+
+	minProcedureSpec := func() *universe.MinProcedureSpec {
+		return &universe.MinProcedureSpec{
+			SelectorConfig: execute.SelectorConfig{Column: "_value"},
+		}
+	}
+	countProcedureSpec := func() *universe.CountProcedureSpec {
+		return &universe.CountProcedureSpec{
+			AggregateConfig: execute.AggregateConfig{Columns: []string{"_value"}},
+		}
+	}
+	sumProcedureSpec := func() *universe.SumProcedureSpec {
+		return &universe.SumProcedureSpec{
+			AggregateConfig: execute.AggregateConfig{Columns: []string{"_value"}},
+		}
+	}
+
+	// ReadGroup -> count => ReadGroup
+	tests = append(tests, plantest.RuleTestCase{
+		Context: ctx,
+		Name: "SimplePassCount",
+		Rules: []plan.Rule{ influxdb.PushDownGroupAggregateRule{}},
+		Before: simplePlanWithAgg( "count", countProcedureSpec() ),
+		After: simpleResult( "count" ),
+	})
+
+	// Rewrite with successors
+	// ReadGroup -> count -> sum {2} => ReadGroup -> count {2}
+	tests = append(tests, plantest.RuleTestCase{
+		Context: ctx,
+		Name: "WithSuccessor1",
+		Rules: []plan.Rule{ influxdb.PushDownGroupAggregateRule{}},
+		Before: &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreateLogicalNode("ReadGroup", readGroup()),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+				plan.CreateLogicalNode("sum", sumProcedureSpec() ),
+				plan.CreateLogicalNode("sum", sumProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{1, 2},
+				{1, 3},
+			},
+		},
+		After: &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreatePhysicalNode("ReadGroup", readGroupAgg("count") ),
+				plan.CreateLogicalNode("sum", sumProcedureSpec() ),
+				plan.CreateLogicalNode("sum", sumProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{0, 2},
+			},
+		},
+	})
+
+	// Cannot replace a ReadGroup that already has an aggregate. This exercises
+	// the check that ReadGroup aggregate is not set.
+	// ReadGroup -> count -> count => ReadGroup -> count
+	tests = append(tests, plantest.RuleTestCase{
+		Context: ctx,
+		Name: "WithSuccessor2",
+		Rules: []plan.Rule{ influxdb.PushDownGroupAggregateRule{}},
+		Before: &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreateLogicalNode("ReadGroup", readGroup()),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{1, 2},
+			},
+		},
+		After: &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreatePhysicalNode("ReadGroup", readGroupAgg("count") ),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+			},
+		},
+	})
+
+	// Bad count column
+	// ReadGroup -> count => NO-CHANGE
+	tests = append(tests, plantest.RuleTestCase{
+		Name: "BadCountCol",
+		Context: ctx,
+		Rules: []plan.Rule{ influxdb.PushDownGroupAggregateRule{}},
+		Before: simplePlanWithAgg( "count", &universe.CountProcedureSpec{
+			AggregateConfig: execute.AggregateConfig{Columns: []string{"_valmoo"}},
+		}),
+		NoChange: true,
+	})
+
+	// No match due to a collapsed node having a successor
+	// ReadGroup -> count 
+	//          \-> min
+	tests = append(tests, plantest.RuleTestCase{
+		Name: "CollapsedWithSuccessor",
+		Context: ctx,
+		Rules: []plan.Rule{ influxdb.PushDownGroupAggregateRule{}},
+		Before: &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreateLogicalNode("ReadGroup", readGroup()),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+				plan.CreateLogicalNode("min", minProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{0, 2},
+			},
+		},
+		NoChange: true,
+	})
+
+	// No pattern match
+	// ReadGroup -> filter -> min -> NO-CHANGE
+	pushableFn1 := executetest.FunctionExpression(t, `(r) => true`)
+
+	makeResolvedFilterFn := func(expr *semantic.FunctionExpression) interpreter.ResolvedFunction {
+		return interpreter.ResolvedFunction{
+			Scope: nil,
+			Fn:    expr,
+		}
+	}
+	noPatternMatch1 := func() *plantest.PlanSpec{
+		return &plantest.PlanSpec{
+			Nodes: []plan.Node{
+				plan.CreateLogicalNode("ReadGroup", readGroup()),
+				plan.CreatePhysicalNode("filter", &universe.FilterProcedureSpec{
+					Fn: makeResolvedFilterFn(pushableFn1),
+				}),
+				plan.CreateLogicalNode("count", countProcedureSpec() ),
+			},
+			Edges: [][2]int{
+				{0, 1},
+				{1, 2},
+			},
+		}
+	}
+	tests = append(tests, plantest.RuleTestCase{
+		Name: "NoPatternMatch",
+		Context: ctx,
+		Rules: []plan.Rule{ influxdb.PushDownWindowAggregateRule{}},
+		Before: noPatternMatch1(),
+		NoChange: true,
+	})
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+			plantest.PhysicalRuleTestHelper(t, &tc)
+		})
+	}
+}