Add vector.FlattenUnions

Add FlattenUnions to vector package that recursive flattens unions into
then dynamic if necessary.

Author: mattnibs

runtime/vam/expr/mapcall.go

@@ -35,16 +35,7 @@ func (m *mapCall) eval(vecs ...vector.Any) vector.Any {
 	}
 	inner = m.lambda.Eval(inner)
 	if d, ok := inner.(*vector.Dynamic); ok {
-		var typs []super.Type
-		for _, vec := range d.Values {
-			typs = append(typs, vec.Type())
-		}
-		utyp, ok := m.sctx.LookupTypeUnion(super.UniqueTypes(typs))
-		if !ok {
-			//XXX need to vector deunion
-			panic(typs)
-		}
-		inner = vector.NewUnion(utyp, d.Tags, d.Values)
+		inner = vector.NewUnionFromDynamic(m.sctx, vector.FlattenUnions(d))
 	}
 	var out vector.Any
 	switch vec := vec.(type) {

runtime/vam/expr/mapexpr.go

@@ -69,12 +69,7 @@ func (m *MapExpr) build(vecs []vector.Any) vector.Any {
 	}
 	out := vb.Build(m.sctx)
 	if d, ok := out.(*vector.Dynamic); ok {
-		utyp, ok := m.sctx.LookupTypeUnion(super.UniqueTypes(typs))
-		if !ok {
-			//XXX need vector deunion
-			panic(typs)
-		}
-		out = &vector.Union{Typ: utyp, Dynamic: d}
+		out = vector.NewUnionFromDynamic(m.sctx, vector.FlattenUnions(d))
 	}
 	return out
 }

runtime/vam/expr/quiet.go

@@ -93,12 +93,7 @@ func (d *Dequiet) dequiet(vec vector.Any) vector.Any {
 	index := b.ToArray()
 	vec = vector.ReversePick(vec, index)
 	out := vector.Combine(vec, index, quiet).(*vector.Dynamic)
-	utyp, ok := d.sctx.LookupTypeUnion([]super.Type{vec.Type(), quiet.Type()})
-	if !ok {
-		// XXX need to deunion vec when it's a union
-		panic(vec.Type())
-	}
-	return vector.NewUnion(utyp, out.Tags, out.Values)
+	return vector.NewUnionFromDynamic(d.sctx, vector.FlattenUnions(out))
 }
 
 func (d *Dequiet) quietTmp(n uint32) vector.Any {

runtime/ztests/expr/map-func.yaml

@@ -34,8 +34,7 @@ output: |
 spq: |
   values map([1,2,3], lambda x: x=1?1:'foo'::(bool|string))
 
-# XXX this will be re-enabled before VENTURA
-vector: false
+vector: true
 
 output: |
   [1,"foo","foo"]::[int64|bool|string]

vector/union.go

@@ -1,6 +1,8 @@
 package vector
 
 import (
+	"slices"
+
 	"github.com/brimdata/super"
 	"github.com/brimdata/super/scode"
 )
@@ -49,3 +51,41 @@ func Deunion(vec Any) Any {
 	}
 	return vec
 }
+
+// FlattenUnions takes a Dynamic and recursively replaces any Union values
+// with their inner values, rewriting tags so that each slot points directly
+// to the leaf value vector.
+func FlattenUnions(d *Dynamic) *Dynamic {
+	hasUnion := slices.ContainsFunc(d.Values, func(vec Any) bool {
+		_, ok := vec.(*Union)
+		return ok
+	})
+	if !hasUnion {
+		return d
+	}
+	bases := make([]uint32, len(d.Values))
+	unions := make([]*Dynamic, len(d.Values))
+	var newValues []Any
+	for i, val := range d.Values {
+		bases[i] = uint32(len(newValues))
+		if u, ok := val.(*Union); ok {
+			flat := FlattenUnions(u.Dynamic)
+			unions[i] = flat
+			newValues = append(newValues, flat.Values...)
+		} else {
+			newValues = append(newValues, val)
+		}
+	}
+	forward := d.ForwardTagMap()
+	newTags := make([]uint32, len(d.Tags))
+	for slot, oldTag := range d.Tags {
+		base := bases[oldTag]
+		if d := unions[oldTag]; d != nil {
+			innerSlot := forward[slot]
+			newTags[slot] = base + d.Tags[innerSlot]
+		} else {
+			newTags[slot] = base
+		}
+	}
+	return NewDynamic(newTags, newValues)
+}