Use nested context in java TestStream

Author: shunping

runners/prism/java/build.gradle

@@ -106,15 +106,14 @@ def sickbayTests = [
 
     // Prism doesn't support multiple TestStreams.
     'org.apache.beam.sdk.testing.TestStreamTest.testMultipleStreams',
-    // Sometimes fails missing a final 'AFTER'. Otherwise, Hangs in ElementManager.FailBundle due to a held stageState lock.
-    'org.apache.beam.sdk.testing.TestStreamTest.testMultiStage',
 
     // GroupIntoBatchesTest tests that fail:
-    // Teststream has bad KV encodings due to using an outer context.
-    'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode',
-    'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInFixedWindow',
-    // sdk worker disconnected
-    'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInGlobalWindow',
+    // Wrong number of elements in windows after GroupIntoBatches: null
+    // Incorrect output collection after GroupIntoBatches: null
+    // 'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode',
+    // 'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInFixedWindow',
+    // varint overflow 4120879041988 on counting step
+    // 'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInGlobalWindow',
     // ShardedKey not yet implemented.
     'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testWithShardedKeyInGlobalWindow',
 

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"
@@ -277,33 +275,14 @@ func executePipeline(ctx context.Context, wks map[string]*worker.W, j *jobservic
 				}
 
 				// Ensure awareness of the coder used for the teststream.
-				cID, err := lpUnknownCoders(pyld.GetCoderId(), coders, comps.GetCoders())
+				_, 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
 				}
-				// Hack for Java Strings in test stream, since it doesn't encode them correctly.
-				forceLP := cID == "StringUtf8Coder" || cID != pyld.GetCoderId()
-				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.
-					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()
-					}
-				}
 
 				tsb := em.AddTestStream(stage.ID, t.Outputs)
 				for _, e := range pyld.GetEvents() {

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

@@ -27,6 +27,7 @@ import (
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/jobservices"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/urns"
 	"golang.org/x/exp/maps"
+	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -152,6 +153,7 @@ func (p *preprocessor) preProcessGraph(comps *pipepb.Components, j *jobservices.
 	}
 
 	// Extract URNs for the given transform.
+	slog.Debug("components in pipeline proto", "proto", prototext.Format(comps))
 
 	keptLeaves := maps.Keys(leaves)
 	sort.Strings(keptLeaves)

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/TestStreamTranslation.java

@@ -123,7 +123,8 @@ static <T> RunnerApi.TestStreamPayload.Event eventToProto(
                   .setTimestamp(element.getTimestamp().getMillis())
                   .setEncodedElement(
                       ByteString.copyFrom(
-                          CoderUtils.encodeToByteArray(coder, element.getValue()))));
+                          CoderUtils.encodeToByteArray(
+                              coder, element.getValue(), Coder.Context.NESTED))));
         }
         return RunnerApi.TestStreamPayload.Event.newBuilder().setElementEvent(builder).build();
       default:
@@ -149,7 +150,8 @@ static <T> TestStream.Event<T> eventFromProto(
             protoEvent.getElementEvent().getElementsList()) {
           decodedElements.add(
               TimestampedValue.of(
-                  CoderUtils.decodeFromByteArray(coder, element.getEncodedElement().toByteArray()),
+                  CoderUtils.decodeFromByteArray(
+                      coder, element.getEncodedElement().toByteArray(), Coder.Context.NESTED),
                   new Instant(element.getTimestamp())));
         }
         return TestStream.ElementEvent.add(decodedElements);