From fbf72c543d5b2d9ed0736fe6c6bd8be92312eea2 Mon Sep 17 00:00:00 2001
From: Matthew Nibecker <hello@mattnibecker.com>
Date: Mon, 3 Feb 2025 13:42:32 -0800
Subject: [PATCH] vam: flatten()

---
 runtime/vam/expr/function/flatten.go          | 33 +++++++++++++++++++
 runtime/vam/expr/function/function.go         |  2 ++
 .../expr/function}/flatten.yaml               |  4 +++
 3 files changed, 39 insertions(+)
 create mode 100644 runtime/vam/expr/function/flatten.go
 rename runtime/{sam/expr/function/ztests => ztests/expr/function}/flatten.yaml (86%)

diff --git a/runtime/vam/expr/function/flatten.go b/runtime/vam/expr/function/flatten.go
new file mode 100644
index 0000000000..5f451b936d
--- /dev/null
+++ b/runtime/vam/expr/function/flatten.go
@@ -0,0 +1,33 @@
+package function
+
+import (
+	"github.com/brimdata/super"
+	samfunc "github.com/brimdata/super/runtime/sam/expr/function"
+	"github.com/brimdata/super/vector"
+	"github.com/brimdata/super/zcode"
+)
+
+type flatten struct {
+	fn *samfunc.Flatten
+}
+
+func newFlatten(zctx *super.Context) *flatten {
+	return &flatten{samfunc.NewFlatten(zctx)}
+}
+
+func (f *flatten) Call(args ...vector.Any) vector.Any {
+	vec := vector.Under(args[0])
+	rtyp := super.TypeRecordOf(vec.Type())
+	if rtyp == nil {
+		return args[0]
+	}
+	builder := vector.NewDynamicBuilder()
+	var b zcode.Builder
+	for i := range vec.Len() {
+		b.Reset()
+		vec.Serialize(&b, i)
+		val := f.fn.Call(nil, []super.Value{super.NewValue(rtyp, b.Bytes().Body())})
+		builder.Write(val)
+	}
+	return builder.Build()
+}
diff --git a/runtime/vam/expr/function/function.go b/runtime/vam/expr/function/function.go
index ebdef8adae..07a481625e 100644
--- a/runtime/vam/expr/function/function.go
+++ b/runtime/vam/expr/function/function.go
@@ -36,6 +36,8 @@ func New(zctx *super.Context, name string, narg int) (expr.Function, field.Path,
 		f = &Error{zctx}
 	case "fields":
 		f = NewFields(zctx)
+	case "flatten":
+		f = newFlatten(zctx)
 	case "floor":
 		f = &Floor{zctx}
 	case "grep":
diff --git a/runtime/sam/expr/function/ztests/flatten.yaml b/runtime/ztests/expr/function/flatten.yaml
similarity index 86%
rename from runtime/sam/expr/function/ztests/flatten.yaml
rename to runtime/ztests/expr/function/flatten.yaml
index f9d09dda3c..8b1f3faca3 100644
--- a/runtime/sam/expr/function/ztests/flatten.yaml
+++ b/runtime/ztests/expr/function/flatten.yaml
@@ -1,5 +1,7 @@
 zed: yield flatten(this)
 
+vector: true
+
 input: |
   127.0.0.1
   {a:1,b:null}
@@ -7,6 +9,7 @@ input: |
   {a:1,b:{c:2,d:{e:4,f:5}}}
   {foo:null({bar:int32})}
   {foo:{bar:2(int32)}}
+  {foo:{bar:null(int32)}}
 
 output: |
   127.0.0.1
@@ -15,3 +18,4 @@ output: |
   [{key:["a"],value:1},{key:["b","c"],value:2},{key:["b","d","e"],value:4},{key:["b","d","f"],value:5}]
   [{key:["foo"],value:null({bar:int32})}]
   [{key:["foo","bar"],value:2(int32)}]
+  [{key:["foo","bar"],value:null(int32)}]