diff --git a/sdks/go/pkg/beam/runners/prism/internal/coders.go b/sdks/go/pkg/beam/runners/prism/internal/coders.go index 885d0eeef436..d326a332b8d3 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/coders.go +++ b/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 } diff --git a/sdks/go/pkg/beam/runners/prism/internal/execute.go b/sdks/go/pkg/beam/runners/prism/internal/execute.go index 7c7526b3d4db..a8e5b364e272 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/execute.go +++ b/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: diff --git a/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go b/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go index 988dd9ec7ed9..7b1ecee19771 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go +++ b/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 {