vam: Fix avg ConsumePartial (#5501)

Closes #5500
brimdata · Nov 25, 2024 · e7d0376 · e7d0376
1 parent 7c930f2
commit e7d0376
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/runtime/vam/expr/agg/avg.go b/runtime/vam/expr/agg/avg.go
@@ -34,8 +34,16 @@ const (
 )
 
 func (a *avg) ConsumeAsPartial(partial vector.Any) {
+	if partial.Len() != 1 {
+		panic("avg: invalid partial")
+	}
+	idx := uint32(0)
+	if view, ok := partial.(*vector.View); ok {
+		idx = view.Index[0]
+		partial = view.Any
+	}
 	rec, ok := partial.(*vector.Record)
-	if !ok || rec.Len() != 1 {
+	if !ok {
 		panic("avg: invalid partial")
 	}
 	si, ok1 := rec.Typ.IndexOfField(sumName)
@@ -48,8 +56,8 @@ func (a *avg) ConsumeAsPartial(partial vector.Any) {
 	if sumVal.Type() != super.TypeFloat64 || countVal.Type() != super.TypeUint64 {
 		panic("avg: invalid partial")
 	}
-	sum, _ := vector.FloatValue(sumVal, 0)
-	count, _ := vector.UintValue(countVal, 0)
+	sum, _ := vector.FloatValue(sumVal, idx)
+	count, _ := vector.UintValue(countVal, idx)
 	a.sum += sum
 	a.count += count
 }

diff --git a/runtime/vam/op/summarize/summarize.go b/runtime/vam/op/summarize/summarize.go
@@ -32,6 +32,7 @@ func New(parent vector.Puller, zctx *super.Context, aggNames []field.Path, aggEx
 	}
 	return &Summarize{
 		parent:      parent,
+		zctx:        zctx,
 		aggs:        aggs,
 		aggExprs:    aggExprs,
 		keyExprs:    keyExprs,