collect(): defuse incoming values (#7171)

Author: mattnibs

compiler/rungen/vop.go

@@ -429,7 +429,7 @@ func (b *Builder) compileVamAgg(agg *dag.AggExpr) (*vamexpr.Aggregator, error) {
 			return nil, err
 		}
 	}
-	pattern, err := vamagg.NewPattern(name, agg.Distinct, agg.Expr != nil)
+	pattern, err := vamagg.NewPattern(b.sctx(), name, agg.Distinct, agg.Expr != nil)
 	if err != nil {
 		return nil, err
 	}

runtime/vam/expr/agg/agg.go

@@ -3,10 +3,11 @@ package agg
 import (
 	"fmt"
 
+	"github.com/brimdata/super"
 	"github.com/brimdata/super/runtime/vam/expr"
 )
 
-func NewPattern(op string, distinct, hasarg bool) (expr.AggPattern, error) {
+func NewPattern(sctx *super.Context, op string, distinct, hasarg bool) (expr.AggPattern, error) {
 	needarg := true
 	var pattern expr.AggPattern
 	switch op {
@@ -25,7 +26,7 @@ func NewPattern(op string, distinct, hasarg bool) (expr.AggPattern, error) {
 		}
 	case "array_agg":
 		pattern = func() expr.AggFunc {
-			return &arrayAgg{}
+			return &arrayAgg{*newCollect(sctx)}
 		}
 	case "blend":
 		pattern = func() expr.AggFunc {
@@ -57,7 +58,7 @@ func NewPattern(op string, distinct, hasarg bool) (expr.AggPattern, error) {
 		}
 	case "collect":
 		pattern = func() expr.AggFunc {
-			return &collect{}
+			return newCollect(sctx)
 		}
 	case "collect_map":
 		pattern = func() expr.AggFunc {

runtime/vam/expr/agg/collect.go

@@ -1,26 +1,31 @@
 package agg
 
 import (
-	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/brimdata/super"
+	"github.com/brimdata/super/runtime/vam/expr"
 	"github.com/brimdata/super/vector"
 	"github.com/brimdata/super/vector/vbuild"
 )
 
 type collect struct {
 	builder *vbuild.DynamicBuilder
+	defuse  *expr.Defuse
+}
+
+func newCollect(sctx *super.Context) *collect {
+	return &collect{defuse: expr.NewDefuse(sctx)}
 }
 
 func (c *collect) NoRip() bool { return true }
 
 func (c *collect) Consume(vec vector.Any) {
-	vector.Apply(vector.ApplyRipUnions, c.consume, vec)
+	vector.Apply(vector.ApplyRipUnions, c.consume, c.defuse.Eval(vec))
 }
 
 func (c *collect) consume(vecs ...vector.Any) vector.Any {
-	vec := filterNones(vecs[0])
-	if vec.Len() == 0 {
-		return vector.NewNone(vecs[0].Len())
+	vec := vecs[0]
+	if vec.Kind() == vector.KindNone || vec.Len() == 0 {
+		return vector.NewNull(vecs[0].Len())
 	}
 	if c.builder == nil {
 		c.builder = vbuild.NewDynamicBuilder()
@@ -29,36 +34,6 @@ func (c *collect) consume(vecs ...vector.Any) vector.Any {
 	return vector.NewNone(vecs[0].Len())
 }
 
-func filterNones(vec vector.Any) vector.Any {
-	switch mask := nonesMask(vec); {
-	case mask.IsEmpty():
-		return vec
-	case mask.GetCardinality() == uint64(vec.Len()):
-		return vector.NewNone(0)
-	default:
-		return vector.ReversePick(vec, mask.ToArray())
-	}
-}
-
-func nonesMask(vec vector.Any) *roaring.Bitmap {
-	vec = vector.Apply(vector.ApplyRipFusions|vector.ApplyRipUnions, func(vecs ...vector.Any) vector.Any {
-		return vecs[0]
-	}, vec)
-	bm := roaring.New()
-	if dynamic, ok := vec.(*vector.Dynamic); ok {
-		for i, vec := range dynamic.Values {
-			if vec.Len() > 0 && vec.Kind() == vector.KindNone {
-				bm.AddMany(dynamic.ReverseTagMap()[i])
-			}
-		}
-		return bm
-	}
-	if vec.Len() > 0 && vec.Kind() == vector.KindNone {
-		bm.AddRange(0, uint64(vec.Len()))
-	}
-	return bm
-}
-
 func (c *collect) Result(sctx *super.Context) vector.Any {
 	if c.builder == nil {
 		atyp := sctx.LookupTypeArray(super.TypeNone)

runtime/ztests/op/aggregate/array_agg.yaml

@@ -12,11 +12,9 @@ input: |
   fusion(1::(int64|null),<int64>)
   fusion(null::(int64|null),<null>)
 
-output-flags: -fusion
-
 output: |
   type foo=int64
-  [{a:1},{a:2},null,{b:1.5},error("missing"),1::foo,fusion(1::(int64|null),<int64>),fusion(null::(int64|null),<null>)]
+  [{a:1},{a:2},null,null,{b:1.5},error("missing"),1::foo,1]
 
 ---
 

runtime/ztests/op/aggregate/collect.yaml

@@ -12,11 +12,9 @@ input: |
   fusion(1::(int64|null),<int64>)
   fusion(null::(int64|null),<null>)
 
-output-flags: -fusion
-
 output: |
   type foo=int64
-  [{a:1},{a:2},null,{b:1.5},error("missing"),1::foo,fusion(1::(int64|null),<int64>),fusion(null::(int64|null),<null>)]
+  [{a:1},{a:2},null,null,{b:1.5},error("missing"),1::foo,1]
 
 ---
 

runtime/ztests/op/fuse.yaml

@@ -191,11 +191,9 @@ input: |
   [fusion({x?:1,y?:_::int64},<{x:int64}>),fusion({x?:_::int64,y?:2},<{y:int64}>)]
   [1,fusion(2::(int64|null),<int64>)]
 
-output-flags: -fusion
-
 output: |
   {x:1}
-  fusion({x?:1,y?:_::int64},<{x:int64}>)
-  fusion({x?:_::int64,y?:2},<{y:int64}>)
+  {x:1}
+  {y:2}
   1
   2