Add vectorized version of unblend()

Author: mattnibs

runtime/vam/expr/function/function.go

@@ -149,7 +149,9 @@ func New(sctx *super.Context, name string, narg int) (expr.Function, error) {
 	case "typeof":
 		f = &TypeOf{sctx}
 	case "unblend":
-		f = &samFunc{sctx, function.NewUnblend(sctx)}
+		f = &anonFunc{sctx, func(sctx *super.Context, vecs ...vector.Any) vector.Any {
+			return expr.Unblend(sctx, vecs[0])
+		}}
 	case "under":
 		f = newUnder(sctx)
 	case "unflatten":
@@ -200,3 +202,12 @@ func (f *samFunc) Call(args ...vector.Any) vector.Any {
 	}
 	return b.Build(f.sctx)
 }
+
+type anonFunc struct {
+	sctx *super.Context
+	fn   func(*super.Context, ...vector.Any) vector.Any
+}
+
+func (a *anonFunc) Call(args ...vector.Any) vector.Any {
+	return a.fn(a.sctx, args...)
+}

runtime/vam/expr/unblend.go

@@ -0,0 +1,259 @@
+package expr
+
+import (
+	"slices"
+
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/vector"
+)
+
+func Unblend(sctx *super.Context, vec vector.Any) vector.Any {
+	switch vec.Kind() {
+	case vector.KindRecord:
+		return unblendRecord(sctx, vec)
+	case vector.KindArray:
+		array := vector.PushdownContainerView(vec).(*vector.Array)
+		tags, inners, offsets := unblendArrayOrSet(sctx, array.Offsets, array.Values)
+		var vals []vector.Any
+		for i, inner := range inners {
+			typ := sctx.LookupTypeArray(inner.Type())
+			vals = append(vals, vector.NewArray(typ, offsets[i], inner))
+		}
+		if len(vals) > 1 {
+			return vector.NewDynamic(tags, vals)
+		}
+		return vals[0]
+	case vector.KindSet:
+		set := vector.PushdownContainerView(vec).(*vector.Set)
+		tags, inners, offsets := unblendArrayOrSet(sctx, set.Offsets, set.Values)
+		var vals []vector.Any
+		for i, inner := range inners {
+			typ := sctx.LookupTypeSet(inner.Type())
+			vals = append(vals, vector.NewSet(typ, offsets[i], inner))
+		}
+		if len(vals) > 1 {
+			return vector.NewDynamic(tags, vals)
+		}
+		return vals[0]
+	case vector.KindMap:
+		return unblendMap(sctx, vector.PushdownContainerView(vec).(*vector.Map))
+	case vector.KindUnion:
+		return vector.Apply(true, func(vecs ...vector.Any) vector.Any {
+			return Unblend(sctx, vecs[0])
+		}, vec)
+	}
+	return vec
+}
+
+func unblendRecord(sctx *super.Context, in vector.Any) vector.Any {
+	vec := in
+	var index []uint32
+	if view, ok := vec.(*vector.View); ok {
+		index = view.Index
+		vec = view.Any
+	}
+	rec := vec.(*vector.Record)
+	if len(rec.Fields) == 0 {
+		return in
+	}
+	fields := slices.Clone(rec.Fields)
+	if index != nil {
+		for i, field := range fields {
+			fields[i] = vector.Pick(field, index)
+		}
+	}
+	for i, field := range fields {
+		fields[i] = Unblend(sctx, field)
+	}
+	return vector.Apply(false, func(vecs ...vector.Any) vector.Any {
+		var fields []super.Field
+		var fvecs []vector.Any
+		for i, vec := range vecs {
+			if _, ok := vec.(*vector.None); ok {
+				continue
+			}
+			fvecs = append(fvecs, vec)
+			fields = append(fields, super.NewField(rec.Typ.Fields[i].Name, vec.Type()))
+		}
+		rtyp := sctx.MustLookupTypeRecord(fields)
+		return vector.NewRecord(rtyp, fvecs, vecs[0].Len())
+	}, fields...)
+}
+
+func unblendArrayOrSet(sctx *super.Context, offsets []uint32, elements vector.Any) ([]uint32, []vector.Any, [][]uint32) {
+	elements = Unblend(sctx, elements)
+	dynamic, ok := elements.(*vector.Dynamic)
+	if !ok {
+		return nil, []vector.Any{elements}, [][]uint32{offsets}
+	}
+	union := vector.NewUnionFromDynamic(sctx, dynamic)
+	slotTypes := typesOfSlotsInList(sctx, union, offsets)
+	// Accumulate unique array types.
+	m := make(map[super.Type][]uint32)
+	for i, typ := range slotTypes {
+		m[typ] = append(m[typ], uint32(i))
+	}
+	dtags := make([]uint32, len(offsets)-1)
+	var inners []vector.Any
+	var offs [][]uint32
+	for typ, index := range m {
+		inner, off := subsetOfList(sctx, union, offsets, index, typ)
+		for _, idx := range index {
+			dtags[idx] = uint32(len(inners))
+		}
+		inners = append(inners, inner)
+		offs = append(offs, off)
+	}
+	return dtags, inners, offs
+}
+
+func unblendMap(sctx *super.Context, vmap *vector.Map) vector.Any {
+	keys := Unblend(sctx, vmap.Keys)
+	if dynamic, ok := keys.(*vector.Dynamic); ok {
+		keys = vector.NewUnionFromDynamic(sctx, dynamic)
+	}
+	vals := Unblend(sctx, vmap.Values)
+	if dynamic, ok := vals.(*vector.Dynamic); ok {
+		vals = vector.NewUnionFromDynamic(sctx, dynamic)
+	}
+	if keys.Kind() != vector.KindUnion && vals.Kind() != vector.KindUnion {
+		return vmap
+	}
+	keySlotTypes := typesOfSlotsInList(sctx, keys, vmap.Offsets)
+	valSlotTypes := typesOfSlotsInList(sctx, vals, vmap.Offsets)
+	type mapType struct {
+		key super.Type
+		val super.Type
+	}
+	// Accumulate unique map types.
+	m := make(map[mapType][]uint32)
+	for i := range vmap.Len() {
+		mtyp := mapType{keySlotTypes[i], valSlotTypes[i]}
+		m[mtyp] = append(m[mtyp], uint32(i))
+	}
+	dtags := make([]uint32, len(vmap.Offsets)-1)
+	var vecs []vector.Any
+	for mtyp, index := range m {
+		keys, offsets := subsetOfList(sctx, keys, vmap.Offsets, index, mtyp.key)
+		vals, _ := subsetOfList(sctx, vals, vmap.Offsets, index, mtyp.val)
+		for _, idx := range index {
+			dtags[idx] = uint32(len(vecs))
+		}
+		typ := sctx.LookupTypeMap(keys.Type(), vals.Type())
+		vecs = append(vecs, vector.NewMap(typ, offsets, keys, vals))
+	}
+	if len(vecs) == 1 {
+		return vecs[0]
+	}
+	return vector.NewDynamic(dtags, vecs)
+}
+
+func subsetOfList(sctx *super.Context, elements vector.Any, parentOffsets, index []uint32, typ super.Type) (vector.Any, []uint32) {
+	if typ == super.TypeNull {
+		nulls := vector.NewNull(uint32(len(index)))
+		offsets := make([]uint32, len(index)+1)
+		return nulls, offsets
+	}
+	var allVals []vector.Any
+	union, ok := elements.(*vector.Union)
+	if ok {
+		allVals = union.Values
+	} else {
+		allVals = []vector.Any{elements}
+	}
+	var subTypes []super.Type
+	utyp, ok := typ.(*super.TypeUnion)
+	if ok {
+		subTypes = slices.Clone(utyp.Types)
+	} else {
+		subTypes = append(subTypes, typ)
+	}
+	subVals := make([]vector.Any, len(subTypes))
+	// map parent union tags to subset union tags
+	tagMap := make([]uint32, len(allVals))
+	for i, typ := range subTypes {
+		idx := slices.IndexFunc(allVals, func(vec vector.Any) bool {
+			return vec.Type() == typ
+		})
+		tagMap[idx] = uint32(i)
+		subVals[i] = allVals[idx]
+	}
+	// Generate:
+	// - offsets for new array
+	// - indexes to create view on values
+	// - tags for union (if applicable)
+	var forwardTags []uint32
+	if union != nil {
+		forwardTags = union.ForwardTagMap()
+	}
+	var tags []uint32
+	indexes := make([][]uint32, len(subTypes))
+	suboffsets := []uint32{0}
+	for _, idx := range index {
+		start, end := parentOffsets[idx], parentOffsets[idx+1]
+		if union != nil {
+			for i, origTag := range union.Tags[start:end] {
+				tag := tagMap[origTag]
+				tags = append(tags, tag)
+				indexes[tag] = append(indexes[tag], forwardTags[start+uint32(i)])
+			}
+		} else {
+			for i := start; i < end; i++ {
+				indexes[0] = append(indexes[0], i)
+			}
+		}
+		suboffsets = append(suboffsets, uint32(len(tags)))
+	}
+	for i, val := range subVals {
+		subVals[i] = vector.Pick(val, indexes[i])
+	}
+	var inner vector.Any
+	if len(subVals) > 1 {
+		inner = vector.NewUnion(sctx.LookupTypeUnion(subTypes), tags, subVals)
+	} else {
+		inner = subVals[0]
+	}
+	return inner, suboffsets
+}
+
+func typesOfSlotsInList(sctx *super.Context, inner vector.Any, offsets []uint32) []super.Type {
+	union, _ := inner.(*vector.Union)
+	var alltypes []super.Type
+	if union != nil {
+		for _, val := range union.Values {
+			alltypes = append(alltypes, val.Type())
+		}
+	} else {
+		alltypes = []super.Type{inner.Type()}
+	}
+	n := uint32(len(offsets) - 1)
+	slotTypes := make([]super.Type, n)
+	for i := range n {
+		if union != nil {
+			slotTypes[i] = typeOfRange(sctx, union, alltypes, offsets[i], offsets[i+1])
+		} else {
+			slotTypes[i] = alltypes[0]
+		}
+	}
+	return slotTypes
+}
+
+func typeOfRange(sctx *super.Context, union *vector.Union, alltypes []super.Type, start, end uint32) super.Type {
+	tags := slices.Clone(union.Tags[start:end])
+	slices.Sort(tags)
+	uniq := slices.Compact(tags)
+	if len(uniq) == 0 {
+		return super.TypeNull
+	}
+	if len(uniq) == 1 {
+		return alltypes[uniq[0]]
+	}
+	if len(uniq) == len(alltypes) {
+		return union.Typ
+	}
+	var types []super.Type
+	for _, tag := range uniq {
+		types = append(types, alltypes[tag])
+	}
+	return sctx.LookupTypeUnion(types)
+}

runtime/ztests/expr/function/unblend.yaml

@@ -1,19 +1,25 @@
-# Test unblend on nested records.
+# Test unblend on records.
 spq: unblend(this)
 
+vector: true
+
 input: |
   {x:{y:2::(int64|bool)},z?:"foo"}
   {x:{y:true::(int64|bool)},z?:_::string}
+  {}
 
 output: |
   {x:{y:2},z:"foo"}
   {x:{y:true}}
+  {}
 
 ---
 
 # Test unblend on root unions.
 spq: unblend(this)
 
+vector: true
+
 input: |
   1::(int64|bool|{x:int64})
   true::(int64|bool|{x:int64})
@@ -29,6 +35,8 @@ output: |
 # Test unblend on arrays.
 spq: unblend(this)
 
+vector: true
+
 input: |
   [1::(int64|bool|string),2::(int64|bool|string),"string"::(int64|bool|string)]
   [1::(int64|bool|string),2::(int64|bool|string),true::(int64|bool|string)]
@@ -45,9 +53,28 @@ output: |
 
 ---
 
+# Test unblend on elements in an array.
+spq: unblend(this)
+
+vector: true
+
+input: |
+  [{x:1::(bool|int64|string)},{x:true::(bool|int64|string)},{x:"foo"::(bool|int64|string)}]
+  [{x:2::(bool|int64|string)},{x:3::(bool|int64|string)},{x:"bar"::(bool|int64|string)}]
+  [{x:true::(bool|int64|string)},{x:false::(bool|int64|string)},{x:"bar"::(bool|int64|string)}]
+
+output: |
+  [{x:1},{x:true},{x:"foo"}]
+  [{x:2},{x:3},{x:"bar"}]
+  [{x:true},{x:false},{x:"bar"}]
+
+---
+
 # Test unblend on sets.
 spq: unblend(this)
 
+vector: true
+
 input: |
   |[1::(int64|bool|string),2::(int64|bool|string),"string"::(int64|bool|string)]|
   |[1::(int64|bool|string),2::(int64|bool|string),true::(int64|bool|string)]|
@@ -71,10 +98,12 @@ input: |
   |{"k1"::(int64|string):"v1"::(int64|string),"k2"::(int64|string):2::(int64|string)}|
   |{"k1"::(int64|string):1::(int64|string),"k2"::(int64|string):2::(int64|string)}|
   |{"k1"::(int64|string):"v1"::(int64|string),2::(int64|string):"v2"::(int64|string)}|
+  |{"k1":"v1"::(int64|string),"k2":"v2"::(int64|string)}|
   |{}|::|{(int64|string):(int64|string)}|
 
 output: |
   |{"k1":"v1","k2":2}|
   |{"k1":1,"k2":2}|
   |{2:"v2","k1":"v1"}|
+  |{"k1":"v1","k2":"v2"}|
   |{}|

vector/array.go

@@ -66,3 +66,37 @@ func Inner(val Any) Any {
 	}
 	panic(val)
 }
+
+func PushdownContainerView(val Any) Any {
+	view, ok := val.(*View)
+	if !ok {
+		return val
+	}
+	switch val := view.Any.(type) {
+	case *Array:
+		inner, offsets := pickList(val.Values, view.Index, val.Offsets)
+		return NewArray(val.Typ, offsets, inner)
+	case *Set:
+		inner, offsets := pickList(val.Values, view.Index, val.Offsets)
+		return NewSet(val.Typ, offsets, inner)
+	case *Map:
+		keys, offsets := pickList(val.Keys, view.Index, val.Offsets)
+		values, _ := pickList(val.Values, view.Index, val.Offsets)
+		return NewMap(val.Typ, offsets, keys, values)
+	default:
+		panic(val)
+	}
+}
+
+func pickList(inner Any, index, offsets []uint32) (Any, []uint32) {
+	newOffsets := []uint32{0}
+	var innerIndex []uint32
+	for _, idx := range index {
+		start, end := offsets[idx], offsets[idx+1]
+		for ; start < end; start++ {
+			innerIndex = append(innerIndex, start)
+		}
+		newOffsets = append(newOffsets, uint32(len(innerIndex)))
+	}
+	return Pick(inner, innerIndex), newOffsets
+}