Add scalar aggregate for keyless aggregations (#6558)

This commit adds a separate runtime operator for scalar aggregates. If
no input is encountered scalar aggregates return default values for each
aggregate value.

Author: mattnibs

book/src/super-sql/aggregates/count.md

@@ -53,7 +53,7 @@ count() by k | sort
 {k:2,count:1}
 ```
 
-A simple count with no input values returns no output:
+A simple count with no input values returns 0:
 ```mdtest-spq
 # spq
 where grep("bar", this) | count()
@@ -62,6 +62,7 @@ where grep("bar", this) | count()
 "foo"
 10.0.0.1
 # expected output
+0
 ```
 
 Count can return an explicit zero when using a `filter` clause in the aggregation:

book/src/tutorials/jq.md

@@ -1072,6 +1072,7 @@ which produces an output like this:
 {reviewers:|["nwt","henridf","mccanne"]|}
 {reviewers:|["mccanne","mattnibs"]|}
 {reviewers:|["henridf","mccanne","mattnibs"]|}
+{reviewers:null}
 {reviewers:|["henridf","mccanne","mattnibs"]|}
 ...
 ```

compiler/rungen/aggregate.go

@@ -12,7 +12,7 @@ import (
 	"github.com/brimdata/super/sbuf"
 )
 
-func (b *Builder) compileAggregate(parent sbuf.Puller, a *dag.AggregateOp) (*aggregate.Op, error) {
+func (b *Builder) compileAggregate(parent sbuf.Puller, a *dag.AggregateOp) (sbuf.Puller, error) {
 	keys, err := b.compileAssignments(a.Keys)
 	if err != nil {
 		return nil, err
@@ -22,6 +22,9 @@ func (b *Builder) compileAggregate(parent sbuf.Puller, a *dag.AggregateOp) (*agg
 		return nil, err
 	}
 	dir := order.Direction(a.InputSortDir)
+	if len(keys) == 0 {
+		return aggregate.NewScalar(b.rctx, parent, names, reducers, a.PartialsIn, a.PartialsOut)
+	}
 	return aggregate.New(b.rctx, parent, keys, names, reducers, a.Limit, dir, a.PartialsIn, a.PartialsOut)
 }
 

compiler/rungen/vop.go

@@ -422,6 +422,9 @@ func (b *Builder) compileVamAggregate(s *dag.AggregateOp, parent vector.Puller)
 		keyNames = append(keyNames, lhs.Path)
 		keyExprs = append(keyExprs, rhs)
 	}
+	if len(keyExprs) == 0 {
+		return aggregate.NewScalar(parent, b.sctx(), aggs, aggNames, aggExprs, s.PartialsIn, s.PartialsOut)
+	}
 	return aggregate.New(parent, b.sctx(), aggNames, aggExprs, aggs, keyNames, keyExprs, s.PartialsIn, s.PartialsOut)
 }
 

runtime/sam/op/aggregate/aggregate.go

@@ -111,7 +111,7 @@ func NewAggregator(ctx context.Context, sctx *super.Context, keyRefs, keyExprs,
 	}, nil
 }
 
-func New(rctx *runtime.Context, parent sbuf.Puller, keys []expr.Assignment, aggNames field.List, aggs []*expr.Aggregator, limit int, inputSortDir order.Direction, partialsIn, partialsOut bool) (*Op, error) {
+func New(rctx *runtime.Context, parent sbuf.Puller, keys []expr.Assignment, aggNames field.List, aggs []*expr.Aggregator, limit int, inputSortDir order.Direction, partialsIn, partialsOut bool) (sbuf.Puller, error) {
 	names := make(field.List, 0, len(keys)+len(aggNames))
 	for _, e := range keys {
 		p, ok := e.LHS.Path()

runtime/sam/op/aggregate/scalar.go

@@ -0,0 +1,90 @@
+package aggregate
+
+import (
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/pkg/field"
+	"github.com/brimdata/super/runtime"
+	"github.com/brimdata/super/runtime/sam/expr"
+	"github.com/brimdata/super/sbuf"
+)
+
+type scalarAggregate struct {
+	sctx        *super.Context
+	parent      sbuf.Puller
+	builder     *super.RecordBuilder
+	aggRefs     []expr.Evaluator
+	aggs        []*expr.Aggregator
+	partialsIn  bool
+	partialsOut bool
+	row         valRow
+}
+
+func NewScalar(rctx *runtime.Context, parent sbuf.Puller, aggNames []field.Path, aggs []*expr.Aggregator, partialsIn, partialsOut bool) (sbuf.Puller, error) {
+	builder, err := super.NewRecordBuilder(rctx.Sctx, aggNames)
+	if err != nil {
+		return nil, err
+	}
+	aggRefs := make([]expr.Evaluator, 0, len(aggNames))
+	for _, fieldName := range aggNames {
+		aggRefs = append(aggRefs, expr.NewDottedExpr(rctx.Sctx, fieldName))
+	}
+	return &scalarAggregate{
+		sctx:        rctx.Sctx,
+		parent:      parent,
+		builder:     builder,
+		aggRefs:     aggRefs,
+		aggs:        aggs,
+		partialsIn:  partialsIn,
+		partialsOut: partialsOut,
+		row:         newValRow(aggs),
+	}, nil
+}
+
+func (s *scalarAggregate) Pull(done bool) (sbuf.Batch, error) {
+	if done {
+		s.row = nil
+		return nil, nil
+	}
+	if s.row == nil {
+		s.row = newValRow(s.aggs)
+		return nil, nil
+	}
+	for {
+		batch, err := s.parent.Pull(false)
+		if err != nil {
+			return nil, err
+		}
+		if batch == nil {
+			return s.result(), nil
+		}
+		for _, val := range batch.Values() {
+			if s.partialsIn {
+				s.row.consumeAsPartial(val, s.aggRefs)
+			} else {
+				s.row.apply(s.sctx, s.aggs, val)
+			}
+		}
+	}
+}
+
+func (s *scalarAggregate) result() sbuf.Batch {
+	var typs []super.Type
+	s.builder.Reset()
+	for _, agg := range s.row {
+		var val super.Value
+		if s.partialsOut {
+			val = agg.ResultAsPartial(s.sctx)
+		} else {
+			val = agg.Result(s.sctx)
+		}
+		typs = append(typs, val.Type())
+		s.builder.Append(val.Bytes())
+	}
+	s.row = nil
+	typ := s.builder.Type(typs)
+	b, err := s.builder.Encode()
+	if err != nil {
+		panic(err)
+	}
+	return sbuf.NewBatch([]super.Value{super.NewValue(typ, b)})
+}

runtime/vam/op/aggregate/scalar.go

@@ -0,0 +1,98 @@
+package aggregate
+
+import (
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/pkg/field"
+	"github.com/brimdata/super/runtime/vam/expr"
+	"github.com/brimdata/super/runtime/vam/expr/agg"
+	"github.com/brimdata/super/vector"
+	"github.com/brimdata/super/vector/bitvec"
+)
+
+type scalarAggregate struct {
+	parent      vector.Puller
+	sctx        *super.Context
+	aggExprs    []expr.Evaluator
+	aggs        []*expr.Aggregator
+	builder     *vector.RecordBuilder
+	partialsIn  bool
+	partialsOut bool
+
+	funcs []agg.Func
+}
+
+func NewScalar(parent vector.Puller, sctx *super.Context, aggs []*expr.Aggregator, aggNames []field.Path, aggExprs []expr.Evaluator, partialsIn, partialsOut bool) (vector.Puller, error) {
+	builder, err := vector.NewRecordBuilder(sctx, aggNames)
+	if err != nil {
+		return nil, err
+	}
+	return &scalarAggregate{
+		parent:      parent,
+		sctx:        sctx,
+		aggs:        aggs,
+		aggExprs:    aggExprs,
+		builder:     builder,
+		partialsIn:  partialsIn,
+		partialsOut: partialsOut,
+		funcs:       newFuncs(aggs),
+	}, nil
+}
+
+func (s *scalarAggregate) Pull(done bool) (vector.Any, error) {
+	if s.funcs == nil {
+		s.funcs = newFuncs(s.aggs)
+		return nil, nil
+	}
+	for {
+		vec, err := s.parent.Pull(done)
+		if err != nil {
+			return nil, err
+		}
+		if vec == nil {
+			return s.result(), nil
+		}
+		var vals []vector.Any
+		if s.partialsIn {
+			for _, e := range s.aggExprs {
+				vals = append(vals, e.Eval(vec))
+			}
+		} else {
+			for _, e := range s.aggs {
+				vals = append(vals, e.Eval(vec))
+			}
+		}
+		vector.Apply(false, func(vecs ...vector.Any) vector.Any {
+			for i, vec := range vecs {
+				if s.partialsIn {
+					s.funcs[i].ConsumeAsPartial(vec)
+				} else {
+					s.funcs[i].Consume(vec)
+				}
+			}
+			return vector.NewConst(super.Null, vecs[0].Len(), bitvec.Zero)
+		}, vals...)
+	}
+}
+
+func newFuncs(aggs []*expr.Aggregator) []agg.Func {
+	var funcs []agg.Func
+	for _, agg := range aggs {
+		funcs = append(funcs, agg.Pattern())
+	}
+	return funcs
+}
+
+func (s *scalarAggregate) result() vector.Any {
+	var vecs []vector.Any
+	for _, f := range s.funcs {
+		b := vector.NewDynamicBuilder()
+		if s.partialsOut {
+			b.Write(f.ResultAsPartial(s.sctx))
+		} else {
+			b.Write(f.Result(s.sctx))
+		}
+		vecs = append(vecs, b.Build())
+	}
+	s.funcs = nil
+	return s.builder.New(vecs, bitvec.Zero)
+}

runtime/ztests/op/aggregate/scalar-empty.yaml

@@ -0,0 +1,24 @@
+spq: |
+  where false
+  | count(*), sum(this), and(this), or(this), min(this), avg(this), collect(this)
+
+vector: true
+
+output: |
+  {count:0,sum:null,and:null::bool,or:null::bool,min:null,avg:null::float64,collect:null}
+
+---
+
+# Test that scalar aggregations with no inputs work in an unnset subquery.
+spq: |
+  unnest this into ( where false | count(*) )
+
+vector: true
+
+input: |
+  [null]
+  [null]
+
+output: |
+  0
+  0