From 0d1de325ea7733619dc95adee9ae9e3baba416c3 Mon Sep 17 00:00:00 2001
From: Matthew Nibecker <hello@mattnibecker.com>
Date: Tue, 4 Aug 2026 10:31:27 -0700
Subject: [PATCH] collect(): preserve input order

This commit preserves input order in the collect() expression by
filtering nones rather than using Apply to filter nones.
---
 runtime/vam/expr/agg/collect.go            | 39 +++++++++++++++++++++-
 runtime/ztests/op/aggregate/array_agg.yaml |  2 +-
 runtime/ztests/op/aggregate/collect.yaml   |  2 +-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/runtime/vam/expr/agg/collect.go b/runtime/vam/expr/agg/collect.go
index c2120a0a76..5a943c6f49 100644
--- a/runtime/vam/expr/agg/collect.go
+++ b/runtime/vam/expr/agg/collect.go
@@ -1,6 +1,7 @@
 package agg
 
 import (
+	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/brimdata/super"
 	"github.com/brimdata/super/runtime/vam/expr"
 	"github.com/brimdata/super/vector"
@@ -19,7 +20,43 @@ func newCollect(sctx *super.Context) *collect {
 func (c *collect) NoRip() bool { return true }
 
 func (c *collect) Consume(vec vector.Any) {
-	vector.Apply(vector.ApplyRipUnions, c.consume, c.defuse.Eval(vec))
+	vec = vector.Apply(vector.ApplyRipUnions, func(vecs ...vector.Any) vector.Any {
+		return vecs[0]
+	}, c.defuse.Eval(vec))
+	if vec = filterNones(vec); vec.Len() == 0 {
+		return
+	}
+	if c.builder == nil {
+		c.builder = vbuild.NewDynamicBuilder()
+	}
+	c.builder.Write(vec)
+}
+
+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 {
+	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) consume(vecs ...vector.Any) vector.Any {
diff --git a/runtime/ztests/op/aggregate/array_agg.yaml b/runtime/ztests/op/aggregate/array_agg.yaml
index 960fd81537..bd64d5aa55 100644
--- a/runtime/ztests/op/aggregate/array_agg.yaml
+++ b/runtime/ztests/op/aggregate/array_agg.yaml
@@ -14,7 +14,7 @@ input: |
 
 output: |
   type foo=int64
-  [{a:1},{a:2},null,null,{b:1.5},error("missing"),1::foo,1]
+  [{a:1},{a:2},null,{b:1.5},error("missing"),1::foo,1,null]
 
 ---
 
diff --git a/runtime/ztests/op/aggregate/collect.yaml b/runtime/ztests/op/aggregate/collect.yaml
index c7405b0cb6..8d6249e23b 100644
--- a/runtime/ztests/op/aggregate/collect.yaml
+++ b/runtime/ztests/op/aggregate/collect.yaml
@@ -14,7 +14,7 @@ input: |
 
 output: |
   type foo=int64
-  [{a:1},{a:2},null,null,{b:1.5},error("missing"),1::foo,1]
+  [{a:1},{a:2},null,{b:1.5},error("missing"),1::foo,1,null]
 
 ---