Move the logic to LP TestStream encoded bytes to transform preparer

sdks/go/pkg/beam/runners/prism/internal/coders.go

@@ -199,11 +199,11 @@ func lpUnknownCoders(cID string, bundle, base map[string]*pipepb.Coder) (string,
 }
 
 // forceLpCoder always add a new LP-coder for a given coder into the "base" map
-func forceLpCoder(cID string, base map[string]*pipepb.Coder) (string, error) {
+func forceLpCoder(cID string, bundle, base map[string]*pipepb.Coder) (string, error) {
 	// First check if we've already added the LP version of this coder to coders already.
 	lpcID := cID + "_flp"
 	// Check if we've done this one before.
-	if _, ok := base[lpcID]; ok {
+	if _, ok := bundle[lpcID]; ok {
 		return lpcID, nil
 	}
 	// Look up the canonical location.
@@ -219,7 +219,7 @@ func forceLpCoder(cID string, base map[string]*pipepb.Coder) (string, error) {
 		},
 		ComponentCoderIds: []string{cID},
 	}
-	base[lpcID] = lpc
+	bundle[lpcID] = lpc
 	return lpcID, nil
 }
 

sdks/go/pkg/beam/runners/prism/internal/execute.go

@@ -16,7 +16,6 @@
 package internal
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -27,7 +26,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
 	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
@@ -270,67 +268,21 @@ func executePipeline(ctx context.Context, wks map[string]*worker.W, j *jobservic
 			case urns.TransformTestStream:
 				// Add a synthetic stage that should largely be unused.
 				em.AddStage(stage.ID, nil, maps.Values(t.GetOutputs()), nil)
+
 				// Decode the test stream, and convert it to the various events for the ElementManager.
 				var pyld pipepb.TestStreamPayload
 				if err := proto.Unmarshal(t.GetSpec().GetPayload(), &pyld); err != nil {
 					return fmt.Errorf("prism error building stage %v - decoding TestStreamPayload: \n%w", stage.ID, err)
 				}
 
-				// Ensure awareness of the coder used for the teststream.
-				cID, err := lpUnknownCoders(pyld.GetCoderId(), coders, comps.GetCoders())
-				if err != nil {
-					panic(err)
-				}
-				mayLP := func(v []byte) []byte {
-					//slog.Warn("teststream bytes", "value", string(v), "bytes", v)
-					return v
-				}
-				// If the TestStream coder needs to be LP'ed or if it is a coder that has different
-				// behaviors between nested context and outer context (in Java SDK), then we must
-				// LP this coder and the TestStream data elements.
-				forceLP := cID != pyld.GetCoderId() ||
-					coders[cID].GetSpec().GetUrn() == urns.CoderStringUTF8 ||
-					coders[cID].GetSpec().GetUrn() == urns.CoderBytes ||
-					coders[cID].GetSpec().GetUrn() == urns.CoderKV
-				if forceLP {
-					// slog.Warn("recoding TestStreamValue", "cID", cID, "newUrn", coders[cID].GetSpec().GetUrn(), "payloadCoder", pyld.GetCoderId(), "oldUrn", coders[pyld.GetCoderId()].GetSpec().GetUrn())
-					// The coder needed length prefixing. For simplicity, add a length prefix to each
-					// encoded element, since we will be sending a length prefixed coder to consume
-					// this anyway. This is simpler than trying to find all the re-written coders after the fact.
-					// This also adds a LP-coder for the original coder in comps.
-					cID, err := forceLpCoder(pyld.GetCoderId(), comps.GetCoders())
-					if err != nil {
-						panic(err)
-					}
-					slog.Debug("teststream: add coder", "coderId", cID)
-
-					mayLP = func(v []byte) []byte {
-						var buf bytes.Buffer
-						if err := coder.EncodeVarInt((int64)(len(v)), &buf); err != nil {
-							panic(err)
-						}
-						if _, err := buf.Write(v); err != nil {
-							panic(err)
-						}
-						//slog.Warn("teststream bytes - after LP", "value", string(v), "bytes", buf.Bytes())
-						return buf.Bytes()
-					}
-
-					// we need to change Coder and Pcollection in comps directly before they are used to build descriptors
-					for _, col := range t.GetOutputs() {
-						oCID := comps.Pcollections[col].CoderId
-						comps.Pcollections[col].CoderId = cID
-						slog.Debug("teststream: rewrite coder for output pcoll", "colId", col, "oldId", oCID, "newId", cID)
-					}
-				}
-
 				tsb := em.AddTestStream(stage.ID, t.Outputs)
 				for _, e := range pyld.GetEvents() {
 					switch ev := e.GetEvent().(type) {
 					case *pipepb.TestStreamPayload_Event_ElementEvent:
 						var elms []engine.TestStreamElement
 						for _, e := range ev.ElementEvent.GetElements() {
-							elms = append(elms, engine.TestStreamElement{Encoded: mayLP(e.GetEncodedElement()), EventTime: mtime.FromMilliseconds(e.GetTimestamp())})
+							// Encoded bytes are already handled in handleTestStream if needed.
+							elms = append(elms, engine.TestStreamElement{Encoded: e.GetEncodedElement(), EventTime: mtime.FromMilliseconds(e.GetTimestamp())})
 						}
 						tsb.AddElementEvent(ev.ElementEvent.GetTag(), elms)
 					case *pipepb.TestStreamPayload_Event_WatermarkEvent:

sdks/go/pkg/beam/runners/prism/internal/handlerunner.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"log/slog"
 	"reflect"
 	"sort"
 	"strings"
@@ -72,6 +73,7 @@ func (*runner) PrepareUrns() []string {
 		urns.TransformRedistributeArbitrarily,
 		urns.TransformRedistributeByKey,
 		urns.TransformFlatten,
+		urns.TransformTestStream,
 	}
 }
 
@@ -82,6 +84,8 @@ func (h *runner) PrepareTransform(tid string, t *pipepb.PTransform, comps *pipep
 		return h.handleFlatten(tid, t, comps)
 	case urns.TransformReshuffle, urns.TransformRedistributeArbitrarily, urns.TransformRedistributeByKey:
 		return h.handleReshuffle(tid, t, comps)
+	case urns.TransformTestStream:
+		return h.handleTestStream(tid, t, comps)
 	default:
 		panic("unknown urn to Prepare: " + t.GetSpec().GetUrn())
 	}
@@ -216,6 +220,111 @@ func (h *runner) handleReshuffle(tid string, t *pipepb.PTransform, comps *pipepb
 	}
 }
 
+func (h *runner) handleTestStream(tid string, t *pipepb.PTransform, comps *pipepb.Components) prepareResult {
+	var pyld pipepb.TestStreamPayload
+	if err := proto.Unmarshal(t.GetSpec().GetPayload(), &pyld); err != nil {
+		panic("Failed to decode TestStreamPayload: " + err.Error())
+	}
+	coders := map[string]*pipepb.Coder{}
+	// Ensure awareness of the coder used for the teststream.
+	cID, err := lpUnknownCoders(pyld.GetCoderId(), coders, comps.GetCoders())
+	if err != nil {
+		panic(err)
+	}
+
+	// If the TestStream coder needs to be LP'ed or if it is a coder that has different
+	// behaviors between nested context and outer context (in Java SDK), then we must
+	// LP this coder and the TestStream data elements.
+	forceLP := (cID != pyld.GetCoderId() && coders[pyld.GetCoderId()].GetSpec().GetUrn() != "beam:go:coder:custom:v1") ||
+		coders[cID].GetSpec().GetUrn() == urns.CoderStringUTF8 ||
+		coders[cID].GetSpec().GetUrn() == urns.CoderBytes ||
+		coders[cID].GetSpec().GetUrn() == urns.CoderKV
+
+	if !forceLP {
+		return prepareResult{SubbedComps: &pipepb.Components{
+			Transforms: map[string]*pipepb.PTransform{tid: t},
+		}}
+	}
+
+	// The coder needed length prefixing. For simplicity, add a length prefix to each
+	// encoded element, since we will be sending a length prefixed coder to consume
+	// this anyway. This is simpler than trying to find all the re-written coders after the fact.
+	// This also adds a LP-coder for the original coder in comps.
+	cID, err = forceLpCoder(pyld.GetCoderId(), coders, comps.GetCoders())
+	if err != nil {
+		panic(err)
+	}
+	slog.Debug("teststream: add coder", "coderId", cID)
+
+	mustLP := func(v []byte) []byte {
+		var buf bytes.Buffer
+		if err := coder.EncodeVarInt((int64)(len(v)), &buf); err != nil {
+			panic(err)
+		}
+		if _, err := buf.Write(v); err != nil {
+			panic(err)
+		}
+		return buf.Bytes()
+	}
+
+	// We need to loop over the events.
+	// For element events, we need to apply the mayLP function to the encoded element.
+	// Then we construct a new payload with the modified events.
+	var newEvents []*pipepb.TestStreamPayload_Event
+	for _, event := range pyld.GetEvents() {
+		switch event.GetEvent().(type) {
+		case *pipepb.TestStreamPayload_Event_ElementEvent:
+			elms := event.GetElementEvent().GetElements()
+			var newElms []*pipepb.TestStreamPayload_TimestampedElement
+			for _, elm := range elms {
+				newElm := proto.Clone(elm).(*pipepb.TestStreamPayload_TimestampedElement)
+				newElm.EncodedElement = mustLP(elm.GetEncodedElement())
+				slog.Debug("handleTestStream: rewrite bytes",
+					"before:", string(elm.GetEncodedElement()),
+					"after:", string(newElm.GetEncodedElement()))
+				newElms = append(newElms, newElm)
+			}
+			newEvents = append(newEvents, &pipepb.TestStreamPayload_Event{
+				Event: &pipepb.TestStreamPayload_Event_ElementEvent{
+					ElementEvent: &pipepb.TestStreamPayload_Event_AddElements{
+						Elements: newElms,
+					},
+				},
+			})
+		default:
+			newEvents = append(newEvents, event)
+		}
+	}
+	newPyld := &pipepb.TestStreamPayload{
+		CoderId:  cID,
+		Events:   newEvents,
+		Endpoint: pyld.GetEndpoint(),
+	}
+	b, err := proto.Marshal(newPyld)
+	if err != nil {
+		panic(fmt.Sprintf("couldn't marshal new test stream payload: %v", err))
+	}
+
+	ts := proto.Clone(t).(*pipepb.PTransform)
+	ts.GetSpec().Payload = b
+
+	pcolSubs := map[string]*pipepb.PCollection{}
+	for _, gi := range ts.GetOutputs() {
+		pcol := comps.GetPcollections()[gi]
+		newPcol := proto.Clone(pcol).(*pipepb.PCollection)
+		newPcol.CoderId = cID
+		slog.Debug("handleTestStream: rewrite coder for output pcoll", "colId", gi, "oldId", pcol.CoderId, "newId", newPcol.CoderId)
+		pcolSubs[gi] = newPcol
+	}
+
+	tSubs := map[string]*pipepb.PTransform{tid: ts}
+	return prepareResult{SubbedComps: &pipepb.Components{
+		Transforms:   tSubs,
+		Pcollections: pcolSubs,
+		Coders:       coders,
+	}}
+}
+
 var _ transformExecuter = (*runner)(nil)
 
 func (*runner) ExecuteUrns() []string {